From 73e3999026c46ef61c77b9e9051c7a5eddec88c2 Mon Sep 17 00:00:00 2001 From: Haibo Huang Date: Mon, 11 Sep 2023 11:43:50 -0700 Subject: [PATCH] Load tensorflow_framework with RTLD_GLOBAL This way libtpu can use symbols from it. PiperOrigin-RevId: 564454296 --- .../core/tpu/tpu_api_dlsym_initializer.cc | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc b/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc index 6f8191364d7..b4a871336b9 100644 --- a/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc +++ b/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc @@ -73,7 +73,29 @@ absl::Status InitializeTpuLibrary(void* library_handle) { return s; } +// Gets the path of current module. It is usually tensorflow_framework.so. +const char* GetCurrentModulePath() { + Dl_info DlInfo; + if (!dladdr((void*)GetCurrentModulePath, &DlInfo)) { + return nullptr; + } + return DlInfo.dli_fname; +} + absl::Status FindAndLoadTpuLibrary() { + // Reopen tensorflow_framework.so as RTLD_GLOBAL. So that libtpu can link to + // Tensorflow C API. + // TODO(b/299313561): Remove this dependency. + const char* so_name = GetCurrentModulePath(); + if (so_name != nullptr) { + LOG(INFO) << "Opening library: " << so_name; + void* tf_lib = dlopen(so_name, RTLD_NOW | RTLD_GLOBAL); + if (tf_lib == nullptr) { + return absl::InternalError( + absl::StrCat("Failed to open libtensorflow ", dlerror())); + } + } + const char* env_value = getenv("TPU_LIBRARY_PATH"); const char* libtpu_path = env_value && strlen(env_value) > 0 ? env_value : "libtpu.so"; -- GitLab