1 | /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | #include <memory> |
17 | |
18 | #include "tensorflow/core/grappler/devices.h" |
19 | #include "tensorflow/core/platform/byte_order.h" |
20 | #include "tensorflow/core/platform/cpu_info.h" |
21 | |
22 | #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
23 | #include "tensorflow/core/common_runtime/gpu/gpu_init.h" |
24 | #include "tensorflow/core/platform/stream_executor.h" |
25 | #endif // GOOGLE_CUDA |
26 | |
27 | namespace tensorflow { |
28 | namespace grappler { |
29 | |
30 | int GetNumAvailableGPUs( |
31 | const std::pair<int, int>& min_cuda_compute_capability) { |
32 | int num_eligible_gpus = 0; |
33 | |
34 | #if TENSORFLOW_USE_ROCM |
35 | if (min_cuda_compute_capability.first != 0 || |
36 | min_cuda_compute_capability.second != 0) { |
37 | LOG(ERROR) << "GetNumAvailableGPUs() should receive zero " |
38 | "min_cuda_compute_capability" ; |
39 | return 0; |
40 | } |
41 | #endif |
42 | #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
43 | if (ValidateGPUMachineManager().ok()) { |
44 | se::Platform* gpu_manager = GPUMachineManager(); |
45 | if (gpu_manager != nullptr) { |
46 | int num_gpus = gpu_manager->VisibleDeviceCount(); |
47 | for (int i = 0; i < num_gpus; i++) { |
48 | #if GOOGLE_CUDA |
49 | auto desc = gpu_manager->DescriptionForDevice(i); |
50 | if (desc.ok()) { |
51 | int min_gpu_core_count = 8; |
52 | if ((*desc)->core_count() >= min_gpu_core_count && |
53 | (*desc)->cuda_compute_capability().IsAtLeast( |
54 | min_cuda_compute_capability.first, |
55 | min_cuda_compute_capability.second)) { |
56 | num_eligible_gpus++; |
57 | } |
58 | } |
59 | #else |
60 | num_eligible_gpus++; |
61 | #endif |
62 | } |
63 | } |
64 | } |
65 | #if GOOGLE_CUDA |
66 | LOG(INFO) |
67 | << "Number of eligible GPUs (core count >= 8, compute capability >= " |
68 | << min_cuda_compute_capability.first << "." |
69 | << min_cuda_compute_capability.second << "): " << num_eligible_gpus; |
70 | #else |
71 | LOG(INFO) << "Number of eligible GPUs: " << num_eligible_gpus; |
72 | #endif |
73 | |
74 | #else // GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
75 | LOG(INFO) |
76 | << "Number of eligible GPUs (core count >= 8, compute capability >= " |
77 | << min_cuda_compute_capability.first << "." |
78 | << min_cuda_compute_capability.second << "): " << num_eligible_gpus |
79 | << " (Note: TensorFlow was not compiled with CUDA or ROCm support)" ; |
80 | #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
81 | return num_eligible_gpus; |
82 | } |
83 | |
84 | int64_t AvailableGPUMemory(int gpu_id) { |
85 | #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
86 | // Look up the device, to see its attributes. |
87 | se::Platform* gpu_platform = GPUMachineManager(); |
88 | CHECK_LT(gpu_id, gpu_platform->VisibleDeviceCount()); |
89 | se::StreamExecutor* se = gpu_platform->ExecutorForDevice(gpu_id).value(); |
90 | int64_t total_memory, available_memory; |
91 | CHECK(se->DeviceMemoryUsage(&available_memory, &total_memory)); |
92 | |
93 | return available_memory; |
94 | #else |
95 | return 0; |
96 | #endif |
97 | } |
98 | |
99 | int GetNumAvailableLogicalCPUCores() { return port::NumSchedulableCPUs(); } |
100 | |
101 | } // end namespace grappler |
102 | } // end namespace tensorflow |
103 | |