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

#pragma once

#include <cstdint>

int32_t get_gpu_count();
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();