1 /* Copyright 2021 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/core/util/autotune_maps/autotune_maps_utils.h"
17
18 #include <string>
19
20 #include "absl/strings/str_format.h"
21 #include "tensorflow/core/platform/errors.h"
22 #include "tensorflow/core/platform/hash.h"
23
24 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
25 #include "tensorflow/stream_executor/gpu/gpu_driver.h"
26 #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
27
28 namespace tensorflow {
29
30 namespace autotune_maps_utils {
31
32 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
33 namespace {
34
35 using ::stream_executor::gpu::GpuDeviceHandle;
36 using ::stream_executor::gpu::GpuDriver;
37
DeviceIdToIdentifierHelper(int device_id)38 StatusOr<string> DeviceIdToIdentifierHelper(int device_id) {
39 GpuDeviceHandle device;
40 TF_RETURN_IF_ERROR(GpuDriver::GetDevice(device_id, &device));
41 std::string device_name;
42 TF_RETURN_IF_ERROR(GpuDriver::GetDeviceName(device, &device_name));
43 int cc_major;
44 int cc_minor;
45 TF_RETURN_IF_ERROR(
46 GpuDriver::GetComputeCapability(&cc_major, &cc_minor, device));
47
48 uint64 device_memory_size = -1;
49 if (!GpuDriver::GetDeviceTotalMemory(device, &device_memory_size)) {
50 return errors::Internal("Failed to get device's total memory");
51 }
52
53 TF_ASSIGN_OR_RETURN(int core_count,
54 GpuDriver::GetMultiprocessorCount(device));
55 return absl::StrFormat("%s sm_%d.%d with %dB RAM and %d cores", device_name,
56 cc_major, cc_minor, device_memory_size, core_count);
57 }
58
59 } // namespace
60
GetDeviceIdToIdentifierMap()61 std::vector<std::string> GetDeviceIdToIdentifierMap() {
62 int device_count = GpuDriver::GetDeviceCount();
63 std::vector<string> map(device_count);
64 for (int device_id = 0; device_id < device_count; device_id++) {
65 StatusOr<string> device_identifier_or_status =
66 DeviceIdToIdentifierHelper(device_id);
67 if (device_identifier_or_status.ok()) {
68 map[device_id] = device_identifier_or_status.ValueOrDie();
69 } else {
70 map[device_id] = "Unknown Graphics Device";
71 }
72 }
73 return map;
74 }
75
DeviceIdToIdentifier(int device_id)76 std::string DeviceIdToIdentifier(int device_id) {
77 // Ensure the static variable is trivially destructible and thus safe to be
78 // destruct in multi-thread setting.
79 static const auto& map =
80 *new std::vector<string>(GetDeviceIdToIdentifierMap());
81 return device_id < map.size() ? map[device_id] : "Unknown Graphics Device";
82 }
83 #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
84
SerializeProtoDeterministic(const protobuf::Message & proto)85 std::string SerializeProtoDeterministic(const protobuf::Message& proto) {
86 std::string proto_serialized_string;
87 protobuf::io::StringOutputStream string_stream(&proto_serialized_string);
88 protobuf::io::CodedOutputStream stream(&string_stream);
89 // Ensure the serialization is deterministic so that equal ConvParameters
90 // have equal serialized strings and therefore equal hash codes.
91 stream.SetSerializationDeterministic(true);
92 proto.SerializeToCodedStream(&stream);
93 return proto_serialized_string;
94 }
95
HashProto(const protobuf::Message & proto)96 uint64 HashProto(const protobuf::Message& proto) {
97 return Hash64(SerializeProtoDeterministic(proto));
98 }
99
100 } // namespace autotune_maps_utils
101 } // namespace tensorflow
102