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"

// 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() == 4096);
    CHECK(cuda_release_pinned_pool() == 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));
}

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