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

#pragma once

#include <cstdint>
#include <string>
#include <vector>

int32_t get_gpu_count();

// 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();
void set_gpu(int32_t dev_id);

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

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