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

#pragma once

#include <chrono>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>

int32_t get_gpu_count();

// From here on get_gpu_count() is 0 without asking CUDA, so nothing in the process creates a context or
// touches a device and every GPU path takes its CPU path, as on a machine without a GPU: what
// CUDA_VISIBLE_DEVICES= does, asked for by the program (rugnux --no-gpu). Call it before anything uses
// CUDA; there is no way back.
void hide_gpus();
// Whether hide_gpus() was called - for a GPU path that is not CUDA's (the Metal integration engine), so
// --no-gpu keeps the process off every GPU.
bool gpus_are_hidden();

// Names of the visible GPUs, in device order and one entry per device, so repeated cards repeat.
// Empty without CUDA and on a machine with no device, which is also what get_gpu_count() == 0 says.
std::vector<std::string> get_gpu_names();

// The same list collapsed for a person: "4x NVIDIA A100-SXM4-80GB", or several such groups separated
// by ", " on a mixed machine. Empty when no GPU is visible.
std::string get_gpu_description();

// The device's total memory in bytes - not what is free at the moment. 0 without CUDA, or when the
// device cannot be queried.
size_t get_gpu_total_memory(int32_t dev_id);

// Free and total memory of the current device in bytes, as cudaMemGetInfo gives them - except on a GPU
// that shares the host's memory (cudaDevAttrIntegrated: GB10, Jetson), whose cudaMemGetInfo free is the
// host's MemFree and leaves out the page cache the kernel hands back on demand. There, on Linux, free is
// the host's MemAvailable less a tenth of the total, at most the total. For waiting on memory and for
// messages only: what runs on the GPU is decided on the total (get_gpu_total_memory), never on this.
// False without CUDA, or when the device cannot be queried.
bool get_gpu_free_memory(size_t &free_bytes, size_t &total_bytes);

void set_gpu(int32_t dev_id);

// Whether kernels on the current device read ordinary (pageable) host memory directly and at memory
// speed, so a host buffer can be handed to a kernel instead of being copied to the card first: true on
// a unified-memory system where the GPU walks the CPU's page tables (DGX Spark GB10, Grace Hopper GH200,
// GB200), false on a discrete card. Asked of the driver once per device. JFJOCH_FORCE_DIRECT_HOST_READ
// in the environment makes it true on every device - for testing the direct path on a discrete card,
// where only page-locked host memory is then readable. false without CUDA.
bool gpu_direct_host_read();

// Pin the calling thread to the next GPU in round-robin order, using a process-wide counter
// (counter++ % get_gpu_count()). Call once per thread; no thread id needed. No-op when no GPU
// is visible. Honours CUDA_VISIBLE_DEVICES via get_gpu_count().
void pin_gpu();

// From here on, pin_gpu() also keeps the calling thread on the CPUs of the NUMA node its GPU is
// attached to (ThreadAffinity.h) - on a machine with one node, or without CUDA, nothing changes.
// Off unless a program asks for it.
void enable_gpu_numa_binding();

// The NUMA node each visible GPU is attached to, in device order; -1 where it is not known or the
// machine has a single node. Empty without CUDA or a device.
std::vector<int> get_gpu_numa_nodes();

// pin_gpu() onto a given device: set_gpu(dev_id), and the NUMA binding above where it is enabled.
// For a worker that takes its card by index rather than round-robin.
void pin_gpu(int32_t dev_id);

// pin_gpu(dev_id) for as long as the object lives; then the calling thread gets back the device and
// the CPUs it had before. For a thread that drives a card only for a while - the merging engine, whose
// calls come from a thread that goes on to other work - where a pin for the rest of its life would
// carry the card's node onto everything that thread does, and starts, afterwards.
class ScopedGpuPin {
public:
    explicit ScopedGpuPin(int32_t dev_id);
    ~ScopedGpuPin();
    ScopedGpuPin(const ScopedGpuPin &) = delete;
    ScopedGpuPin &operator=(const ScopedGpuPin &) = delete;
private:
    int32_t prev_device = -1;   // -1: nothing to give back
    std::vector<int> prev_cpus;  // empty: the CPUs were not changed
};

// Have every GPU's host threads BLOCK in a synchronisation (cudaDeviceScheduleBlockingSync) instead
// of spinning on a core until the device finishes. Must be called before anything creates a CUDA
// context; it creates none itself. No-op without CUDA.
void set_gpu_blocking_sync();

// Create the CUDA context of every device from first_device on, one thread per device, and return
// when they all exist - for a caller that runs it beside other work, so that the first use of a card
// does not wait for its context. No-op without CUDA.
void create_gpu_contexts(int32_t first_device);

// Destroy every device's CUDA context, one thread per device, for a program about to exit with no
// CUDA work left: at exit the runtime destroys them one after another, a fifth of a second per card.
// No-op without CUDA.
void release_gpu_contexts();

// Drop the error CUDA has recorded for the calling thread. Call it where a CUDA failure has been
// HANDLED - a device route that fell back to the host, an indexing attempt whose failure was turned
// into a result - because the error otherwise stays as the thread's last error and the next
// cuda_err(cudaGetLastError()) after some later kernel launch reports it, over work that went fine.
// A sticky error (an illegal access, say) is not cleared by this, and nothing here pretends it is:
// the context is gone in that case and every later call fails on its own. No-op without CUDA.
void cuda_clear_error();

// Throw if the CUDA context is lost - a sticky error (an illegal access, say) that every later call
// on this device will fail with. For a device route about to fall back to the host: there is no
// fallback from that, and falling back only moves the failure somewhere less clear. A handled,
// non-sticky failure passes. No-op without CUDA.
void cuda_throw_if_context_lost();

// GPU work that runs beside the main line of a process and gives all of its device memory back when it
// ends: a probe pass run on a copy of the run while the run itself scales and merges. Hold one for as
// long as that work runs. On a small card it waits first for the main line's image passes (GPUImagePass)
// and for other work beside to end. With after_merges it waits also for the main line's merges
// (GPUMerge) - unless `bytes`, what this work holds of the card at its most, and what the merges said
// they hold together fit within 80% of the card's total (on a card that pages what does not fit, like a
// WDDM card, the rest is the margin). Decided on those bytes and the total, never on what is free.
// Build-independent - it only counts.
class GPUWorkBeside {
public:
    explicit GPUWorkBeside(bool after_merges = false, size_t bytes = 0);
    ~GPUWorkBeside();
    GPUWorkBeside(const GPUWorkBeside &) = delete;
    GPUWorkBeside &operator=(const GPUWorkBeside &) = delete;
};

// Wait until no GPUWorkBeside is held, for at most `timeout`. True once none is - at once if none was.
bool wait_for_gpu_work_beside(std::chrono::seconds timeout);

// The image pass of a run's main line - not of a probe holding a GPUWorkBeside: the stretch that holds
// the most of the card, and that never waits for a probe, so a probe that does not fit beside it can
// wait for it to end. Hold one for as long as the image pass runs. On a card too small for both, it
// waits first for the work beside to end. Build-independent - it only counts.
class GPUImagePass {
public:
    GPUImagePass();
    ~GPUImagePass();
    GPUImagePass(const GPUImagePass &) = delete;
    GPUImagePass &operator=(const GPUImagePass &) = delete;
};

// The scaling, merging and model validation of a run's main line - not of a probe. A probe that the main
// line never waits for (the speculative geometry probe) waits for it where the two do not fit the card
// together: on a card that pages what does not fit rather than failing the allocation, the two would not
// fail but crawl. Until Holds says how much of the card the merge holds, it counts as not leaving room for
// any. It waits for nothing itself, so nothing that the main line waits for inside its merge - the
// rotation-scale walk's probes - can wait in a circle with it. Build-independent - it only counts.
class GPUMerge {
    std::optional<size_t> bytes_;
public:
    GPUMerge();
    ~GPUMerge();
    GPUMerge(const GPUMerge &) = delete;
    GPUMerge &operator=(const GPUMerge &) = delete;
    // What the merge holds of the card from now on, at its most (an estimate fitted as an upper bound).
    void Holds(size_t bytes);
};

// Whether the card is too small for a probe pass beside the main line's image pass (12 GB, decided on the
// first card's total memory), so that they take turns (GPUWorkBeside). Build-independent.
bool gpu_takes_turns();

// Wait until no GPUImagePass is held - at once if none was. With no timeout: an image pass waits for no
// probe, so it ends.
void wait_for_gpu_image_passes();