From 6a2f00362eed16cfcc792d8fc716c19fa9108ea4 Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Tue, 11 Aug 2020 09:50:22 -0700
Subject: [PATCH] Add TpuTopologyExternal::version() and
 TpuVersionEnumToString().

PiperOrigin-RevId: 326040796
Change-Id: I3033d21271db465e4d6c57f5a3fe3c0c8913d1aa
---
 tensorflow/core/tpu/tpu_library_init_fns.inc      |  2 ++
 tensorflow/stream_executor/tpu/c_api_decl.h       |  6 ++++++
 .../stream_executor/tpu/tpu_executor_c_api.h      |  2 ++
 tensorflow/stream_executor/tpu/tpu_topology.cc    | 15 +++++++++++++++
 tensorflow/stream_executor/tpu/tpu_topology.h     |  3 +++
 5 files changed, 28 insertions(+)

diff --git a/tensorflow/core/tpu/tpu_library_init_fns.inc b/tensorflow/core/tpu/tpu_library_init_fns.inc
index a27dbea2388..9ac4fb9ec6d 100644
--- a/tensorflow/core/tpu/tpu_library_init_fns.inc
+++ b/tensorflow/core/tpu/tpu_library_init_fns.inc
@@ -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);
diff --git a/tensorflow/stream_executor/tpu/c_api_decl.h b/tensorflow/stream_executor/tpu/c_api_decl.h
index bca5f254ad1..a7b4c372e18 100644
--- a/tensorflow/stream_executor/tpu/c_api_decl.h
+++ b/tensorflow/stream_executor/tpu/c_api_decl.h
@@ -31,6 +31,12 @@ enum TpuCoreTypeEnum {
   kEmbeddingV2,
 };
 
+enum TpuVersionEnum {
+  kUnknownTpuVersion,
+  kTpuV2,
+  kTpuV3,
+};
+
 typedef struct SE_Status SE_Status;
 
 typedef struct SE_Platform SE_Platform;
diff --git a/tensorflow/stream_executor/tpu/tpu_executor_c_api.h b/tensorflow/stream_executor/tpu/tpu_executor_c_api.h
index c498244cc6e..149a00615a9 100644
--- a/tensorflow/stream_executor/tpu/tpu_executor_c_api.h
+++ b/tensorflow/stream_executor/tpu/tpu_executor_c_api.h
@@ -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);
diff --git a/tensorflow/stream_executor/tpu/tpu_topology.cc b/tensorflow/stream_executor/tpu/tpu_topology.cc
index 6c885b229ec..c86b399b34e 100644
--- a/tensorflow/stream_executor/tpu/tpu_topology.cc
+++ b/tensorflow/stream_executor/tpu/tpu_topology.cc
@@ -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
diff --git a/tensorflow/stream_executor/tpu/tpu_topology.h b/tensorflow/stream_executor/tpu/tpu_topology.h
index 07e9afc7d81..5219ba7017b 100644
--- a/tensorflow/stream_executor/tpu/tpu_topology.h
+++ b/tensorflow/stream_executor/tpu/tpu_topology.h
@@ -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