diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 6b4874a8393..2b16801f6ed 100644 --- a/tensorflow/core/BUILD +++ b/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", diff --git a/tensorflow/core/framework/load_library.cc b/tensorflow/core/framework/load_library.cc index b9e33b148f7..c223eac4722 100644 --- a/tensorflow/core/framework/load_library.cc +++ b/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(); } diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD index 48a9a229d2a..5d1b7e1101f 100644 --- a/tensorflow/core/tpu/BUILD +++ b/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"], +) diff --git a/tensorflow/core/tpu/tpu_library_loader.cc b/tensorflow/core/tpu/tpu_library_loader.cc new file mode 100644 index 00000000000..bfd9fe29efe --- /dev/null +++ b/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 diff --git a/tensorflow/core/tpu/tpu_library_loader.h b/tensorflow/core/tpu/tpu_library_loader.h new file mode 100644 index 00000000000..35a7dd7c9be --- /dev/null +++ b/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_