Adds necessary hooks to load a TPU-specific shared library.

PiperOrigin-RevId: 312601701
Change-Id: I1ae43d253d1734c30ffefe4d4062c82639d7a4d1

tensorflow/core/BUILD

@@ -2254,6 +2254,7 @@ tf_cuda_library(
         "//tensorflow/core/platform/default/build_config:platformlib",
         "//tensorflow/core/profiler/lib:annotated_traceme",
         "//tensorflow/core/profiler/lib:traceme",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/core/util:einsum_op_util",
         "//tensorflow/core/util:padding",
         "//tensorflow/core/util:port",

tensorflow/core/framework/load_library.cc

@@ -21,6 +21,9 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/mem.h"
+#if !defined(IS_MOBILE_PLATFORM)
+#include "tensorflow/core/tpu/tpu_library_loader.h"
+#endif  // IS_MOBILE_PLATFORM
 
 namespace tensorflow {
 
@@ -97,6 +100,17 @@ Status LoadLibrary(const char* library_filename, void** result,
   *buf = str_buf;
   *len = str.length();
 
+#if !defined(IS_MOBILE_PLATFORM)
+  // Determine if this library is a TPU library, and if so, calls the TPU
+  // initialization functions to populate function tables, etc...
+  void* unused_symbol;
+  if (env->GetSymbolFromLibrary(library.handle, "TfTpu_Initialize",
+                                &unused_symbol)
+          .ok()) {
+    TF_RETURN_IF_ERROR(tensorflow::tpu::InitializeTPULibrary(library.handle));
+  }
+#endif  // IS_MOBILE_PLATFORM
+
   *result = library.handle;
   return Status::OK();
 }

tensorflow/core/tpu/BUILD

@@ -91,3 +91,11 @@ cc_library(
         "//tensorflow/c:tf_status",
     ],
 )
+
+cc_library(
+    name = "tpu_library_loader",
+    srcs = ["tpu_library_loader.cc"],
+    hdrs = ["tpu_library_loader.h"],
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = ["//tensorflow/core/platform:status"],
+)

tensorflow/core/tpu/tpu_library_loader.cc

@@ -0,0 +1,30 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/tpu/tpu_library_loader.h"
+
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace tpu {
+
+Status InitializeTPULibrary(void* library) {
+  // TODO(frankchn): dlsym the loaded library and populate a struct with the
+  // relevant C APIs necessary for TPUs.
+  return Status::OK();
+}
+
+}  // namespace tpu
+}  // namespace tensorflow

tensorflow/core/tpu/tpu_library_loader.h

@@ -0,0 +1,29 @@
+/* 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_TPU_LIBRARY_LOADER_H_
+#define TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
+
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace tpu {
+
+Status InitializeTPULibrary(void* library);
+
+}  // namespace tpu
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_