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

Add TpuTopologyExternal::version() and TpuVersionEnumToString().

Browse Source 
PiperOrigin-RevId: 326040796
Change-Id: I3033d21271db465e4d6c57f5a3fe3c0c8913d1aa
This commit is contained in:
Skye Wanderman-Milne 2020-08-11 09:50:22 -07:00 committed by TensorFlower Gardener
parent 0f35ef2abc
commit 6a2f00362e
5 changed files with 28 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/tpu/tpu_library_init_fns.inc
View File

@ -178,6 +178,8 @@ tensorflow::Status SetExecutorStructFn(void* library_handle) {
TFTPU_SET_FN(executor_fn, TpuTopology_NumCores);
TFTPU_SET_FN(executor_fn, TpuTopology_Cores);
TFTPU_SET_FN(executor_fn, TpuTopology_IdForHost);
TFTPU_SET_FN(executor_fn, TpuTopology_Version);
TFTPU_SET_FN(executor_fn, TpuCoreLocation_ChipCoordinates);
TFTPU_SET_FN(executor_fn, TpuCoreLocation_HostCoordinates);
TFTPU_SET_FN(executor_fn, TpuCoreLocation_Index);

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

@ -31,6 +31,12 @@ enum TpuCoreTypeEnum {
kEmbeddingV2,
};
enum TpuVersionEnum {
kUnknownTpuVersion,
kTpuV2,
kTpuV3,
};
typedef struct SE_Status SE_Status;
typedef struct SE_Platform SE_Platform;

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

@ -211,6 +211,7 @@ void TpuTopology_Cores(SE_TpuTopology* tpu_topology,
TpuCoreTypeEnum tpu_core_type,
SE_TpuTopology_Core** cores);
int TpuTopology_IdForHost(SE_TpuTopology* tpu_topology, int x, int y, int z);
TpuVersionEnum TpuTopology_Version(SE_TpuTopology* tpu_topology);
void TpuCoreLocation_ChipCoordinates(SE_TpuTopology_Core* tpu_core_location,
int* x, int* y, int* z);
void TpuCoreLocation_HostCoordinates(SE_TpuTopology_Core* tpu_core_location,
@ -367,6 +368,7 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_NumCores);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_Cores);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_IdForHost);
TFTPU_ADD_FN_IN_STRUCT(TpuTopology_Version);
TFTPU_ADD_FN_IN_STRUCT(TpuCoreLocation_ChipCoordinates);
TFTPU_ADD_FN_IN_STRUCT(TpuCoreLocation_HostCoordinates);

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

@ -95,5 +95,20 @@ int TpuTopologyExternal::IdForHost(TpuDimensionsExternal host) const {
host.y, host.z);
}
TpuVersionEnum TpuTopologyExternal::version() const {
return tpu::ExecutorApiFn()->TpuTopology_VersionFn(topology_);
}
std::string TpuVersionEnumToString(TpuVersionEnum version) {
switch (version) {
case kUnknownTpuVersion:
return "Unknown TPU version";
case kTpuV2:
return "TPU v2";
case kTpuV3:
return "TPU v3";
}
}
} // namespace tpu
} // namespace tensorflow

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

@ -74,11 +74,14 @@ class TpuTopologyExternal {
int index) const;
std::vector<TpuCoreLocationExternal> cores(TpuCoreTypeEnum core_type) const;
int IdForHost(TpuDimensionsExternal host) const;
TpuVersionEnum version() const;
private:
SE_TpuTopology* topology_;
};
std::string TpuVersionEnumToString(TpuVersionEnum version);
} // namespace tpu
} // namespace tensorflow