TPU library internal change.
PiperOrigin-RevId: 315627327
Change-Id: Ib792cb2983f64f1e4a4a14a4acb96230d4a1f622
diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index d9d843b..6bfd7b6 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -287,6 +287,14 @@
)
cc_library(
+ name = "tpu_util_c_api_hdrs",
+ hdrs = ["tpu_util_c_api.h"],
+ deps = [
+ "//tensorflow/stream_executor/tpu:proto_helper",
+ ],
+)
+
+cc_library(
name = "tpu_util",
srcs = ["tpu_util.cc"],
hdrs = ["tpu_util.h"],
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
index 53e79aa..8b8ab3a 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
+++ b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
@@ -92,9 +92,6 @@
const char* data,
size_t size);
-// Checks if whether a TPU compilation is enabled.
-bool TpuCompile_IsTpuCompilationEnabled();
-
// Executes the computations using XLA TPU compiler and returns TPU programs
// ready for execution.
void TpuCompile_CompileAheadOfTime(
@@ -108,12 +105,6 @@
const XLA_TpuMeshState* mesh_state,
TpuSerializedProto* serialized_device_assignment, SE_Status* status);
-// Converts an XLA `Shape` into its equivalent TPU `Shape` representation.
-void TpuCompile_ToTpuShapeRepresentation(
- TpuSerializedProto serialized_xla_shape, int data_type,
- bool use_fast_memory, TpuSerializedProto* serialized_tensor_shape,
- SE_Status* status);
-
} // extern "C"
#endif // TENSORFLOW_CORE_TPU_KERNELS_TPU_COMPILE_C_API_H_
diff --git a/tensorflow/core/tpu/kernels/tpu_util_c_api.h b/tensorflow/core/tpu/kernels/tpu_util_c_api.h
new file mode 100644
index 0000000..4d99244
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_util_c_api.h
@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TPU_KERNELS_TPU_UTIL_C_API_H_
+#define TENSORFLOW_CORE_TPU_KERNELS_TPU_UTIL_C_API_H_
+
+#include "tensorflow/stream_executor/tpu/proto_helper.h"
+
+typedef struct SE_Status SE_Status;
+
+extern "C" {
+
+// Checks if whether a TPU compilation is enabled.
+bool TpuCompile_IsTpuCompilationEnabled();
+
+// Converts an XLA `Shape` into its equivalent TPU `Shape` representation.
+void TpuCompile_ToTpuShapeRepresentation(
+ TpuSerializedProto serialized_xla_shape, int data_type,
+ bool use_fast_memory, TpuSerializedProto* serialized_tensor_shape,
+ SE_Status* status);
+
+} // extern "C"
+
+#endif // TENSORFLOW_CORE_TPU_KERNELS_TPU_UTIL_C_API_H_
