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