#include <cuda.h>
#include <cuda_runtime.h>

void initialisation_cuda()
{
    char* local_rank_env;
    int local_rank;
    cudaError_t cudaRet;

     /* Récupération du rang local du processus via la variable d'environnement
        positionnée par Slurm, l'utilisation de MPI_Comm_rank n'étant pas encore
        possible puisque cette routine est utilisée AVANT l'initialisation de MPI */
    local_rank_env = getenv("SLURM_LOCALID");

    if (local_rank_env) {
        local_rank = atoi(local_rank_env);
        /* Définition du GPU à utiliser pour chaque processus MPI */
        cudaRet = cudaSetDevice(local_rank);
        if(cudaRet != CUDA_SUCCESS) {
            printf("Erreur: cudaSetDevice a échoué\n");
            exit(1);
        }
    } else {
        printf("Erreur : impossible de déterminer le rang local du processus\n");
        exit(1);
    }
}