From 9839be926d37456eb09b499cc7d2b64d1d81dc84 Mon Sep 17 00:00:00 2001
From: Jose Baiocchi <jbaiocchi@google.com>
Date: Fri, 9 Oct 2020 14:06:55 -0700
Subject: [PATCH] Add tf_profiler_pywrap_cc_library_wrapper macro

PiperOrigin-RevId: 336361876
Change-Id: I63941a552f9833d0004896be3b1eef9acef7ac45
---
 .../core/profiler/builds/build_config.bzl     |  2 ++
 .../core/profiler/builds/oss/build_config.bzl | 12 +++++++
 tensorflow/core/profiler/lib/BUILD            | 32 +++++++++----------
 tensorflow/core/profiler/rpc/BUILD            | 27 +++++++++-------
 tensorflow/core/profiler/rpc/client/BUILD     | 24 +++++++-------
 .../profiler/rpc/profiler_service_impl.cc     |  1 -
 tensorflow/python/profiler/internal/BUILD     | 19 ++++-------
 7 files changed, 63 insertions(+), 54 deletions(-)

diff --git a/tensorflow/core/profiler/builds/build_config.bzl b/tensorflow/core/profiler/builds/build_config.bzl
index 5560baa9b17..bd20787b398 100644
--- a/tensorflow/core/profiler/builds/build_config.bzl
+++ b/tensorflow/core/profiler/builds/build_config.bzl
@@ -3,9 +3,11 @@
 load(
     "//tensorflow/core/profiler/builds/oss:build_config.bzl",
     _tf_profiler_alias = "tf_profiler_alias",
+    _tf_profiler_pybind_cc_library_wrapper = "tf_profiler_pybind_cc_library_wrapper",
 )
 
 tf_profiler_alias = _tf_profiler_alias
+tf_profiler_pybind_cc_library_wrapper = _tf_profiler_pybind_cc_library_wrapper
 
 def if_profiler_oss(if_true, if_false = []):
     return select({
diff --git a/tensorflow/core/profiler/builds/oss/build_config.bzl b/tensorflow/core/profiler/builds/oss/build_config.bzl
index 1dcfd0e3291..8849aa4cff4 100644
--- a/tensorflow/core/profiler/builds/oss/build_config.bzl
+++ b/tensorflow/core/profiler/builds/oss/build_config.bzl
@@ -3,5 +3,17 @@
 TF profiler build macros for use in OSS.
 """
 
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library")
+
 def tf_profiler_alias(target_dir, name):
     return target_dir + "oss:" + name
+
+def tf_profiler_pybind_cc_library_wrapper(name, actual, **kwargs):
+    """Wrapper for cc_library used by tf_python_pybind_extension.
+
+    This wrapper ensures that cc libraries headers are made available to pybind
+    code, without creating ODR violations in the dynamically linked case.  The
+    symbols in these deps symbols should be linked to, and exported by, the core
+    pywrap_tensorflow_internal.so
+    """
+    cc_header_only_library(name = name, deps = [actual], **kwargs)
diff --git a/tensorflow/core/profiler/lib/BUILD b/tensorflow/core/profiler/lib/BUILD
index 8ba43a01504..cf4ba6f84a9 100644
--- a/tensorflow/core/profiler/lib/BUILD
+++ b/tensorflow/core/profiler/lib/BUILD
@@ -1,9 +1,12 @@
 load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow/core/platform:build_config_root.bzl", "if_static")
 load("//tensorflow:tensorflow.bzl", "if_not_android", "tf_cc_test", "tf_cuda_library")
-load("//tensorflow:tensorflow.bzl", "filegroup")  # buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "tf_pybind_cc_library_wrapper")  # buildifier: disable=same-origin-load
-load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_copts")
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load(
+    "//tensorflow/core/profiler/builds:build_config.bzl",
+    "tf_profiler_copts",
+    "tf_profiler_pybind_cc_library_wrapper",
+)
 
 package(
     default_visibility = [
@@ -13,13 +16,10 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
-tf_pybind_cc_library_wrapper(
-    name = "profiler_session_headers",
-    visibility = [
-        "//tensorflow/core/profiler/rpc:__pkg__",
-        "//tensorflow/python/profiler/internal:__pkg__",
-    ],
-    deps = [":profiler_session"],
+tf_profiler_pybind_cc_library_wrapper(
+    name = "profiler_session_for_pybind",
+    actual = ":profiler_session",
+    visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
 )
 
 cc_library(
@@ -68,13 +68,13 @@ cc_library(
     alwayslink = True,
 )
 
-tf_pybind_cc_library_wrapper(
-    name = "local_profiler_headers",
+tf_profiler_pybind_cc_library_wrapper(
+    name = "local_profiler_for_pybind",
+    actual = ":local_profiler",
     visibility = [
         "//tensorflow/core/profiler/rpc:__pkg__",
         "//tensorflow/python/profiler/internal:__pkg__",
     ],
-    deps = [":local_profiler"],
 )
 
 cc_library(
@@ -179,10 +179,10 @@ tf_cuda_library(
     alwayslink = True,
 )
 
-tf_pybind_cc_library_wrapper(
-    name = "traceme_headers",
+tf_profiler_pybind_cc_library_wrapper(
+    name = "traceme_for_pybind",
+    actual = ":traceme",
     visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
-    deps = [":traceme"],
 )
 
 cc_library(
diff --git a/tensorflow/core/profiler/rpc/BUILD b/tensorflow/core/profiler/rpc/BUILD
index f5cae8d4651..d532ad90bb9 100644
--- a/tensorflow/core/profiler/rpc/BUILD
+++ b/tensorflow/core/profiler/rpc/BUILD
@@ -1,8 +1,12 @@
 load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_external_workspace_visible")  # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")  # buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "tf_pybind_cc_library_wrapper")  # buildifier: disable=same-origin-load
-load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_alias", "tf_profiler_copts")
+load(
+    "//tensorflow/core/profiler/builds:build_config.bzl",
+    "tf_profiler_alias",
+    "tf_profiler_copts",
+    "tf_profiler_pybind_cc_library_wrapper",
+)
 
 package(
     default_visibility = ["//tensorflow/core/profiler:internal"],
@@ -25,12 +29,12 @@ exports_files(
     visibility = ["//tensorflow/core/profiler/rpc:__subpackages__"],
 )
 
+# Linked to pywrap_tensorflow.
 cc_library(
     name = "profiler_service_impl",
     srcs = ["profiler_service_impl.cc"],
     hdrs = ["profiler_service_impl.h"],
     copts = tf_profiler_copts(),
-    features = ["-layering_check"],
     visibility = tf_external_workspace_visible(
         [
             "//tensorflow/core/data/service:__pkg__",
@@ -41,13 +45,20 @@ cc_library(
     deps = [
         "//tensorflow/core:lib",
         "//tensorflow/core/profiler:profiler_service_proto_cc",
-        "//tensorflow/core/profiler/lib:profiler_session_headers",
+        "//tensorflow/core/profiler/lib:profiler_session",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/memory",
         tf_grpc_cc_dependency(),
     ],
 )
 
+tf_profiler_pybind_cc_library_wrapper(
+    name = "profiler_server_for_pybind",
+    actual = ":profiler_server_impl",
+    visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
+)
+
 cc_library(
     name = "profiler_server_impl",
     srcs = ["profiler_server.cc"],
@@ -68,11 +79,3 @@ cc_library(
     ],
     alwayslink = True,
 )
-
-tf_pybind_cc_library_wrapper(
-    name = "profiler_server_headers",
-    visibility = [
-        "//tensorflow/python/profiler/internal:__pkg__",
-    ],
-    deps = [":profiler_server_impl"],
-)
diff --git a/tensorflow/core/profiler/rpc/client/BUILD b/tensorflow/core/profiler/rpc/client/BUILD
index 381a4d59b5c..9111e16b494 100644
--- a/tensorflow/core/profiler/rpc/client/BUILD
+++ b/tensorflow/core/profiler/rpc/client/BUILD
@@ -1,8 +1,11 @@
 load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
-load("//tensorflow:tensorflow.bzl", "tf_pybind_cc_library_wrapper")  # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
-load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_copts")
+load(
+    "//tensorflow/core/profiler/builds:build_config.bzl",
+    "tf_profiler_copts",
+    "tf_profiler_pybind_cc_library_wrapper",
+)
 
 # For platform specific build config
 load(
@@ -19,14 +22,14 @@ cc_library(
     srcs = ["capture_profile.cc"],
     hdrs = ["capture_profile.h"],
     copts = tf_profiler_copts(),
-    features = ["-layering_check"],
     visibility = [
         "//tensorflow/python/profiler/internal:__pkg__",
     ],
     deps = [
-        ":profiler_client_headers",
+        ":profiler_client_for_pybind",
         ":save_profile",
         "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
         "//tensorflow/core/profiler:profiler_analysis_proto_cc",
         "//tensorflow/core/profiler:profiler_options_proto_cc",
         "//tensorflow/core/profiler:profiler_service_proto_cc",
@@ -53,10 +56,9 @@ cc_library(
     ],
 )
 
-tf_pybind_cc_library_wrapper(
-    name = "profiler_client_headers",
-    visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
-    deps = [":profiler_client"],
+tf_profiler_pybind_cc_library_wrapper(
+    name = "profiler_client_for_pybind",
+    actual = ":profiler_client",
 )
 
 cc_library(
@@ -113,9 +115,9 @@ cc_library(
 tf_cc_test(
     name = "profiler_client_test",
     srcs = ["profiler_client_test.cc"],
-    features = ["-layering_check"],
     tags = ["external"],  # So that test suite reruns unconditionally.
     deps = [
+        ":profiler_client",
         ":profiler_client_impl",  # for oss
         ":profiler_client_test_util",
         "@com_google_absl//absl/time",
@@ -131,10 +133,9 @@ cc_library(
     srcs = ["remote_profiler_session_manager.cc"],
     hdrs = ["remote_profiler_session_manager.h"],
     copts = tf_profiler_copts(),
-    features = ["-layering_check"],
     visibility = ["//tensorflow/core/profiler:internal"],
     deps = [
-        ":profiler_client_headers",
+        ":profiler_client_for_pybind",
         ":save_profile",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
@@ -150,7 +151,6 @@ cc_library(
 tf_cc_test(
     name = "remote_profiler_session_manager_test",
     srcs = ["remote_profiler_session_manager_test.cc"],
-    features = ["-layering_check"],
     tags = ["external"],  # So that test suite reruns unconditionally.
     deps = [
         ":profiler_client_impl",  # for oss
diff --git a/tensorflow/core/profiler/rpc/profiler_service_impl.cc b/tensorflow/core/profiler/rpc/profiler_service_impl.cc
index 41bd0c7584b..cd94b1b28cf 100644
--- a/tensorflow/core/profiler/rpc/profiler_service_impl.cc
+++ b/tensorflow/core/profiler/rpc/profiler_service_impl.cc
@@ -27,7 +27,6 @@ limitations under the License.
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/status.h"
-#include "tensorflow/core/profiler/lib/profiler_interface.h"
 #include "tensorflow/core/profiler/lib/profiler_session.h"
 #include "tensorflow/core/profiler/profiler_service.grpc.pb.h"
 #include "tensorflow/core/profiler/profiler_service.pb.h"
diff --git a/tensorflow/python/profiler/internal/BUILD b/tensorflow/python/profiler/internal/BUILD
index 6b7b9a6db49..47d34aa000b 100644
--- a/tensorflow/python/profiler/internal/BUILD
+++ b/tensorflow/python/profiler/internal/BUILD
@@ -1,11 +1,6 @@
-# buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "cuda_py_test")
-
-# buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
-
-# buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "tf_python_pybind_extension")
+load("//tensorflow:tensorflow.bzl", "cuda_py_test")  # buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")  # buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_python_pybind_extension")  # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "py_test")
 load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_copts")
 
@@ -102,13 +97,12 @@ cc_library(
     name = "traceme_wrapper",
     hdrs = ["traceme_wrapper.h"],
     copts = tf_profiler_copts(),
-    features = ["-layering_check"],
     visibility = [
         "//tensorflow/compiler/xla/python:__pkg__",
     ],
     deps = [
         "//tensorflow/core:lib",
-        "//tensorflow/core/profiler/lib:traceme_headers",
+        "//tensorflow/core/profiler/lib:traceme_for_pybind",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@pybind11",
@@ -119,7 +113,6 @@ tf_python_pybind_extension(
     name = "_pywrap_profiler",
     srcs = ["profiler_wrapper.cc"],
     copts = tf_profiler_copts(),
-    features = ["-layering_check"],
     module_name = "_pywrap_profiler",
     visibility = [
         "//tensorflow/python/eager:__pkg__",
@@ -129,9 +122,9 @@ tf_python_pybind_extension(
         "//tensorflow/core:lib",
         "//tensorflow/core/profiler/convert:xplane_to_tools_data",
         "//tensorflow/core/profiler/convert:xplane_to_trace_events",
-        "//tensorflow/core/profiler/lib:profiler_session_headers",
+        "//tensorflow/core/profiler/lib:profiler_session_for_pybind",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
-        "//tensorflow/core/profiler/rpc:profiler_server_headers",
+        "//tensorflow/core/profiler/rpc:profiler_server_for_pybind",
         "//tensorflow/core/profiler/rpc/client:capture_profile",
         "//tensorflow/core/profiler/rpc/client:save_profile",
         "//tensorflow/python:pybind11_status",