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

// NUMA node the given GPU is attached to, read from sysfs; -1 when unknown, when there is no such
// device, or on a platform where the question does not apply.
int get_gpu_numa_node(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().
//
// With JFJOCH_PIN_CPU_TO_GPU_NODE set in the environment it also confines the thread to the cores of
// that GPU's NUMA node, so the memory it allocates lands beside the GPU it feeds. Off by default -
// it also takes the other sockets' cores away from the thread, and which way that trade goes depends
// on the machine.
void pin_gpu();