Use TF env for loading libs dynamically

tensorflow · May 16, 2020 · 9d0a6c8 · 9d0a6c8
1 parent 7cb7a3f
commit 9d0a6c8
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 16 deletions.
diff --git a/tensorflow-core/tensorflow-core-api/.bazelrc b/tensorflow-core/tensorflow-core-api/.bazelrc
@@ -128,8 +128,8 @@ build --define open_source_build=true
 test --define open_source_build=true
 
 # For workaround https://github.com/bazelbuild/bazel/issues/8772 with Bazel >= 0.29.1
-#build --java_toolchain=//third_party/toolchains/java:tf_java_toolchain
-#build --host_java_toolchain=//third_party/toolchains/java:tf_java_toolchain
+build --java_toolchain=@org_tensorflow//third_party/toolchains/java:tf_java_toolchain
+build --host_java_toolchain=@org_tensorflow//third_party/toolchains/java:tf_java_toolchain
 
 # Please note that MKL on MacOS or windows is still not supported.
 # If you would like to use a local MKL instead of downloading, please set the

diff --git a/tensorflow-core/tensorflow-core-api/BUILD b/tensorflow-core/tensorflow-core-api/BUILD
@@ -61,13 +61,15 @@ java_proto_library(
     deps = ["@org_tensorflow//tensorflow/core:protos_all"]
 )
 
-cc_binary(
+tf_cc_binary(
     name = "libcustom_op_test.so",
     srcs = ["src/bazel/test/my_test_op.cc"],
-    copts = ["-fPIC"],
     linkshared = 1,
+    linkopts = select({
+        "@org_tensorflow//tensorflow:windows": [],
+        "//conditions:default": ["-lm"],
+    }),
     deps = [
         "@org_tensorflow//tensorflow/core:framework",
-        "@org_tensorflow//tensorflow:libtensorflow_framework_import_lib"
     ]
 )
diff --git a/tensorflow-core/tensorflow-core-api/src/bazel/api_def/import/api_import.cc b/tensorflow-core/tensorflow-core-api/src/bazel/api_def/import/api_import.cc
@@ -16,7 +16,6 @@
 #include <string>
 #include <vector>
 #include <iostream>
-#include <dlfcn.h>
 
 #include "tensorflow/core/framework/op_gen_lib.h"
 #include "tensorflow/core/framework/op.h"
@@ -153,11 +152,13 @@ int main(int argc, char* argv[]) {
   port::InitMain(usage.c_str(), &argc, &argv);
   QCHECK(parsed_flags_ok && !java_api_dir.empty()
 		  && !tf_src_dir.empty() && !tf_lib_path.empty()) << usage;
-  void* tf_lib_handle = dlopen(tf_lib_path.c_str(), RTLD_NOW);  // Register TF ops
+
+  Env* env = Env::Default();
+  void* tf_lib_handle;
+  TF_CHECK_OK(env->LoadLibrary(argv[1], &tf_lib_handle));  // This registers all TF ops
   OpList op_defs;
   OpRegistry::Global()->Export(false, &op_defs);
   ApiDefMap python_api_map(op_defs);
-  Env* env = Env::Default();
 
   // Load Python API defs
   string base_api_dir = tf_src_dir + "/tensorflow/core/api_def/base_api";
@@ -320,6 +321,5 @@ int main(int argc, char* argv[]) {
   } else {
     LOG(INFO) << "All resolved!";
   }
-  dlclose(tf_lib_handle);
   return 0;
 }
diff --git a/tensorflow-core/tensorflow-core-api/src/bazel/op_generator/op_gen_main.cc b/tensorflow-core/tensorflow-core-api/src/bazel/op_generator/op_gen_main.cc
@@ -15,7 +15,6 @@
 
 #include <string>
 #include <vector>
-#include <dlfcn.h>
 
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/lib/core/status.h"
@@ -73,17 +72,15 @@ int main(int argc, char* argv[]) {
   QCHECK(parsed_flags_ok && !output_dir.empty() && argc > 1) << usage;
   std::vector<tensorflow::string> api_dirs = tensorflow::str_util::Split(
       api_dirs_str, ",", tensorflow::str_util::SkipEmpty());
-  std::vector<void*> ops_libs_handles;
-  ops_libs_handles.reserve(argc - 1);
+
+  tensorflow::Env* env = tensorflow::Env::Default();
+  void* ops_libs_handles[50];
   for (int i = 1; i < argc; ++i) {
-	  ops_libs_handles.push_back(dlopen(argv[i], RTLD_NOW));
+	  TF_CHECK_OK(env->LoadLibrary(argv[1], &ops_libs_handles[i - 1]));
   }
   tensorflow::java::OpGenerator generator(api_dirs);
   tensorflow::OpList ops;
   tensorflow::OpRegistry::Global()->Export(false, &ops);
   TF_CHECK_OK(generator.Run(ops, base_package, output_dir));
-  for (void* ops_lib_handle : ops_libs_handles) {
-	  dlclose(ops_lib_handle);
-  }
   return 0;
 }