Snippet to find number of GPUs on system and pick one for CUDA use

cudaDevices.cu

/* get number of GPUs on local host */
int getGPUCount(){
  int ans=0;
  if(cudaGetDeviceCount(&ans) != cudaSuccess){
    return 0;
  }
  return ans;
}

/* for multithreaded or mpi tasks, pick a GPU 
 * without conflict. Return unique GPU id for 
 * a given rank, or -1 if error

 * rank:  MPI task or thread ID
 * ngpus: number of GPUs on local host
 */
int pickGPU(int rank, int ngpus){
  int gpuID = rank%ngpus; //any better way?
  if(cudaSetDevice(gpuID) != cudaSuccess){
    return -1;
  }
  return gpuID;
}