[TF1] Port of dynamic pywrap - pywrap_tf_internal loads TF symbols from libtf_cc.so instead of duplicating them

Author: trevor-m

.bazelrc

@@ -195,6 +195,7 @@ build --define=grpc_no_ares=true
 # Modular TF build options
 build:dynamic_kernels --define=dynamic_loaded_kernels=true
 build:dynamic_kernels --copt=-DAUTOLOAD_DYNAMIC_KERNELS
+build:dynamic_pywrap --define=dynamic_pywrap=true
 
 # Build TF with C++ 17 features.
 build:c++17 --cxxopt=-std=c++1z
@@ -216,6 +217,9 @@ build:short_logs --output_filter=DONT_MATCH_ANYTHING
 build --cxxopt=-std=c++14
 build --host_cxxopt=-std=c++14
 
+# Use dynamic pywrap by default
+build --config=dynamic_pywrap
+
 # Instruction set optimizations
 # TODO(gunan): Create a feature in toolchains for avx/avx2 to
 #   avoid having to define linux/win separately.

tensorflow/BUILD

@@ -31,6 +31,7 @@ load(
     "//third_party/mkl:build_defs.bzl",
     "if_mkl_ml",
 )
+load("//tensorflow/core/platform:default/build_config_root.bzl", "if_dynamic_pywrap")
 
 package(
     default_visibility = [":internal"],
@@ -464,6 +465,16 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+# Enables pywrap_tensorflow_internal to load TensorFlow symbols from libtensorflow_cc instead of duplicating them.
+config_setting(
+    name = "dynamic_pywrap",
+    define_values = {
+        "dynamic_pywrap": "true",
+        "framework_shared_object": "true",
+    },
+    visibility = ["//visibility:public"],
+)
+
 # DO NOT ADD ANY NEW EXCEPTIONS TO THIS LIST!
 # Instead, please use public APIs or public build rules TF provides.
 # If you need functionality that is not exposed, we will work with you to expand our public APIs.
@@ -676,7 +687,15 @@ tf_cc_shared_object(
         "//tensorflow/contrib/tensor_forest:model_ops_lib",
         "//tensorflow/contrib/tensor_forest:forest_proto_impl",
         "//tensorflow/contrib/tensor_forest/kernels/v4:decision-tree-resource_impl",
-    ] + if_ngraph(["@ngraph_tf//:ngraph_tf"]),
+    ] + if_ngraph(["@ngraph_tf//:ngraph_tf"]) + if_dynamic_pywrap(["//tensorflow/python:tensorflow_impl"]),
+)
+
+cc_import(
+    name = "libtensorflow_cc_import_lib",
+    shared_library = select({
+        "//tensorflow:macos": ":libtensorflow_cc.dylib",
+        "//conditions:default": ":libtensorflow_cc.so",
+    }),
 )
 
 # ** Targets for Windows build (start) **

tensorflow/c/BUILD

@@ -31,6 +31,18 @@ filegroup(
     visibility = ["//tensorflow:__subpackages__"],
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "c_api_internal.h",
+        "checkpoint_reader.h",
+        "tf_status_helper.h",
+        "tf_status_internal.h",
+        "tf_tensor_internal.h",
+    ],
+    visibility = ["//tensorflow/python:__subpackages__"],
+)
+
 filegroup(
     name = "srcs",
     srcs = glob(
@@ -629,11 +641,17 @@ tf_cc_test(
 tf_cuda_library(
     name = "python_api",
     srcs = ["python_api.cc"],
-    hdrs = ["python_api.h"],
+    hdrs = [
+        "python_api.h",
+        ":headers",
+        ":pywrap_required_hdrs",
+        "//tensorflow/core:protos_all_proto_text_hdrs",
+        "//tensorflow/core:pywrap_required_hdrs",
+    ],
     visibility = ["//tensorflow/python:__pkg__"],
     deps = [
-        ":c_api",
-        ":c_api_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
         # TODO(b/74620627): remove when _USE_C_SHAPES is removed
         "//tensorflow/python:cpp_shape_inference_proto_cc",
     ],

tensorflow/c/eager/BUILD

@@ -21,6 +21,18 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "c_api_internal.h",
+        "c_api_experimental.h",
+        "tape.h",
+    ],
+    visibility = [
+        "//tensorflow/python:__subpackages__",
+    ],
+)
+
 tf_cuda_library(
     name = "c_api",
     srcs = [

tensorflow/contrib/tensor_forest/BUILD

@@ -81,7 +81,9 @@ tf_custom_op_library(
     ] + if_static(
         extra_deps = [],
         otherwise = [
-            ":libforestprotos.so",
+            # Proto implementations are already in libtensorflow_cc.so due to Triton
+            # ":libforestprotos.so",
+            "//tensorflow:libtensorflow_cc.so.1",
         ],
     ),
     deps = [
@@ -235,7 +237,9 @@ tf_custom_op_library(
     ] + if_static(
         extra_deps = [],
         otherwise = [
-            ":libforestprotos.so",
+            # Proto implementations are already in libtensorflow_cc.so due to Triton
+            # ":libforestprotos.so",
+            "//tensorflow:libtensorflow_cc.so.1",
         ],
     ),
     deps = [":model_ops_lib"],
@@ -330,7 +334,9 @@ tf_custom_op_library(
     ] + if_static(
         extra_deps = [],
         otherwise = [
-            ":libforestprotos.so",
+            # Proto implementations are already in libtensorflow_cc.so due to Triton
+            # ":libforestprotos.so",
+            "//tensorflow:libtensorflow_cc.so.1",
         ],
     ),
     deps = [":stats_ops_lib"],

tensorflow/core/BUILD

@@ -3199,6 +3199,20 @@ CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [
     "graph/quantize_training.h",
 ] + if_mkl(["graph/mkl_graph_util.h"])
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = CORE_CPU_LIB_HEADERS + [
+        "common_runtime/gpu/gpu_id.h",
+        "framework/op_gen_lib.h",
+        "public/session.h",
+        "public/session_options.h",
+    ],
+    visibility = [
+        "//tensorflow/c:__pkg__",
+        "//tensorflow/python:__subpackages__",
+    ],
+)
+
 tf_cuda_library(
     name = "core_cpu_impl",
     srcs = [

tensorflow/core/common_runtime/eager/BUILD

@@ -17,6 +17,22 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "attr_builder.h",
+        "context.h",
+        "eager_executor.h",
+        "eager_operation.h",
+        "kernel_and_device.h",
+        "tensor_handle.h",
+        "tensor_handle_data.h",
+    ],
+    visibility = [
+        "//tensorflow/python:__subpackages__",
+    ],
+)
+
 tf_cuda_library(
     name = "eager_executor",
     srcs = [

tensorflow/core/distributed_runtime/BUILD

@@ -31,6 +31,22 @@ filegroup(
     ]),
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "call_options.h",
+        "message_wrappers.h",
+        "rendezvous_mgr_interface.h",
+        "server_lib.h",
+        "worker_cache.h",
+        "worker_env.h",
+        "worker_interface.h",
+    ],
+    visibility = [
+        "//tensorflow/python:__subpackages__",
+    ],
+)
+
 cc_library(
     name = "partial_run_mgr",
     srcs = ["partial_run_mgr.cc"],

tensorflow/core/distributed_runtime/eager/BUILD

@@ -12,6 +12,18 @@ package(
 
 exports_files(["LICENSE"])
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "eager_client.h",
+        "remote_tensor_handle.h",
+        "remote_tensor_handle_data.h",
+    ],
+    visibility = [
+        "//tensorflow/python:__subpackages__",
+    ],
+)
+
 cc_library(
     name = "remote_tensor_handle",
     hdrs = ["remote_tensor_handle.h"],

tensorflow/core/grappler/BUILD

@@ -4,6 +4,19 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+exports_files(["grappler_item_builder.h"])
+
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "devices.h",
+        "grappler_item.h",
+        "grappler_item_builder.h",
+        "utils.h",
+    ],
+    visibility = ["//tensorflow/python:__subpackages__"],
+)
+
 cc_library(
     name = "op_types",
     srcs = ["op_types.cc"],

tensorflow/core/grappler/clusters/BUILD

@@ -10,6 +10,17 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "cluster.h",
+        "single_machine.h",
+        "virtual_cluster.h",
+        "utils.h",
+    ],
+    visibility = ["//tensorflow/python:__subpackages__"],
+)
+
 config_setting(
     name = "xsmm",
     licenses = ["notice"],

tensorflow/core/grappler/costs/BUILD

@@ -11,6 +11,23 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "analytical_cost_estimator.h",
+        "cost_estimator.h",
+        "graph_properties.h",
+        "measuring_cost_estimator.h",
+        "op_level_cost_estimator.h",
+        "op_context.h",
+        "virtual_scheduler.h",
+        "virtual_placer.h",
+        "graph_memory.h",
+        "utils.h",
+    ],
+    visibility = ["//tensorflow/python:__subpackages__"],
+)
+
 filegroup(
     name = "graph_properties_testdata",
     srcs = glob([

tensorflow/core/grappler/graph_analyzer/BUILD

@@ -4,6 +4,11 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+exports_files(
+    ["graph_analyzer_tool.h"],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "graph_analyzer_lib",
     srcs = [

tensorflow/core/grappler/optimizers/BUILD

@@ -11,6 +11,15 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "meta_optimizer.h",
+        "graph_optimizer.h",
+    ],
+    visibility = ["//tensorflow/python:__subpackages__"],
+)
+
 cc_library(
     name = "static_schedule",
     srcs = ["static_schedule.cc"],

tensorflow/core/grappler/utils/BUILD

@@ -8,6 +8,8 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+exports_files(["topological_sort.h"])
+
 cc_library(
     name = "scc",
     srcs = ["scc.cc"],

tensorflow/core/grappler/verifiers/BUILD

@@ -4,6 +4,8 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+exports_files(["graph_verifier.h"])
+
 cc_library(
     name = "graph_verifier",
     hdrs = [

tensorflow/core/platform/default/build_config_root.bzl

@@ -95,3 +95,9 @@ def if_dynamic_kernels(extra_deps, otherwise = []):
         str(Label("//tensorflow:dynamic_loaded_kernels")): extra_deps,
         "//conditions:default": otherwise,
     })
+
+def if_dynamic_pywrap(extra_deps, otherwise = []):
+    return select({
+        str(Label("//tensorflow:dynamic_pywrap")): extra_deps,
+        "//conditions:default": otherwise,
+    })