| 28 | namespace grappler { |
| 29 | |
| 30 | int GetNumAvailableGPUs( |
| 31 | const std::pair<int, int>& min_cuda_compute_capability) { |
| 32 | int num_eligible_gpus = 0; |
| 33 | #if GOOGLE_CUDA |
| 34 | if (ValidateGPUMachineManager().ok()) { |
| 35 | se::Platform* gpu_manager = GPUMachineManager(); |
| 36 | if (gpu_manager != nullptr) { |
| 37 | int num_gpus = gpu_manager->VisibleDeviceCount(); |
| 38 | for (int i = 0; i < num_gpus; i++) { |
| 39 | auto desc_status = gpu_manager->DescriptionForDevice(i); |
| 40 | if (desc_status.ok()) { |
| 41 | auto desc = desc_status.ConsumeValueOrDie(); |
| 42 | int cc_major = 0; |
| 43 | int cc_minor = 0; |
| 44 | desc->cuda_compute_capability(&cc_major, &cc_minor); |
| 45 | std::pair<int, int> cuda_compute_capability(cc_major, cc_minor); |
| 46 | int min_gpu_core_count = 8; |
| 47 | if (desc->core_count() >= min_gpu_core_count && |
| 48 | cuda_compute_capability >= min_cuda_compute_capability) { |
| 49 | num_eligible_gpus++; |
| 50 | } |
| 51 | } |
| 52 | } |
| 53 | } |
| 54 | } |
| 55 | LOG(INFO) |
| 56 | << "Number of eligible GPUs (core count >= 8, compute capability >= " |
| 57 | << min_cuda_compute_capability.first << "." |
| 58 | << min_cuda_compute_capability.second << "): " << num_eligible_gpus; |
| 59 | #else |
| 60 | LOG(INFO) |
| 61 | << "Number of eligible GPUs (core count >= 8, compute capability >= " |
| 62 | << min_cuda_compute_capability.first << "." |
| 63 | << min_cuda_compute_capability.second << "): " << num_eligible_gpus |
| 64 | << " (Note: TensorFlow was not compiled with CUDA support)"; |
| 65 | #endif // GOOGLE_CUDA |
| 66 | return num_eligible_gpus; |
| 67 | } |
| 68 | |
| 69 | int64 AvailableGPUMemory(int gpu_id) { |
| 70 | #if GOOGLE_CUDA |