34 #include <cuda_runtime.h>
47 bool CompareCudaDevice(
const cudaDeviceProp& d1,
const cudaDeviceProp& d2) {
48 bool result = (d1.major > d2.major) ||
49 ((d1.major == d2.major) && (d1.minor > d2.minor)) ||
50 ((d1.major == d2.major) && (d1.minor == d2.minor) &&
51 (d1.multiProcessorCount > d2.multiProcessorCount));
58 int num_cuda_devices = 0;
59 cudaError_t error = cudaGetDeviceCount(&num_cuda_devices);
60 if (error != cudaSuccess) {
64 return num_cuda_devices;
69 CHECK_GT(num_cuda_devices, 0) <<
"No CUDA devices available";
71 int selected_gpu_index = -1;
73 selected_gpu_index = gpu_index;
75 std::vector<cudaDeviceProp> all_devices(num_cuda_devices);
76 for (
int device_id = 0; device_id < num_cuda_devices; ++device_id) {
77 cudaGetDeviceProperties(&all_devices[device_id], device_id);
79 std::sort(all_devices.begin(), all_devices.end(), CompareCudaDevice);
81 cudaChooseDevice(&selected_gpu_index, all_devices.data()));
82 VLOG(2) <<
"Found " << num_cuda_devices <<
" CUDA device(s), "
83 <<
"selected device " << selected_gpu_index <<
" with name "
84 << all_devices[selected_gpu_index].name;
87 CHECK_GE(selected_gpu_index, 0);
88 CHECK_LT(selected_gpu_index, num_cuda_devices)
89 <<
"Invalid CUDA GPU selected";
91 cudaDeviceProp device;
92 cudaGetDeviceProperties(&device, selected_gpu_index);
#define CUDA_SAFE_CALL(error)
void SetBestCudaDevice(const int gpu_index)