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 <optional>
#include <thread>

#ifdef JFJOCH_USE_CUDA

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

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

// 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);
}

// 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");

    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}};

    // The whole card, held by the work beside.
    std::optional<std::vector<CudaDevicePtr<uint8_t>>> held;
    const auto fill = [&held] {
        held.emplace();
        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;
                }
            }
    };
    IndexerThreadPool probe(settings, IndexerConstruction::OnFirstUse, /*beside=*/true);
    IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse);

    fill();
    {
        GPUWorkBeside work;
        CHECK(probe.Run(x, recip).error != nullptr);
    }
    {
        std::optional<GPUImagePass> image_pass(std::in_place);
        std::thread ends([&] {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
            held.reset();
            image_pass.reset();
        });
        GPUWorkBeside work;
        CHECK(probe.Run(x, recip).error == nullptr);
        ends.join();
    }

    fill();
    {
        std::optional<GPUWorkBeside> work(std::in_place);
        std::thread ends([&] {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
            held.reset();
            work.reset();
        });
        CHECK(pool.Run(x, recip).error == nullptr);
        ends.join();
    }
}

#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();
}