1 /* Copyright 2020 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/stream_executor/tpu/tpu_topology.h"
17
18 #include "tensorflow/core/tpu/tpu_api.h"
19
20 namespace tensorflow {
21 namespace tpu {
22
chip_coordinates() const23 TpuDimensionsExternal TpuCoreLocationExternal::chip_coordinates() const {
24 int x, y, z;
25 tpu::ExecutorApiFn()->TpuCoreLocation_ChipCoordinatesFn(core_location_, &x,
26 &y, &z);
27 return {x, y, z};
28 }
29
host_coordinates() const30 TpuDimensionsExternal TpuCoreLocationExternal::host_coordinates() const {
31 int x, y, z;
32 tpu::ExecutorApiFn()->TpuCoreLocation_HostCoordinatesFn(core_location_, &x,
33 &y, &z);
34 return {x, y, z};
35 }
36
index() const37 int32 TpuCoreLocationExternal::index() const {
38 return tpu::ExecutorApiFn()->TpuCoreLocation_IndexFn(core_location_);
39 }
40
Id() const41 int32 TpuCoreLocationExternal::Id() const {
42 return tpu::ExecutorApiFn()->TpuCoreLocation_IdFn(core_location_);
43 }
44
Id() const45 int32 TpuHostLocationExternal::Id() const {
46 return tpu::ExecutorApiFn()->TpuHostLocation_IdFn(host_location_);
47 }
48
Cores(TpuCoreTypeEnum core_type) const49 std::vector<TpuCoreLocationExternal> TpuHostLocationExternal::Cores(
50 TpuCoreTypeEnum core_type) const {
51 int num_cores = tpu::ExecutorApiFn()->TpuHostLocation_NumCoresFn(
52 host_location_, core_type);
53 std::vector<SE_TpuTopology_Core*> core_ptrs(num_cores);
54 tpu::ExecutorApiFn()->TpuHostLocation_CoresFn(host_location_, core_type,
55 core_ptrs.data());
56 std::vector<TpuCoreLocationExternal> result;
57 result.reserve(num_cores);
58 for (SE_TpuTopology_Core* ptr : core_ptrs) {
59 result.emplace_back(ptr);
60 }
61 return result;
62 }
63
LogicalDevicesPerHost(TpuCoreTypeEnum core_type) const64 int32 TpuTopologyExternal::LogicalDevicesPerHost(
65 TpuCoreTypeEnum core_type) const {
66 return tpu::ExecutorApiFn()->TpuTopology_LogicalDevicesPerHostFn(topology_,
67 core_type);
68 }
69
LogicalDevicesPerChip(TpuCoreTypeEnum core_type) const70 int32 TpuTopologyExternal::LogicalDevicesPerChip(
71 TpuCoreTypeEnum core_type) const {
72 return tpu::ExecutorApiFn()->TpuTopology_LogicalDevicesPerChipFn(topology_,
73 core_type);
74 }
75
HostCount() const76 int32 TpuTopologyExternal::HostCount() const {
77 return tpu::ExecutorApiFn()->TpuTopology_HostCountFn(topology_);
78 }
79
ChipsPerHost() const80 int32 TpuTopologyExternal::ChipsPerHost() const {
81 return tpu::ExecutorApiFn()->TpuTopology_ChipsPerHostFn(topology_);
82 }
83
chip_bounds() const84 TpuTopologyChipBoundsExternal TpuTopologyExternal::chip_bounds() const {
85 return {tpu::ExecutorApiFn()->TpuTopology_ChipBounds_XFn(topology_),
86 tpu::ExecutorApiFn()->TpuTopology_ChipBounds_YFn(topology_),
87 tpu::ExecutorApiFn()->TpuTopology_ChipBounds_ZFn(topology_)};
88 }
89
HasChip(int x,int y,int z) const90 bool TpuTopologyExternal::HasChip(int x, int y, int z) const {
91 return tpu::ExecutorApiFn()->TpuTopology_HasChipFn(topology_, x, y, z);
92 }
93
CoreForId(TpuCoreTypeEnum core_type,int id) const94 TpuCoreLocationExternal TpuTopologyExternal::CoreForId(
95 TpuCoreTypeEnum core_type, int id) const {
96 return TpuCoreLocationExternal(
97 tpu::ExecutorApiFn()->TpuTopology_CoreForIdFn(topology_, core_type, id));
98 }
99
Core(TpuCoreTypeEnum core_type,int x,int y,int z,int index) const100 TpuCoreLocationExternal TpuTopologyExternal::Core(TpuCoreTypeEnum core_type,
101 int x, int y, int z,
102 int index) const {
103 return TpuCoreLocationExternal(tpu::ExecutorApiFn()->TpuTopology_CoreFn(
104 topology_, core_type, x, y, z, index));
105 }
106
cores(TpuCoreTypeEnum core_type) const107 std::vector<TpuCoreLocationExternal> TpuTopologyExternal::cores(
108 TpuCoreTypeEnum core_type) const {
109 int num_cores =
110 tpu::ExecutorApiFn()->TpuTopology_NumCoresFn(topology_, core_type);
111 std::vector<SE_TpuTopology_Core*> core_ptrs(num_cores);
112 tpu::ExecutorApiFn()->TpuTopology_CoresFn(topology_, core_type,
113 core_ptrs.data());
114 std::vector<TpuCoreLocationExternal> result;
115 result.reserve(num_cores);
116 for (SE_TpuTopology_Core* ptr : core_ptrs) {
117 result.emplace_back(ptr);
118 }
119 return result;
120 }
121
IdForHost(TpuDimensionsExternal host) const122 int TpuTopologyExternal::IdForHost(TpuDimensionsExternal host) const {
123 return tpu::ExecutorApiFn()->TpuTopology_IdForHostFn(topology_, host.x,
124 host.y, host.z);
125 }
126
version() const127 TpuVersionEnum TpuTopologyExternal::version() const {
128 return tpu::ExecutorApiFn()->TpuTopology_VersionFn(topology_);
129 }
130
TpuVersionEnumToString(TpuVersionEnum version)131 std::string TpuVersionEnumToString(TpuVersionEnum version) {
132 switch (version) {
133 case kUnknownTpuVersion:
134 return "Unknown TPU version";
135 case kTpuV2:
136 return "TPU v2";
137 case kTpuV3:
138 return "TPU v3";
139 case kTpuV4:
140 return "TPU v4";
141 }
142 }
143
144 } // namespace tpu
145 } // namespace tensorflow
146