Jungfraujoch/image_analysis/indexing/CUDAMemHelpers.h

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <cuda_runtime.h>
#include <cufft.h>
#include <algorithm>
#include <cstdio>
#include <map>
#include <mutex>
#include <stdexcept>
#include <string>
#include <vector>
#include "../common/JFJochException.h"
#include "../common/CUDAWrapper.h"

// The message for GPU memory that could not be had: what was asked (bytes, 0 where unknown), what the
// card has free and, where it is merely full, the usual causes and the remedy - the same advice the
// scaling's own out-of-memory gives. A context lost to another fault fails here too and keeps its reason.
inline std::string GPUMemoryFailure(const char *what, size_t bytes, bool out_of_memory, const char *reason) {
    char buf[512];
    if (bytes > 0)
        std::snprintf(buf, sizeof buf, "%s (%.1f MB) failed (%s)", what, bytes / 1e6, reason);
    else
        std::snprintf(buf, sizeof buf, "%s failed (%s)", what, reason);
    std::string msg(buf);
    if (!out_of_memory)
        return msg;
    size_t free_b = 0, total_b = 0;
    if (get_gpu_free_memory(free_b, total_b)) {
        std::snprintf(buf, sizeof buf, ", with %.1f of the %.1f GB card free", free_b / 1e9, total_b / 1e9);
        msg += buf;
    }
    return msg + ". Either another program is using this GPU, or this data set needs more GPU memory than "
                 "the card has - free the GPU and run again, or run on the CPU (rugnux --no-gpu "
                 "..., or a CPU-only build)";
}

class CudaStream {
    cudaStream_t stream_ = nullptr;
public:
    // Non-blocking by default: a stream created with cudaStreamDefault synchronises against the legacy
    // NULL stream, so any NULL-stream operation anywhere in the process serialises every worker's GPU
    // work against every other's. With one engine per worker thread that costs most of the parallelism.
    CudaStream(unsigned int flags = cudaStreamNonBlocking, int priority = 0) {
        if (cudaStreamCreateWithPriority(&stream_, flags, priority) != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  "Failed to create CUDA stream");
    }
    // The stream's queued work is waited for before it goes. cudaStreamDestroy alone returns at once and
    // lets that work run on, while the owner's buffers are freed next: a pooled buffer's free is ordered
    // on the allocation stream (CudaDevicePtr), not on this one, so its memory could be handed to
    // another engine, or returned to the driver, under a kernel or copy still writing it.
    ~CudaStream() {
        if (stream_) {
            cudaStreamSynchronize(stream_);
            cudaStreamDestroy(stream_);
        }
    }
    // Move-only type
    CudaStream(CudaStream&& other) noexcept : stream_(other.stream_) { other.stream_ = nullptr; }
    CudaStream& operator=(CudaStream&& other) noexcept {
        if (this != &other) {
            if (stream_) {
                cudaStreamSynchronize(stream_);
                cudaStreamDestroy(stream_);
            }
            stream_ = other.stream_;
            other.stream_ = nullptr;
        }
        return *this;
    }
    CudaStream(const CudaStream&) = delete;
    CudaStream& operator=(const CudaStream&) = delete;

    operator cudaStream_t() const { return stream_; }
    cudaStream_t get() const { return stream_; }
};

// A timing event, so a phase that is queued on a stream can still report how long the device spent
// on it. cudaEventDisableTiming is deliberately NOT used - timing is the whole point here.
class CudaEvent {
    cudaEvent_t event_ = nullptr;
public:
    CudaEvent() {
        if (cudaEventCreate(&event_) != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  "Failed to create CUDA event");
    }
    ~CudaEvent() {
        if (event_) cudaEventDestroy(event_);
    }
    CudaEvent(CudaEvent&& other) noexcept : event_(other.event_) { other.event_ = nullptr; }
    CudaEvent& operator=(CudaEvent&& other) noexcept {
        if (this != &other) {
            if (event_) cudaEventDestroy(event_);
            event_ = other.event_;
            other.event_ = nullptr;
        }
        return *this;
    }
    CudaEvent(const CudaEvent&) = delete;
    CudaEvent& operator=(const CudaEvent&) = delete;

    operator cudaEvent_t() const { return event_; }
    cudaEvent_t get() const { return event_; }
};

// How much freed memory the pool keeps rather than returning it to the driver. Returning it puts the
// next allocation straight back on the path this exists to avoid, so hold the per-worker engine
// buffers; but the card also has to fit the merge, which asks for several gigabytes of its own after
// the image loop, so do not hold everything.
constexpr uint64_t CUDA_MEM_POOL_RELEASE_THRESHOLD = 1ull << 30;

// The stream device allocations are ordered on: one per (thread, device), borrowed on first use and
// handed back when the thread exits. The streams themselves are never destroyed.
//
// cudaMalloc and cudaFree are on CUDA's implicit-synchronization list - each one synchronises the
// device across every stream - so one analysis engine per worker thread, each making a few dozen of
// them, serialises every worker against every other AND stalls the workers already processing
// images. cudaMallocAsync/cudaFreeAsync take the stream-ordered path instead and do not.
//
// The stream is deliberately not owned by the engine: cudaFreeAsync must be ordered on a stream that
// is still alive, and an engine's own stream can be destroyed before the buffers it allocated. It is
// keyed by device because a thread pinned to one GPU must not order work on another's stream.
//
// Nor is it owned by the thread, for the same reason: a buffer can outlive the thread that allocated
// it - the shadow finder builds its accumulator on a throw-away thread and frees it from another,
// long after - so the thread only borrows the stream. But it does hand it back. The broker starts
// fresh worker threads for every data collection, and a stream holds about half a megabyte of device
// memory, so a stream per thread ever started is a leak that a long-running broker does not survive.
// Handed back, the streams number as many as the threads that were ever alive at once.
namespace jfjoch_cuda_allocation_streams {
    struct Idle {
        std::mutex m;
        std::map<int, std::vector<CudaStream>> streams;   // by device
    };

    // Never destroyed: a thread can still be handing a stream back while the process exits, and
    // destroying streams during static destruction would call into a runtime that may be gone.
    inline Idle &idle() {
        static Idle *i = new Idle;
        return *i;
    }

    struct Borrowed {
        std::map<int, CudaStream> streams;   // by device
        ~Borrowed() {
            std::lock_guard lock(idle().m);
            for (auto &[device, stream] : streams)
                idle().streams[device].push_back(std::move(stream));
        }
    };
}

inline cudaStream_t cuda_allocation_stream() {
    int device = 0;
    if (cudaGetDevice(&device) != cudaSuccess)
        return nullptr;
    thread_local jfjoch_cuda_allocation_streams::Borrowed borrowed;
    const auto it = borrowed.streams.find(device);
    if (it != borrowed.streams.end())
        return it->second;

    auto &idle = jfjoch_cuda_allocation_streams::idle();
    {
        std::lock_guard lock(idle.m);
        auto &streams = idle.streams[device];
        if (!streams.empty()) {
            CudaStream stream = std::move(streams.back());
            streams.pop_back();
            return borrowed.streams.emplace(device, std::move(stream)).first->second;
        }
    }

    try {
        CudaStream stream;
        cudaMemPool_t pool = nullptr;
        if (cudaDeviceGetDefaultMemPool(&pool, device) == cudaSuccess) {
            uint64_t threshold = CUDA_MEM_POOL_RELEASE_THRESHOLD;
            cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);
        }
        return borrowed.streams.emplace(device, std::move(stream)).first->second;
    } catch (const JFJochException &) {
        // No stream: the caller allocates synchronously instead, so this is handled - see
        // CudaDevicePtr for why it must then not stay behind as the thread's last error.
        cudaGetLastError();
        return nullptr;
    }
}

// Return the current device's pool memory that no buffer holds to the driver. Frees ordered on a
// stream count only once the stream has reached them, hence the device-wide wait first.
inline void cuda_trim_allocation_pool() {
    int device = 0;
    cudaMemPool_t pool = nullptr;
    if (cudaGetDevice(&device) != cudaSuccess || cudaDeviceGetDefaultMemPool(&pool, device) != cudaSuccess) {
        cudaGetLastError();
        return;
    }
    cudaDeviceSynchronize();
    cudaMemPoolTrimTo(pool, 0);
    cudaGetLastError();
}

// The same on every card: for work that has ended and whose threads were pinned round-robin (pin_gpu).
inline void cuda_trim_allocation_pools() {
    for (int32_t d = 0; d < get_gpu_count(); d++) {
        ScopedGpuPin pin(d);
        cuda_trim_allocation_pool();
    }
}

// Whether `bytes` more device memory can be had on the current device, for a step whose buffers are
// sized by the data set and which ends the run when they do not fit. The free memory counts what the
// pool holds with no buffer in it (up to CUDA_MEM_POOL_RELEASE_THRESHOLD) as taken, so that is handed
// back before the step gives up; and GPU work running beside it in this process (GPUWorkBeside) is
// waited for, as its memory comes back when it ends. Only when the step runs depends on this, never
// what it computes. `free_bytes` is the last reading, for the message of a step that does not fit.
inline bool cuda_memory_available(size_t bytes, size_t &free_bytes) {
    const auto read_free = [&] {
        size_t total_bytes = 0;
        if (!get_gpu_free_memory(free_bytes, total_bytes))
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read free device memory");
        return bytes <= free_bytes;
    };
    if (read_free())
        return true;
    cuda_trim_allocation_pool();
    if (read_free())
        return true;
    // That work is a probe pass of seconds, so this long means it is stuck: the step fails as it would
    // have without it.
    wait_for_gpu_work_beside(std::chrono::minutes(10));
    cuda_trim_allocation_pool();
    return read_free();
}

// A pooled buffer's free is ordered on the allocation stream, not on the stream of the engine that
// used it, and the pool may hand the memory to another engine - or return it to the driver - as soon
// as that free is reached, which on an idle allocation stream is at once. So an engine with pooled
// buffers waits for its own stream before they go: in its destructor (a CudaStream it owns does that
// itself when it is destroyed before the buffers, i.e. declared after them) and before it reallocates.
// Otherwise a kernel or copy still queued writes memory that is someone else's, or unmapped.
//
// Which allocator a device buffer uses. Pooled is what every per-worker buffer wants. Synchronous is
// for a buffer that is freed by a thread other than the one whose work read it: cudaFreeAsync orders
// the free on the freeing side's allocation stream, which says nothing about kernels queued
// elsewhere, while cudaFree synchronises the whole device and therefore cannot be early.
// SynchronousAfterWorkBeside is Synchronous for a step of the main line - the merge's helpers - whose
// buffers may be asked for while a probe pass runs beside it on the same card (GPUWorkBeside): one that
// does not fit waits for that work to end and is asked for once more, as the merge's own arrays are
// (RotationScaleMergeGPU). Never for the work beside itself, which would wait for itself.
enum class CudaAlloc { Pooled, Synchronous, SynchronousAfterWorkBeside };

template <typename T>
class CudaDevicePtr {
    T* ptr_ = nullptr;
    // The stream the allocation was ordered on, and the one the free must be ordered on. Null when
    // the pool was not used, so the two always pair up.
    cudaStream_t stream_ = nullptr;

    void release() {
        if (!ptr_) return;
        if (stream_) cudaFreeAsync(ptr_, stream_);
        else         cudaFree(ptr_);
        ptr_ = nullptr;
    }

public:
    CudaDevicePtr() = default;

    explicit CudaDevicePtr(size_t count, CudaAlloc alloc = CudaAlloc::Pooled) {
        if (alloc == CudaAlloc::Pooled)
            stream_ = cuda_allocation_stream();
        if (stream_) {
            if (cudaMallocAsync(&ptr_, count * sizeof(T), stream_) == cudaSuccess) {
                // The allocation is ordered on this stream, so anything that wants to use the memory
                // from another stream has to be ordered after it. Waiting for it once here is that
                // ordering, and it leaves the pointer as freely usable as cudaMalloc's would have been.
                cudaStreamSynchronize(stream_);
                return;
            }
            // The pool failing is handled - by the synchronous allocation below - so it must not stay
            // behind as the thread's last error. Every kernel launch is followed by cudaGetLastError(),
            // which would report it there: an "out of memory" thrown over an image whose buffers were
            // all allocated and whose kernel launched.
            cudaGetLastError();
            stream_ = nullptr;
        }
        if (cudaMalloc(&ptr_, count * sizeof(T)) == cudaSuccess)
            return;
        // cudaMalloc cannot take memory the pool is holding: what pooled buffers handed back stays
        // reserved by the pool, up to gigabytes after the image loop's engines are gone, until the
        // pool is trimmed. Hand it back to the driver and try once more; the failed attempt must not
        // stay behind as the last error (see above).
        cudaGetLastError();
        cuda_trim_allocation_pool();
        cudaError_t err = cudaMalloc(&ptr_, count * sizeof(T));
        if (err == cudaErrorMemoryAllocation && alloc == CudaAlloc::SynchronousAfterWorkBeside) {
            cudaGetLastError();
            // That work is a probe pass of seconds, so this long means it is stuck.
            wait_for_gpu_work_beside(std::chrono::minutes(10));
            cuda_trim_allocation_pool();
            err = cudaMalloc(&ptr_, count * sizeof(T));
        }
        // With the reason: a context lost to another kernel's fault fails here too, and must not read
        // as a card that is merely full.
        if (err != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  GPUMemoryFailure("Allocating GPU memory", count * sizeof(T),
                                                   err == cudaErrorMemoryAllocation, cudaGetErrorString(err)));
    }
    ~CudaDevicePtr() {
        release();
    }
    // Move-only type
    CudaDevicePtr(CudaDevicePtr&& other) noexcept : ptr_(other.ptr_), stream_(other.stream_) {
        other.ptr_ = nullptr;
        other.stream_ = nullptr;
    }
    CudaDevicePtr& operator=(CudaDevicePtr&& other) noexcept {
        if (this != &other) {
            release();
            ptr_ = other.ptr_;
            stream_ = other.stream_;
            other.ptr_ = nullptr;
            other.stream_ = nullptr;
        }
        return *this;
    }
    CudaDevicePtr(const CudaDevicePtr&) = delete;
    CudaDevicePtr& operator=(const CudaDevicePtr&) = delete;

    T* get() const { return ptr_; }
    operator T*() const { return ptr_; }
};

// The transform a cuFFT plan is made for, as cufftMakePlanMany takes it; an empty embed is its nullptr.
struct CudaFFTShape {
    std::vector<int> n;
    std::vector<int> inembed;
    int istride = 1, idist = 0;
    std::vector<int> onembed;
    int ostride = 1, odist = 0;
    cufftType type = CUFFT_R2C;
    int batch = 1;
};

// While a CudaFFTPlanScope is alive, cuFFT plans are kept when their user is done with them, by device
// and shape, and lent to the next user of that shape; when the last scope ends they are destroyed, and
// outside any scope a plan is destroyed as soon as its user is done. Making a plan is costly and
// serialised: cuFFT holds a process-wide lock while it loads the plan's kernel modules (each plan loads
// its own, two plans of one shape included), which took over a second a plan on a busy card, and on
// Windows compiles them from PTX first; and it frees and loads them all again for the next plan of the
// same shape. A plan is never shared by two users at once - cuFFT is thread safe only across distinct
// plans.
//
// The scope is one processing run (rugnux), so nothing is kept from one run for the next, nor past the
// device reset at the end of a run. A kept plan holds no work area: the user supplies one (CudaFFTPlan),
// so an idle plan holds only its tables and modules.
namespace jfjoch_cufft_plans {
    struct Plan {
        cufftHandle handle = 0;
        size_t work_bytes = 0;
    };

    struct Idle {
        std::mutex m;
        std::map<std::vector<int>, std::vector<Plan>> plans;   // by key()
        int scopes = 0;
    };

    inline Idle &idle() {
        static Idle *i = new Idle;
        return *i;
    }

    // The current device and the shape.
    inline std::vector<int> key(const CudaFFTShape &s) {
        int device = 0;
        if (cudaGetDevice(&device) != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read the CUDA device");
        std::vector<int> k{device, static_cast<int>(s.n.size())};
        k.insert(k.end(), s.n.begin(), s.n.end());
        k.push_back(static_cast<int>(s.inembed.size()));
        k.insert(k.end(), s.inembed.begin(), s.inembed.end());
        k.push_back(static_cast<int>(s.onembed.size()));
        k.insert(k.end(), s.onembed.begin(), s.onembed.end());
        k.insert(k.end(), {s.istride, s.idist, s.ostride, s.odist, static_cast<int>(s.type), s.batch});
        return k;
    }

    inline Plan make(const CudaFFTShape &s) {
        Plan p;
        cufftResult r = cufftCreate(&p.handle);
        if (r == CUFFT_SUCCESS) {
            r = cufftSetAutoAllocation(p.handle, 0);
            if (r == CUFFT_SUCCESS)
                r = cufftMakePlanMany(p.handle, static_cast<int>(s.n.size()), const_cast<int *>(s.n.data()),
                                      s.inembed.empty() ? nullptr : const_cast<int *>(s.inembed.data()),
                                      s.istride, s.idist,
                                      s.onembed.empty() ? nullptr : const_cast<int *>(s.onembed.data()),
                                      s.ostride, s.odist, s.type, s.batch, &p.work_bytes);
            if (r != CUFFT_SUCCESS)
                cufftDestroy(p.handle);
        }
        // Said as a memory failure whatever the code: on a full card (two runs sharing it) the plan
        // failed as CUFFT_INTERNAL_ERROR, not CUFFT_ALLOC_FAILED.
        if (r != CUFFT_SUCCESS)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  GPUMemoryFailure("Creating the CUFFT plan (cufftMakePlanMany)", 0, true,
                                                   ("cuFFT error " + std::to_string(r)).c_str()));
        // The plan's setup can still be queued on the NULL stream when cufftMakePlanMany returns, and a
        // non-blocking stream (CudaStream) does not wait for it: a transform launched on one straight
        // away read the plan's tables before they were there - an illegal address in 3 of 4 runs where
        // a rotation pool built its FFT indexers at transform length 3072. Waiting here orders it.
        if (cudaStreamSynchronize(nullptr) != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  "Creating the CUFFT plan: setup on the NULL stream failed");
        return p;
    }

    // An idle plan of this key, or a new one.
    inline Plan borrow(const std::vector<int> &k, const CudaFFTShape &s) {
        {
            std::lock_guard lock(idle().m);
            auto it = idle().plans.find(k);
            if (it != idle().plans.end() && !it->second.empty()) {
                const Plan p = it->second.back();
                it->second.pop_back();
                return p;
            }
        }
        return make(s);
    }

    // On the plan's own device (the first element of its key).
    inline void destroy(const std::vector<int> &k, const Plan &p) {
        int current = 0;
        cudaGetDevice(&current);
        cudaSetDevice(k[0]);
        cufftDestroy(p.handle);
        cudaSetDevice(current);
    }

    inline void give_back(const std::vector<int> &k, const Plan &p) {
        {
            std::lock_guard lock(idle().m);
            if (idle().scopes > 0) {
                idle().plans[k].push_back(p);
                return;
            }
        }
        destroy(k, p);
    }

    // How many plans are kept, idle.
    inline size_t kept() {
        std::lock_guard lock(idle().m);
        size_t n = 0;
        for (const auto &[k, plans] : idle().plans)
            n += plans.size();
        return n;
    }
}

// See above: cuFFT plans are kept for reuse while one of these is alive, and destroyed when the last ends.
class CudaFFTPlanScope {
public:
    CudaFFTPlanScope() {
        std::lock_guard lock(jfjoch_cufft_plans::idle().m);
        jfjoch_cufft_plans::idle().scopes++;
    }
    ~CudaFFTPlanScope() {
        std::map<std::vector<int>, std::vector<jfjoch_cufft_plans::Plan>> plans;
        {
            std::lock_guard lock(jfjoch_cufft_plans::idle().m);
            if (--jfjoch_cufft_plans::idle().scopes == 0)
                plans.swap(jfjoch_cufft_plans::idle().plans);
        }
        for (const auto &[k, list] : plans)
            for (const auto &p : list)
                jfjoch_cufft_plans::destroy(k, p);
    }
    CudaFFTPlanScope(const CudaFFTPlanScope &) = delete;
    CudaFFTPlanScope &operator=(const CudaFFTPlanScope &) = delete;
};

// A cuFFT plan of `shape` on the current device, borrowed from the plans kept above (or made) for as long
// as this lives, its transforms on `stream`. Its work area is `work_area` (at least WorkBytes() long), or, given
// none, one of its own. The plan goes back once the stream has finished with it.
class CudaFFTPlan {
    std::vector<int> key_;
    jfjoch_cufft_plans::Plan plan_;
    cudaStream_t stream_ = nullptr;
    CudaDevicePtr<char> work_;

    void release() {
        if (!plan_.handle)
            return;
        cudaStreamSynchronize(stream_);
        work_ = CudaDevicePtr<char>();
        jfjoch_cufft_plans::give_back(key_, plan_);
        plan_ = {};
    }

public:
    CudaFFTPlan() = default;

    explicit CudaFFTPlan(const CudaFFTShape &shape, cudaStream_t stream = nullptr, void *work_area = nullptr)
            : key_(jfjoch_cufft_plans::key(shape)), plan_(jfjoch_cufft_plans::borrow(key_, shape)), stream_(stream) {
        try {
            if (work_area == nullptr) {
                work_ = CudaDevicePtr<char>(std::max<size_t>(plan_.work_bytes, 1));
                work_area = work_;
            }
            if (cufftSetWorkArea(plan_.handle, work_area) != CUFFT_SUCCESS
                || cufftSetStream(plan_.handle, stream_) != CUFFT_SUCCESS)
                throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Setting up the CUFFT plan failed");
        } catch (...) {
            release();
            throw;
        }
    }

    // The work area a plan of `shape` on the current device needs - read from the kept plan, which is
    // made here if there is none, so that within a scope the plan asked is the plan used.
    static size_t WorkBytes(const CudaFFTShape &shape) {
        const std::vector<int> k = jfjoch_cufft_plans::key(shape);
        const jfjoch_cufft_plans::Plan p = jfjoch_cufft_plans::borrow(k, shape);
        jfjoch_cufft_plans::give_back(k, p);
        return p.work_bytes;
    }

    ~CudaFFTPlan() {
        release();
    }
    // Move-only type
    CudaFFTPlan(CudaFFTPlan&& other) noexcept
            : key_(std::move(other.key_)), plan_(other.plan_), stream_(other.stream_), work_(std::move(other.work_)) {
        other.plan_ = {};
    }
    CudaFFTPlan& operator=(CudaFFTPlan&& other) noexcept {
        if (this != &other) {
            release();
            key_ = std::move(other.key_);
            plan_ = other.plan_;
            stream_ = other.stream_;
            work_ = std::move(other.work_);
            other.plan_ = {};
        }
        return *this;
    }
    CudaFFTPlan(const CudaFFTPlan&) = delete;
    CudaFFTPlan& operator=(const CudaFFTPlan&) = delete;

    size_t WorkBytes() const { return plan_.work_bytes; }
    operator cufftHandle() const { return plan_.handle; }
    cufftHandle get() const { return plan_.handle; }
};

// Page-locked host memory is slow to get and to give back: cudaHostAlloc and cudaFreeHost map and lock
// every page under the process's memory-map lock, so the calls of all threads queue behind each other -
// measured at 1-3 ms a call on a workstation and 20-30 ms on a four-socket host. The engines that hold it
// are built per worker per step, which made that thousands of calls a run. So a freed buffer is kept, in a
// power-of-two size class, for the next engine to take. Up to CUDA_PINNED_POOL_LIMIT bytes are kept, a
// buffer larger than CUDA_PINNED_POOL_MAX_BUFFER (a whole frame) is not kept at all, and
// cuda_release_pinned_pool hands everything back at the end of a run.
//
// A buffer taken from the pool holds whatever its last owner left in it. And it goes back to the pool the
// moment its owner frees it, so - as with a pooled device buffer (CudaDevicePtr) - the owner waits for
// its stream before it does: a copy still queued into the buffer would write someone else's memory.
constexpr size_t CUDA_PINNED_POOL_LIMIT = 512ull << 20;
constexpr size_t CUDA_PINNED_POOL_MAX_BUFFER = 64ull << 20;

namespace jfjoch_cuda_pinned_pool {
    struct Pool {
        std::mutex m;
        std::map<size_t, std::vector<void *>> free;   // by size class
        size_t kept_bytes = 0;
    };

    // Never destroyed, as the allocation streams above: a buffer can be freed during static destruction.
    inline Pool &pool() {
        static Pool *p = new Pool;
        return *p;
    }

    // What is actually allocated for `bytes`: the next power of two from 4 kB, so that an engine whose
    // buffers grew a little still finds them; and exactly `bytes` for a buffer that is never kept.
    inline size_t allocated_bytes(size_t bytes) {
        if (bytes > CUDA_PINNED_POOL_MAX_BUFFER)
            return bytes;
        size_t c = 4096;
        while (c < bytes)
            c *= 2;
        return c;
    }
}

inline void *cuda_pinned_alloc(size_t bytes) {
    using namespace jfjoch_cuda_pinned_pool;
    const size_t size = allocated_bytes(bytes);
    {
        std::lock_guard lock(pool().m);
        auto it = pool().free.find(size);
        if (it != pool().free.end() && !it->second.empty()) {
            void *ptr = it->second.back();
            it->second.pop_back();
            pool().kept_bytes -= size;
            return ptr;
        }
    }
    // Portable: a buffer freed by a worker on one card is taken next by a worker on another.
    void *ptr = nullptr;
    if (cudaHostAlloc(&ptr, size, cudaHostAllocPortable) != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "Failed to allocate pinned host memory");
    return ptr;
}

inline void cuda_pinned_free(void *ptr, size_t bytes) noexcept {
    using namespace jfjoch_cuda_pinned_pool;
    const size_t size = allocated_bytes(bytes);
    if (size <= CUDA_PINNED_POOL_MAX_BUFFER) {
        std::lock_guard lock(pool().m);
        if (pool().kept_bytes + size <= CUDA_PINNED_POOL_LIMIT) {
            pool().free[size].push_back(ptr);
            pool().kept_bytes += size;
            return;
        }
    }
    cudaFreeHost(ptr);
}

// Hand every kept buffer back to the driver. Returns how many bytes that was.
inline size_t cuda_release_pinned_pool() {
    using namespace jfjoch_cuda_pinned_pool;
    std::map<size_t, std::vector<void *>> free;
    size_t bytes = 0;
    {
        std::lock_guard lock(pool().m);
        free.swap(pool().free);
        bytes = pool().kept_bytes;
        pool().kept_bytes = 0;
    }
    for (auto &[size, ptrs] : free)
        for (void *ptr : ptrs)
            cudaFreeHost(ptr);
    return bytes;
}

template <typename T>
class CudaHostPtr {
    T* ptr_ = nullptr;
    size_t count_ = 0;
public:
    CudaHostPtr() = default;
    explicit CudaHostPtr(size_t count)
        : ptr_(static_cast<T *>(cuda_pinned_alloc(count * sizeof(T)))), count_(count) {}
    ~CudaHostPtr() {
        if (ptr_) cuda_pinned_free(ptr_, count_ * sizeof(T));
    }
    // Move-only type
    CudaHostPtr(CudaHostPtr&& other) noexcept : ptr_(other.ptr_), count_(other.count_) { other.ptr_ = nullptr; }
    CudaHostPtr& operator=(CudaHostPtr&& other) noexcept {
        if (this != &other) {
            if (ptr_) cuda_pinned_free(ptr_, count_ * sizeof(T));
            ptr_ = other.ptr_;
            count_ = other.count_;
            other.ptr_ = nullptr;
        }
        return *this;
    }
    CudaHostPtr(const CudaHostPtr&) = delete;
    CudaHostPtr& operator=(const CudaHostPtr&) = delete;

    T* get() const { return ptr_; }
    operator T*() const { return ptr_; }
};

// A std::vector in pooled pinned memory, for the host side of an engine's per-frame copies: the copies
// are genuinely asynchronous, and the memory comes from the pool instead of being registered and
// unregistered with every engine (see cuda_pinned_alloc).
template <typename T>
struct CudaPinnedAllocator {
    using value_type = T;
    CudaPinnedAllocator() = default;
    template <typename U> CudaPinnedAllocator(const CudaPinnedAllocator<U> &) {}
    T *allocate(size_t n) { return static_cast<T *>(cuda_pinned_alloc(n * sizeof(T))); }
    void deallocate(T *p, size_t n) noexcept { cuda_pinned_free(p, n * sizeof(T)); }
    template <typename U> bool operator==(const CudaPinnedAllocator<U> &) const { return true; }
};

template <typename T>
using PinnedVector = std::vector<T, CudaPinnedAllocator<T>>;

template <typename T>
class CudaRegisteredVector {
    std::vector<T>* vec_ = nullptr;
    bool registered_ = false;

    static void registerPtr(void* ptr, size_t bytes, unsigned int flags) {
        cudaError_t err = cudaHostRegister(ptr, bytes, flags);
        if (err != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "cudaHostRegister failed");
    }
    static void unregisterPtr(void* ptr) {
        cudaError_t err = cudaHostUnregister(ptr);
        if (err != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "cudaHostUnregister failed");
    }
    // Unchecked, for the destructor and the move-assignment: both are noexcept, so throwing out of
    // them aborts the process instead of reporting the failure - and teardown is exactly where
    // cudaHostUnregister fails (after a device reset, or while another exception is unwinding).
    // The other destructors in this header ignore their teardown status for the same reason.
    static void unregisterPtrNoThrow(void* ptr) noexcept {
        cudaHostUnregister(ptr);
    }

public:
    // Non-owning wrapper. Does NOT provide accessors to the vector.
    CudaRegisteredVector() = default;

    CudaRegisteredVector(std::vector<T>& vec, unsigned int flags = cudaHostRegisterDefault)
        : vec_(&vec)
    {
        if (!vec.empty()) {
            registerPtr(vec.data(), vec.size() * sizeof(T), flags);
            registered_ = true;
        }
    }

    ~CudaRegisteredVector() {
        if (registered_ && vec_ && !vec_->empty()) {
            unregisterPtrNoThrow(vec_->data());
        }
    }

    // Move-only
    CudaRegisteredVector(CudaRegisteredVector&& other) noexcept
        : vec_(other.vec_), registered_(other.registered_) {
        other.vec_ = nullptr;
        other.registered_ = false;
    }

    CudaRegisteredVector& operator=(CudaRegisteredVector&& other) noexcept {
        if (this != &other) {
            // Clean current registration
            if (registered_ && vec_ && !vec_->empty()) {
                unregisterPtrNoThrow(vec_->data());
            }
            vec_ = other.vec_;
            registered_ = other.registered_;
            other.vec_ = nullptr;
            other.registered_ = false;
        }
        return *this;
    }

    CudaRegisteredVector(const CudaRegisteredVector&) = delete;
    CudaRegisteredVector& operator=(const CudaRegisteredVector&) = delete;

    // Re-register after vector capacity/size change. Caller must ensure
    // the vector is not registered at the moment of mutation.
    void rebind(std::vector<T>& vec, unsigned int flags = cudaHostRegisterDefault) {
        // Unregister previous if needed
        if (registered_ && vec_ && !vec_->empty()) {
            unregisterPtr(vec_->data());
        }
        vec_ = &vec;
        if (!vec.empty()) {
            registerPtr(vec.data(), vec.size() * sizeof(T), flags);
            registered_ = true;
        } else {
            registered_ = false;
        }
    }

    // Explicit unregister (optional).
    void unregister() {
        if (registered_ && vec_ && !vec_->empty()) {
            unregisterPtr(vec_->data());
            registered_ = false;
        }
    }

    bool isRegistered() const { return registered_; }
};