experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add HostCount() / ChipsPerHost() API to TPUTopologyExternal.

Browse Source 
PiperOrigin-RevId: 335962331
Change-Id: I4d9a3f8770039348dfc77616333d80c68f493bda
This commit is contained in:
A. Unique TensorFlower 2020-10-07 15:11:13 -07:00 committed by TensorFlower Gardener
parent 0406aa2fcc
commit d1d4c51fcd
4 changed files with 20 additions and 1 deletions

5
tensorflow/stream_executor/tpu/BUILD
View File

@ -3,7 +3,10 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
default_visibility = ["//tensorflow/core/tpu:__subpackages__"],
default_visibility = [
"//learning/brain/experimental/dtensor:__subpackages__",
"//tensorflow/core/tpu:__subpackages__",
],
licenses = ["notice"], # Apache 2.0
)

6
tensorflow/stream_executor/tpu/tpu_executor_c_api.h
View File

@ -246,6 +246,9 @@ int TpuTopology_LogicalDevicesPerHost(SE_TpuTopology* tpu_topology,
TpuCoreTypeEnum tpu_core_type);
int TpuTopology_LogicalDevicesPerChip(SE_TpuTopology* tpu_topology,
TpuCoreTypeEnum tpu_core_type);
int TpuTopology_HostCount(SE_TpuTopology* tpu_topology);
int TpuTopology_ChipsPerHost(SE_TpuTopology* tpu_topology);
int TpuTopology_ChipBounds_X(SE_TpuTopology* tpu_topology);
int TpuTopology_ChipBounds_Y(SE_TpuTopology* tpu_topology);
int TpuTopology_ChipBounds_Z(SE_TpuTopology* tpu_topology);
@ -436,6 +439,9 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_LogicalDevicesPerHost);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_LogicalDevicesPerChip);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_HostCount);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_ChipsPerHost);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_ChipBounds_X);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_ChipBounds_Y);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_ChipBounds_Z);

8
tensorflow/stream_executor/tpu/tpu_topology.cc
View File

@ -73,6 +73,14 @@ int32 TpuTopologyExternal::LogicalDevicesPerChip(
core_type);
}
int32 TpuTopologyExternal::HostCount() const {
return tpu::ExecutorApiFn()->TpuTopology_HostCountFn(topology_);
}
int32 TpuTopologyExternal::ChipsPerHost() const {
return tpu::ExecutorApiFn()->TpuTopology_ChipsPerHostFn(topology_);
}
TpuTopologyChipBoundsExternal TpuTopologyExternal::chip_bounds() const {
return {tpu::ExecutorApiFn()->TpuTopology_ChipBounds_XFn(topology_),
tpu::ExecutorApiFn()->TpuTopology_ChipBounds_YFn(topology_),

2
tensorflow/stream_executor/tpu/tpu_topology.h
View File

@ -71,6 +71,8 @@ class TpuTopologyExternal {
: topology_(topology) {}
int32 LogicalDevicesPerHost(TpuCoreTypeEnum core_type) const;
int32 LogicalDevicesPerChip(TpuCoreTypeEnum core_type) const;
int32 HostCount() const;
int32 ChipsPerHost() const;
TpuTopologyChipBoundsExternal chip_bounds() const;
bool HasChip(int x, int y, int z) const;
TpuCoreLocationExternal Core(int x, int y, int z, TpuCoreTypeEnum core_type,