From 9d0a6c87f7b0454f174511ea55ca00a451f82e4f Mon Sep 17 00:00:00 2001 From: Karl Lessard Date: Fri, 15 May 2020 15:12:54 -0400 Subject: [PATCH] Use TF env for loading libs dynamically --- tensorflow-core/tensorflow-core-api/.bazelrc | 4 ++-- tensorflow-core/tensorflow-core-api/BUILD | 8 +++++--- .../src/bazel/api_def/import/api_import.cc | 8 ++++---- .../src/bazel/op_generator/op_gen_main.cc | 11 ++++------- 4 files changed, 15 insertions(+), 16 deletions(-) diff --git a/tensorflow-core/tensorflow-core-api/.bazelrc b/tensorflow-core/tensorflow-core-api/.bazelrc index 6f672e7d087..64fe001564f 100644 --- 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 index 31849ee31be..f9af9f62211 100644 --- 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 index f4fba282cf0..fb55d23258e 100644 --- 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 #include #include -#include #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 index 42d0b6d1926..f35a84b8178 100644 --- 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 #include -#include #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 api_dirs = tensorflow::str_util::Split( api_dirs_str, ",", tensorflow::str_util::SkipEmpty()); - std::vector 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; }