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

Add a TpuTopologyExternal::CoreForId method.

Browse Source 
PiperOrigin-RevId: 354581402
Change-Id: I7dc8f1b987e9f90a4f7d2d751de573f8b01eeb70
This commit is contained in:
Wenhao Jia 2021-01-29 12:02:09 -08:00 committed by TensorFlower Gardener
parent 1a46fdc4a2
commit 597741a69d
5 changed files with 20 additions and 8 deletions

2
tensorflow/core/tpu/graph_rewrite/distributed_tpu_rewrite_pass.cc
View File

@ -1482,7 +1482,7 @@ static Status ParseDeviceAssignmentAttr(
") are not valid for the current TPU topology");
}
tpu::TpuCoreLocationExternal core_location =
tpu_topology.Core(x, y, z, kTensorCore, core);
tpu_topology.Core(kTensorCore, x, y, z, core);
if (replica_assignment(x, y, z, core) != -1) {
return errors::InvalidArgument("Duplicate coordinates (", x, ",", y,

1
tensorflow/core/tpu/tpu_executor_init_fns.inc
View File

@ -119,6 +119,7 @@ tensorflow::Status SetExecutorStructFn(void* library_handle) {
TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Y);
TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Z);
TFTPU_SET_FN(executor_fn, TpuTopology_HasChip);
TFTPU_SET_FN(executor_fn, TpuTopology_CoreForId);
TFTPU_SET_FN(executor_fn, TpuTopology_Core);
TFTPU_SET_FN(executor_fn, TpuTopology_NumCores);
TFTPU_SET_FN(executor_fn, TpuTopology_Cores);

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

@ -251,9 +251,12 @@ int TpuTopology_ChipBounds_X(SE_TpuTopology* tpu_topology);
int TpuTopology_ChipBounds_Y(SE_TpuTopology* tpu_topology);
int TpuTopology_ChipBounds_Z(SE_TpuTopology* tpu_topology);
bool TpuTopology_HasChip(SE_TpuTopology* tpu_topology, int x, int y, int z);
SE_TpuTopology_Core* TpuTopology_Core(SE_TpuTopology* tpu_topology, int x,
int y, int z,
TpuCoreTypeEnum tpu_core_type, int index);
SE_TpuTopology_Core* TpuTopology_CoreForId(SE_TpuTopology* tpu_topology,
TpuCoreTypeEnum tpu_core_type,
int id);
SE_TpuTopology_Core* TpuTopology_Core(SE_TpuTopology* tpu_topology,
TpuCoreTypeEnum tpu_core_type, int x,
int y, int z, int index);
int TpuTopology_NumCores(SE_TpuTopology* tpu_topology,
TpuCoreTypeEnum tpu_core_type);
// 'cores' should be a preallocated array of size TpuTopology_NumCores.
@ -457,6 +460,7 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_ChipBounds_Y);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_ChipBounds_Z);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_HasChip);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_CoreForId);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_Core);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_NumCores);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_Cores);

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

@ -91,11 +91,17 @@ bool TpuTopologyExternal::HasChip(int x, int y, int z) const {
return tpu::ExecutorApiFn()->TpuTopology_HasChipFn(topology_, x, y, z);
}
TpuCoreLocationExternal TpuTopologyExternal::Core(int x, int y, int z,
TpuCoreTypeEnum core_type,
TpuCoreLocationExternal TpuTopologyExternal::CoreForId(
TpuCoreTypeEnum core_type, int id) const {
return TpuCoreLocationExternal(
tpu::ExecutorApiFn()->TpuTopology_CoreForIdFn(topology_, core_type, id));
}
TpuCoreLocationExternal TpuTopologyExternal::Core(TpuCoreTypeEnum core_type,
int x, int y, int z,
int index) const {
return TpuCoreLocationExternal(tpu::ExecutorApiFn()->TpuTopology_CoreFn(
topology_, x, y, z, core_type, index));
topology_, core_type, x, y, z, index));
}
std::vector<TpuCoreLocationExternal> TpuTopologyExternal::cores(

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

@ -75,7 +75,8 @@ class TpuTopologyExternal {
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,
TpuCoreLocationExternal CoreForId(TpuCoreTypeEnum core_type, int id) const;
TpuCoreLocationExternal Core(TpuCoreTypeEnum core_type, int x, int y, int z,
int index) const;
std::vector<TpuCoreLocationExternal> cores(TpuCoreTypeEnum core_type) const;
int IdForHost(TpuDimensionsExternal host) const;