xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/device_description.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2015 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 "tensorflow/compiler/xla/stream_executor/device_description.h"
17 
18 #include <algorithm>
19 
20 #include "absl/strings/str_cat.h"
21 #include "tensorflow/compiler/xla/stream_executor/lib/human_readable.h"
22 #include "tensorflow/compiler/xla/stream_executor/lib/mathutil.h"
23 
24 namespace stream_executor {
25 
26 static const uint64_t kUninitializedUint64 = -1ULL;
27 /* static */ const char *DeviceDescription::kUndefinedString = "<undefined>";
28 
DeviceDescription()29 DeviceDescription::DeviceDescription()
30     : device_vendor_(kUndefinedString),
31       platform_version_(kUndefinedString),
32       driver_version_(kUndefinedString),
33       runtime_version_(kUndefinedString),
34       pci_bus_id_(kUndefinedString),
35       name_(kUndefinedString),
36       thread_dim_limit_(kUninitializedUint64, kUninitializedUint64,
37                         kUninitializedUint64),
38       block_dim_limit_(kUninitializedUint64, kUninitializedUint64,
39                        kUninitializedUint64),
40       threads_per_core_limit_(kUninitializedUint64),
41       threads_per_block_limit_(kUninitializedUint64),
42       threads_per_warp_(kUninitializedUint64),
43       registers_per_core_limit_(kUninitializedUint64),
44       registers_per_block_limit_(kUninitializedUint64),
45       device_address_bits_(kUninitializedUint64),
46       device_memory_size_(kUninitializedUint64),
47       memory_bandwidth_(kUninitializedUint64),
48       shared_memory_per_core_(kUninitializedUint64),
49       shared_memory_per_block_(kUninitializedUint64),
50       clock_rate_ghz_(-1.0),
51       numa_node_(-1),
52       core_count_(-1),
53       ecc_enabled_(false) {}
54 
ToMap() const55 std::unique_ptr<std::map<std::string, std::string>> DeviceDescription::ToMap()
56     const {
57   std::unique_ptr<std::map<std::string, std::string>> owned_result{
58       new std::map<std::string, std::string>};
59   std::map<std::string, std::string> &result = *owned_result;
60   result["Device Vendor"] = device_vendor();
61   result["Platform Version"] = platform_version();
62   result["Driver Version"] = driver_version();
63   result["Runtime Version"] = runtime_version();
64   result["PCI bus ID"] = pci_bus_id_;
65   result["Device Name"] = name_;
66 
67   const ThreadDim &thread_dim = thread_dim_limit();
68   result["ThreadDim Limit"] =
69       absl::StrCat(thread_dim.x, ",", thread_dim.y, ",", thread_dim.z);
70   const BlockDim &block_dim = block_dim_limit();
71   result["BlockDim Limit"] =
72       absl::StrCat(block_dim.x, ",", block_dim.y, ",", block_dim.z);
73 
74   result["Threads Per Core Limit"] = absl::StrCat(threads_per_core_limit());
75   result["Threads Per Block Limit"] = absl::StrCat(threads_per_block_limit());
76   result["Registers Per Block Limit"] =
77       absl::StrCat(registers_per_block_limit());
78 
79   result["Device Address Bits"] = absl::StrCat(device_address_bits());
80   result["Device Memory Size"] =
81       port::HumanReadableNumBytes::ToString(device_memory_size());
82   result["Memory Bandwidth"] = absl::StrCat(
83       port::HumanReadableNumBytes::ToString(memory_bandwidth_), "/s");
84 
85   result["Shared Memory Per Core"] =
86       port::HumanReadableNumBytes::ToString(shared_memory_per_core_);
87   result["Shared Memory Per Block"] =
88       port::HumanReadableNumBytes::ToString(shared_memory_per_block_);
89 
90   result["Clock Rate GHz"] = absl::StrCat(clock_rate_ghz());
91 
92   result["CUDA Compute Capability"] = cuda_compute_capability().ToString();
93 
94   result["AMDGPU GCN Arch Name"] = rocm_compute_capability().gcn_arch_name();
95 
96   result["NUMA Node"] = absl::StrCat(numa_node());
97   result["Core Count"] = absl::StrCat(core_count());
98   result["ECC Enabled"] = absl::StrCat(ecc_enabled());
99   return owned_result;
100 }
101 
102 namespace internal {
103 
DeviceDescriptionBuilder()104 DeviceDescriptionBuilder::DeviceDescriptionBuilder()
105     : device_description_(new DeviceDescription) {}
106 
107 }  // namespace internal
108 
cuda_compute_capability() const109 CudaComputeCapability DeviceDescription::cuda_compute_capability() const {
110   return cuda_compute_capability_;
111 }
112 
rocm_compute_capability() const113 RocmComputeCapability DeviceDescription::rocm_compute_capability() const {
114   return rocm_compute_capability_;
115 }
116 
ThreadDimOk(const DeviceDescription & device_description,const ThreadDim & thread_dim)117 bool ThreadDimOk(const DeviceDescription &device_description,
118                  const ThreadDim &thread_dim) {
119   const int64_t total_threads = thread_dim.x * thread_dim.y * thread_dim.z;
120   const int64_t threads_per_block_limit =
121       device_description.threads_per_block_limit();
122   if (total_threads > threads_per_block_limit) {
123     VLOG(2) << "exceeded total-thread-per-block limit: " << total_threads
124             << " vs limit " << threads_per_block_limit;
125     return false;
126   }
127 
128   const auto &limit = device_description.thread_dim_limit();
129   bool ok = thread_dim.x <= limit.x && thread_dim.y <= limit.y &&
130             thread_dim.z <= limit.z;
131   if (!ok) {
132     VLOG(2) << "thread dim " << thread_dim.ToString()
133             << " exceeds limit constraints of " << limit.ToString();
134   }
135   return ok;
136 }
137 
DivideCeil(uint64 x,uint64 y)138 uint64_t DivideCeil(uint64 x, uint64 y) {
139   return port::MathUtil::CeilOfRatio(x, y);
140 }
141 
CalculateDimensionality(const DeviceDescription & device_description,int64_t element_count,int64_t * threads_per_block,int64_t * block_count)142 void CalculateDimensionality(const DeviceDescription &device_description,
143                              int64_t element_count, int64_t *threads_per_block,
144                              int64_t *block_count) {
145   *threads_per_block = device_description.threads_per_block_limit();
146   *block_count = port::MathUtil::CeilOfRatio(element_count, *threads_per_block);
147   if (*block_count == 1) {
148     CHECK_LE(element_count, *threads_per_block);
149     *threads_per_block = element_count;
150   }
151 }
152 
153 }  // namespace stream_executor
154