Jungfraujoch/tests/CUDAMemHelpersTest.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#include <atomic>
#include <optional>
#include <thread>

#ifdef JFJOCH_USE_CUDA

#include <cstring>
#include <future>
#include <latch>
#include <set>
#include <thread>

#include "../image_analysis/indexing/CUDAMemHelpers.h"
#include "../image_analysis/indexing/IndexerThreadPool.h"

// Whether the current GPU shares the host's memory (GB10, Jetson). A test that fills the card until an
// allocation fails would there fill the host's RAM - every process on the machine runs out with it.
static bool gpu_is_integrated() {
    int device = 0, integrated = 0;
    return cudaGetDevice(&device) == cudaSuccess
           && cudaDeviceGetAttribute(&integrated, cudaDevAttrIntegrated, device) == cudaSuccess && integrated;
}

// The broker starts fresh worker threads for every data collection. A stream per thread ever started
// is a device-memory leak that ends a long-running broker, so threads that follow one another have to
// end up on the same stream.
TEST_CASE("CudaAllocationStream_ReusedByLaterThreads", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    std::set<cudaStream_t> seen;
    for (int i = 0; i < 100; i++)
        std::thread([&seen] {
            CudaDevicePtr<uint8_t> buffer(1 << 16);
            seen.insert(cuda_allocation_stream());
        }).join();

    CHECK(seen.size() == 1);
    CHECK(*seen.begin() != nullptr);
}

TEST_CASE("CudaAllocationStream_DistinctForConcurrentThreads", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    // Neither thread may exit - and hand its stream back - before the other has taken its own.
    std::latch both_have_one(2);
    auto take = [&both_have_one] {
        const cudaStream_t stream = cuda_allocation_stream();
        both_have_one.count_down();
        both_have_one.wait();
        return stream;
    };
    auto a = std::async(std::launch::async, take);
    auto b = std::async(std::launch::async, take);
    const cudaStream_t stream_a = a.get();
    const cudaStream_t stream_b = b.get();

    CHECK(stream_a != nullptr);
    CHECK(stream_b != nullptr);
    CHECK(stream_a != stream_b);
}

// The shadow finder builds its GPU accumulator on a throw-away thread and frees it from another one
// long after. The free is ordered on the allocating thread's stream, so that stream has to outlive
// the thread - including while a later thread has borrowed it.
TEST_CASE("CudaDevicePtr_FreedAfterAllocatingThreadExited", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    auto buffer = std::async(std::launch::async, [] {
        return std::make_unique<CudaDevicePtr<uint8_t>>(1 << 20);
    }).get();
    std::thread([] { CudaDevicePtr<uint8_t> other(1 << 16); }).join();

    REQUIRE(cudaMemset(buffer->get(), 0, 1 << 20) == cudaSuccess);
    buffer.reset();
    CHECK(cudaDeviceSynchronize() == cudaSuccess);
    CHECK(cudaGetLastError() == cudaSuccess);
}

// A pool that cannot serve the request is not an error: the buffer comes from cudaMalloc instead. It
// must then not stay behind as the thread's last error, or the cudaGetLastError() after the next
// kernel launch throws "out of memory" over work that went fine.
TEST_CASE("CudaDevicePtr_HandledPoolFailureLeavesNoError", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    int device = 0;
    REQUIRE(cudaGetDevice(&device) == cudaSuccess);

    cudaMemPoolProps props{};
    props.allocType = cudaMemAllocationTypePinned;
    props.location.type = cudaMemLocationTypeDevice;
    props.location.id = device;
    props.maxSize = 4 << 20;

    cudaMemPool_t previous = nullptr, capped = nullptr;
    REQUIRE(cudaDeviceGetMemPool(&previous, device) == cudaSuccess);
    REQUIRE(cudaMemPoolCreate(&capped, &props) == cudaSuccess);
    REQUIRE(cudaDeviceSetMemPool(device, capped) == cudaSuccess);
    {
        CudaDevicePtr<uint8_t> buffer(64 << 20);
        CHECK(buffer.get() != nullptr);
        CHECK(cudaGetLastError() == cudaSuccess);
    }
    REQUIRE(cudaDeviceSetMemPool(device, previous) == cudaSuccess);
    cudaMemPoolDestroy(capped);
}

// Engines built one after another take the pinned buffers the previous one gave back - also a little
// larger ones, within the size class - while a frame-sized buffer is not kept at all.
TEST_CASE("CudaPinnedPool_ReusedByTheNextEngine", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    cuda_release_pinned_pool();
    void *first = nullptr;
    {
        CudaHostPtr<float> a(1000);
        first = a.get();
    }
    {
        PinnedVector<float> b(1020);
        CHECK(static_cast<void *>(b.data()) == first);
    }
    {
        CudaHostPtr<uint8_t> large(CUDA_PINNED_POOL_MAX_BUFFER + 1);
        large.get()[CUDA_PINNED_POOL_MAX_BUFFER] = 1;
    }
    CHECK(cuda_release_pinned_pool() == CUDA_PINNED_POOL_SLAB_BUFFERS * CUDA_PINNED_POOL_MIN_BUFFER);
    CHECK(cuda_release_pinned_pool() == 0);
}

// The small buffers are pieces of one slab, which goes back to the driver only once every piece is back.
TEST_CASE("CudaPinnedPool_SlabKeptWhileAPieceIsInUse", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    cuda_release_pinned_pool();
    {
        CudaHostPtr<float> a(1000);
        CudaHostPtr<float> b(1000);
        CHECK(std::abs(static_cast<uint8_t *>(static_cast<void *>(a.get()))
                       - static_cast<uint8_t *>(static_cast<void *>(b.get()))) == CUDA_PINNED_POOL_MIN_BUFFER);
        CHECK(cuda_release_pinned_pool() == 0);
        a.get()[999] = 1.0f;
    }
    CHECK(cuda_release_pinned_pool() == CUDA_PINNED_POOL_SLAB_BUFFERS * CUDA_PINNED_POOL_MIN_BUFFER);
}

// Within a scope a cuFFT plan is made once per shape and shared: two users at once get the same plan,
// each transforming on its own stream with its own work area, bit for bit as on a plan of its own; the
// plan is kept when they are done, and the next user gets it. When the scope ends, every kept plan is
// destroyed; outside a scope a plan goes with its last user.
TEST_CASE("CudaFFTPlan_SharedAndKept", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    REQUIRE(jfjoch_cufft_plans::kept() == 0);
    auto scope = std::make_unique<CudaFFTPlanScope>();

    const int n = 1000, batch = 7, nc = n / 2 + 1;
    const CudaFFTShape shape{.n = {n}, .idist = n, .odist = nc, .type = CUFFT_R2C, .batch = batch};
    std::vector<float> input(n * batch);
    for (size_t i = 0; i < input.size(); i++)
        input[i] = static_cast<float>((i * 7919) % 1013) - 500.0f;
    CudaDevicePtr<float> d_in(input.size());
    REQUIRE(cudaMemcpy(d_in, input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice) == cudaSuccess);

    auto transform = [&](const CudaFFTPlan &plan, cudaStream_t stream) {
        CudaDevicePtr<cufftComplex> d_out(nc * batch);
        REQUIRE(plan.R2C(d_in, d_out) == CUFFT_SUCCESS);
        std::vector<cufftComplex> out(nc * batch);
        REQUIRE(cudaMemcpyAsync(out.data(), d_out, out.size() * sizeof(cufftComplex), cudaMemcpyDeviceToHost,
                                stream) == cudaSuccess);
        REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
        return out;
    };

    CudaStream stream_a, stream_b, stream_c;
    std::vector<cufftComplex> first;
    cufftHandle handle = 0;
    {
        CudaFFTPlan a(shape, stream_a);
        CudaFFTPlan b(shape, stream_b);
        handle = a.Handle();
        CHECK(b.Handle() == handle);
        CHECK(jfjoch_cufft_plans::kept() == 1);
        first = transform(a, stream_a);
        const std::vector<cufftComplex> second = transform(b, stream_b);
        CHECK(memcmp(second.data(), first.data(), first.size() * sizeof(cufftComplex)) == 0);
    }
    CHECK(jfjoch_cufft_plans::kept() == 1);

    const size_t work = CudaFFTPlan::WorkBytes(shape);
    CudaDevicePtr<char> work_area(std::max<size_t>(work, 1));
    CudaFFTPlan c(shape, stream_c, work_area);
    CHECK(c.Handle() == handle);
    CHECK(c.WorkBytes() == work);
    const std::vector<cufftComplex> again = transform(c, stream_c);
    REQUIRE(again.size() == first.size());
    CHECK(memcmp(again.data(), first.data(), first.size() * sizeof(cufftComplex)) == 0);

    // The scope ends while c still uses the plan: it goes with c.
    scope.reset();
    CHECK(jfjoch_cufft_plans::kept() == 1);
    c = CudaFFTPlan();
    CHECK(jfjoch_cufft_plans::kept() == 0);

    CHECK(CudaFFTPlan::WorkBytes(shape) == work);
    { CudaFFTPlan d(shape); }
    CHECK(jfjoch_cufft_plans::kept() == 0);
}

// An FFT indexer that finds the card full waits for the GPU work beside it to give its memory back, and
// is built then: the main line's for the probes, a probe's for the main line's image pass - never for a
// probe, itself included. With nothing to wait for it fails at once, as one that does not fit always has.
TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    if (gpu_is_integrated())
        SKIP("Fills the card, which on an integrated GPU is the host's memory");

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT).IndexingThreads(1);
    DiffractionExperiment x;
    x.ImportIndexingSettings(settings);
    const std::vector<Coord> recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}};

    // Every card, held by the work beside: an indexer worker takes its card round-robin (pin_gpu), so on
    // a machine with several the one it lands on is not the test thread's. Topped up right before each
    // Run, so memory something else gave back in between does not let the indexer fit.
    const int gpus = get_gpu_count();
    std::vector<std::vector<CudaDevicePtr<uint8_t>>> held(gpus);
    const auto fill = [&] {
        for (int d = 0; d < gpus; d++) {
            ScopedGpuPin pin(d);
            for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2)
                while (true) {
                    try {
                        held[d].emplace_back(chunk, CudaAlloc::Synchronous);
                    } catch (const JFJochException &) {
                        break;
                    }
                }
        }
    };
    const auto release = [&] {
        for (int d = 0; d < gpus; d++) {
            ScopedGpuPin pin(d);
            held[d].clear();
        }
    };
    // The pools' indexers go back to the allocation pool when they are destroyed; trim it after them, so
    // the next test finds the cards free.
    {
        IndexerThreadPool probe(settings, IndexerConstruction::OnFirstUse, /*beside=*/true);
        IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse);

        {
            GPUWorkBeside work;
            fill();
            CHECK(probe.Run(x, recip).error != nullptr);
        }
        {
            std::optional<GPUImagePass> image_pass(std::in_place);
            GPUWorkBeside work;
            fill();
            std::thread ends([&] {
                std::this_thread::sleep_for(std::chrono::milliseconds(500));
                release();
                image_pass.reset();
            });
            CHECK(probe.Run(x, recip).error == nullptr);
            ends.join();
        }
        {
            std::optional<GPUWorkBeside> work(std::in_place);
            fill();
            std::thread ends([&] {
                std::this_thread::sleep_for(std::chrono::milliseconds(500));
                release();
                work.reset();
            });
            CHECK(pool.Run(x, recip).error == nullptr);
            ends.join();
        }
    }
    for (int d = 0; d < gpus; d++) {
        ScopedGpuPin pin(d);
        cuda_trim_allocation_pool();
    }
}

// A merge helper's buffer that finds the card full while a probe pass holds the rest of it waits for
// that work to end and its memory to come back (CudaAlloc::SynchronousAfterWorkBeside); a plain
// synchronous one fails at once.
TEST_CASE("CudaDevicePtr_AfterWorkBesideWaitsForItsMemory", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");
    if (gpu_is_integrated())
        SKIP("Fills the card, which on an integrated GPU is the host's memory");
    std::vector<CudaDevicePtr<uint8_t>> held;
    const auto fill = [&] {
        for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2)
            while (true) {
                try {
                    held.emplace_back(chunk, CudaAlloc::Synchronous);
                } catch (const JFJochException &) {
                    break;
                }
            }
    };
    constexpr size_t BYTES = size_t(64) << 20;
    {
        std::optional<GPUWorkBeside> work(std::in_place);
        fill();
        CHECK_THROWS_AS(CudaDevicePtr<uint8_t>(BYTES, CudaAlloc::Synchronous), JFJochException);
        std::thread ends([&] {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
            held.clear();
            work.reset();
        });
        CHECK_NOTHROW(CudaDevicePtr<uint8_t>(BYTES, CudaAlloc::SynchronousAfterWorkBeside));
        ends.join();
    }
    cuda_trim_allocation_pool();
}

// The free memory waited on is cudaMemGetInfo's on a discrete card. On one that shares the host's memory
// (GB10, Jetson) it also counts the page cache the host hands back, less a tenth of the card for the host:
// never less than cudaMemGetInfo's. Allocations between the readings move them, hence the 1 GB slack.
TEST_CASE("GetGpuFreeMemory_IntegratedCountsThePageCache", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    constexpr double SLACK = 1e9;
    size_t free_bytes = 0, total_bytes = 0;
    REQUIRE(get_gpu_free_memory(free_bytes, total_bytes));
    size_t cuda_free = 0, cuda_total = 0;
    REQUIRE(cudaMemGetInfo(&cuda_free, &cuda_total) == cudaSuccess);
    CHECK(total_bytes == cuda_total);
    CHECK(free_bytes <= total_bytes);
    CHECK(double(free_bytes) >= double(cuda_free) - SLACK);
    if (!gpu_is_integrated())
        CHECK(double(free_bytes) == Catch::Approx(double(cuda_free)).margin(SLACK));
}


// The pool's indexers go back to the card - not only to the allocation pool - when it is asked, and the
// next frame builds them again.
TEST_CASE("IndexerThreadPool_ReleaseIndexersGivesTheCardBack", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT).IndexingThreads(1);
    DiffractionExperiment x;
    x.ImportIndexingSettings(settings);
    const std::vector<Coord> recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}};

    // What the allocation pools of all cards hand out: an indexer worker takes its card round-robin.
    const auto pooled = [] {
        uint64_t sum = 0;
        for (int d = 0; d < get_gpu_count(); d++) {
            cudaMemPool_t pool = nullptr;
            uint64_t used = 0;
            if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess
                && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrUsedMemCurrent, &used) == cudaSuccess)
                sum += used;
        }
        return sum;
    };
    const auto reserved = [] {
        uint64_t sum = 0;
        for (int d = 0; d < get_gpu_count(); d++) {
            cudaMemPool_t pool = nullptr;
            uint64_t r = 0;
            if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess
                && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemCurrent, &r) == cudaSuccess)
                sum += r;
        }
        return sum;
    };
    const uint64_t before = pooled();
    const uint64_t reserved_before = reserved();
    {
        IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse);
        const IndexerResult first = pool.Run(x, recip);
        REQUIRE(first.error == nullptr);
        REQUIRE(pooled() > before + (uint64_t(100) << 20));

        pool.ReleaseIndexers();
        CHECK(pooled() <= before + (uint64_t(16) << 20));
        CHECK(reserved() <= reserved_before + (uint64_t(64) << 20));

        const IndexerResult again = pool.Run(x, recip);
        CHECK(again.error == nullptr);
        CHECK(again.executed == first.executed);
        CHECK(again.lattice.size() == first.lattice.size());
    }
    for (int d = 0; d < get_gpu_count(); d++) {
        ScopedGpuPin pin(d);
        cuda_trim_allocation_pool();
    }
}

// jfjoch_broker keeps the FFT indexers from one rugnux pass and collection to the next: a pool's worker
// borrows the one kept on its card instead of building its own, gives it back when the pool ends, and
// the card gets the memory back only with the last FFTIndexerReuse.
TEST_CASE("IndexerThreadPool_FFTIndexerKeptAcrossPools", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT).IndexingThreads(1);
    DiffractionExperiment x;
    x.ImportIndexingSettings(settings);
    const std::vector<Coord> recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}};

    const auto pooled = [] {
        uint64_t sum = 0;
        for (int d = 0; d < get_gpu_count(); d++) {
            cudaMemPool_t pool = nullptr;
            uint64_t used = 0;
            if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess
                && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrUsedMemCurrent, &used) == cudaSuccess)
                sum += used;
        }
        return sum;
    };

    IndexerResult fresh;
    {
        IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse);
        fresh = pool.Run(x, recip);
        REQUIRE(fresh.error == nullptr);
    }

    const uint64_t before = pooled();
    {
        FFTIndexerReuse reuse;
        for (int d = 0; d < get_gpu_count(); d++) {
            ScopedGpuPin pin(d);
            KeepFFTIndexers(settings, 1);
        }
        const uint64_t kept = pooled();
        REQUIRE(kept > before + (uint64_t(100) << 20));
        for (int i = 0; i < 2; i++) {
            IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse);
            const IndexerResult r = pool.Run(x, recip);
            CHECK(pooled() == kept);
            CHECK(r.error == nullptr);
            CHECK(r.executed == fresh.executed);
            CHECK(r.lattice.size() == fresh.lattice.size());
        }
        CHECK(pooled() == kept);

        // A settings change in the broker: the kept indexers go while the reuse scope lives on
        ReleaseKeptFFTIndexers();
        CHECK(pooled() <= before + (uint64_t(16) << 20));
    }
    CHECK(pooled() <= before + (uint64_t(16) << 20));
    for (int d = 0; d < get_gpu_count(); d++) {
        ScopedGpuPin pin(d);
        cuda_trim_allocation_pool();
    }
}

#endif

// A merge that finds the card full waits for the GPU work running beside it: until the last one
// ends, and not forever.
TEST_CASE("GPUWorkBeside_WaitEndsWithTheWork", "[CUDAMemHelpers]") {
    CHECK(wait_for_gpu_work_beside(std::chrono::seconds(0)));

    std::optional<GPUWorkBeside> work;
    work.emplace();
    CHECK_FALSE(wait_for_gpu_work_beside(std::chrono::seconds(0)));

    std::thread ends([&work] {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        work.reset();
    });
    CHECK(wait_for_gpu_work_beside(std::chrono::seconds(60)));
    ends.join();
}

// The speculative geometry probe waits for the main line's merge until the merge has said how much of the
// card it holds, and then only where the two do not fit within 80% of the card's total; the rest of the
// work beside never waits for a merge. With no card there is nothing to wait for.
TEST_CASE("GPUWorkBeside_AfterMergesWaitsWhereTheBytesDoNotFit", "[CUDAMemHelpers]") {
    const bool card = get_gpu_count() > 0;
    // A probe started on a thread of its own; true once it is past its wait.
    const auto start_probe = [](size_t bytes, std::atomic<bool> &started) {
        return std::thread([bytes, &started] {
            GPUWorkBeside beside(/*after_merges=*/true, bytes);
            started = true;
        });
    };

    std::optional<GPUMerge> merge(std::in_place);
    { GPUWorkBeside other; }   // does not wait

    // Not said yet: it waits, and starts once the merge says it holds nothing.
    std::atomic<bool> started = false;
    std::thread probe = start_probe(1, started);
    if (card) {
        std::this_thread::sleep_for(std::chrono::milliseconds(200));
        CHECK_FALSE(started);
    }
    merge->Holds(0);
    probe.join();
    CHECK(started);

    // The merge holds the whole card: it waits until the merge ends.
    merge->Holds(get_gpu_total_memory(0));
    started = false;
    probe = start_probe(1, started);
    if (card) {
        std::this_thread::sleep_for(std::chrono::milliseconds(200));
        CHECK_FALSE(started);
    }
    merge.reset();
    probe.join();
    CHECK(started);
}