Set 'TPU_SHARED_LIBRARY_COMPILE_LINK 1' by default in the env where it has access to libtpu.so.
PiperOrigin-RevId: 293272729
Change-Id: Ief4c4fb01153334070a7496a3d0fe34c89c3adc4
diff --git a/tensorflow/compiler/xla/python/tpu_driver/BUILD b/tensorflow/compiler/xla/python/tpu_driver/BUILD
index 57246a2..08da1c2 100644
--- a/tensorflow/compiler/xla/python/tpu_driver/BUILD
+++ b/tensorflow/compiler/xla/python/tpu_driver/BUILD
@@ -74,26 +74,6 @@
)
cc_library(
- name = "direct_tpu_driver_local",
- srcs = ["direct_tpu_driver.cc"],
- defines = ["TPU_SHARED_LIBRARY_COMPILE_LINK"],
- deps = [
- ":tpu_driver",
- "@com_google_absl//absl/strings:str_format",
- "//tensorflow/compiler/xla:statusor",
- "//tensorflow/core/platform:logging",
- "//tensorflow/compiler/xla:status",
- "//tensorflow/compiler/xla:util",
- "//tensorflow/compiler/xla:xla_data_proto_cc",
- "//tensorflow/compiler/xla/service:hlo_proto_cc",
- ":tpu_service_proto_cc",
- ":tpu_driver_proto_cc",
- "//tensorflow/compiler/xla/python/tpu_driver/client:libtpu",
- ] + external_deps(),
- alwayslink = 1,
-)
-
-cc_library(
name = "direct_tpu_driver",
srcs = ["direct_tpu_driver.cc"],
deps = [
diff --git a/tensorflow/compiler/xla/python/tpu_driver/client/BUILD b/tensorflow/compiler/xla/python/tpu_driver/client/BUILD
index b5f1a83..cdfd1ca 100644
--- a/tensorflow/compiler/xla/python/tpu_driver/client/BUILD
+++ b/tensorflow/compiler/xla/python/tpu_driver/client/BUILD
@@ -22,7 +22,6 @@
"//tensorflow/compiler/xla/python:local_client",
"//tensorflow/compiler/xla/python:semaphore",
"//tensorflow/compiler/xla/python/tpu_driver",
- "//tensorflow/compiler/xla/python/tpu_driver:direct_tpu_driver",
"//tensorflow/compiler/xla/python/tpu_driver:grpc_tpu_driver",
"//tensorflow/compiler/xla/python/tpu_driver:recording_tpu_driver",
"//tensorflow/compiler/xla/python/tpu_driver:tpu_driver_proto_cc",
diff --git a/tensorflow/compiler/xla/python/tpu_driver/direct_tpu_driver.cc b/tensorflow/compiler/xla/python/tpu_driver/direct_tpu_driver.cc
index 54f2ddc..3e4626c 100644
--- a/tensorflow/compiler/xla/python/tpu_driver/direct_tpu_driver.cc
+++ b/tensorflow/compiler/xla/python/tpu_driver/direct_tpu_driver.cc
@@ -27,6 +27,11 @@
namespace tpu_driver {
namespace {
+// Enable the macro by default in the env where the libtpu.so is available.
+#ifdef PLATFORM_GOOGLE
+#define TPU_SHARED_LIBRARY_COMPILE_LINK 1
+#endif
+
xla::Status CreateXlaStatus(::TpuStatus* status) {
if (status->code == tensorflow::error::OK) {
return xla::Status::OK();
