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
GetNumAvailableGPUs(const std::pair<int,int> & min_cuda_compute_capability)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_status = gpu_manager->DescriptionForDevice(i);
50 if (desc_status.ok()) {
51 auto desc = desc_status.ConsumeValueOrDie();
52 int cc_major = 0;
53 int cc_minor = 0;
54 desc->cuda_compute_capability(&cc_major, &cc_minor);
55 std::pair<int, int> cuda_compute_capability(cc_major, cc_minor);
56 int min_gpu_core_count = 8;
57 if (desc->core_count() >= min_gpu_core_count &&
58 cuda_compute_capability >= min_cuda_compute_capability) {
59 num_eligible_gpus++;
60 }
61 }
62 #else
63 num_eligible_gpus++;
64 #endif
65 }
66 }
67 }
68 #if GOOGLE_CUDA
69 LOG(INFO)
70 << "Number of eligible GPUs (core count >= 8, compute capability >= "
71 << min_cuda_compute_capability.first << "."
72 << min_cuda_compute_capability.second << "): " << num_eligible_gpus;
73 #else
74 LOG(INFO) << "Number of eligible GPUs: " << num_eligible_gpus;
75 #endif
76
77 #else // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
78 LOG(INFO)
79 << "Number of eligible GPUs (core count >= 8, compute capability >= "
80 << min_cuda_compute_capability.first << "."
81 << min_cuda_compute_capability.second << "): " << num_eligible_gpus
82 << " (Note: TensorFlow was not compiled with CUDA or ROCm support)";
83 #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
84 return num_eligible_gpus;
85 }
86
AvailableGPUMemory(int gpu_id)87 int64 AvailableGPUMemory(int gpu_id) {
88 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
89 // Look up the device, to see its attributes.
90 se::Platform* gpu_platform = GPUMachineManager();
91 CHECK_LT(gpu_id, gpu_platform->VisibleDeviceCount());
92 se::StreamExecutor* se = gpu_platform->ExecutorForDevice(gpu_id).ValueOrDie();
93 int64 total_memory, available_memory;
94 CHECK(se->DeviceMemoryUsage(&available_memory, &total_memory));
95
96 return available_memory;
97 #else
98 return 0;
99 #endif
100 }
101
GetNumAvailableLogicalCPUCores()102 int GetNumAvailableLogicalCPUCores() { return port::NumSchedulableCPUs(); }
103
104 } // end namespace grappler
105 } // end namespace tensorflow
106