chore: import upstream snapshot with attribution

tensorflow/python/profiler/internal/BUILD

@@ -0,0 +1,263 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("@xla//xla/tsl:tsl.bzl", "if_macos", "internal_visibility")
+load("@xla//xla/tsl:tsl.default.bzl", "tsl_pybind_extension")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "get_compatible_with_portable", "tf_python_pybind_extension")
+load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/python/profiler:__subpackages__"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "flops_registry",
+    srcs = ["flops_registry.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:graph_util",
+        "//tensorflow/python/framework:ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_test(
+    name = "flops_registry_test",
+    srcs = ["flops_registry_test.py"],
+    strict_deps = True,
+    deps = [
+        ":flops_registry",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "model_analyzer_testlib",
+    srcs = ["model_analyzer_testlib.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_grad",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:rnn",
+        "//tensorflow/python/ops:rnn_cell",
+        "//tensorflow/python/ops:tensor_array_grad",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/profiler:model_analyzer",
+        "//tensorflow/python/training:gradient_descent",
+        "//tensorflow/python/util:_pywrap_tfprof",
+        "//tensorflow/python/util:compat",
+    ],
+)
+
+py_test(
+    name = "print_model_analysis_test",
+    srcs = ["print_model_analysis_test.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "run_metadata_test",
+    srcs = ["run_metadata_test.py"],
+    tags = [
+        "no_gpu",  # b/138442728
+        "no_pip",
+    ],
+    xla_enable_strict_auto_jit = False,  # Node names are different with autojit
+    deps = [
+        ":model_analyzer_testlib",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/profiler:model_analyzer",
+        "//tensorflow/python/profiler:option_builder",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_traceme",
+    srcs = ["traceme_wrapper.cc"],
+    copts = tf_profiler_copts(),
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_traceme.pyi",
+    ],
+    visibility = [
+        "//perftools/accelerators/xprof/xprofilez/integration_tests:__pkg__",
+        "//tensorflow/python:__pkg__",
+        "//tensorflow/python/profiler:__subpackages__",
+        "//tensorflow/tools/pip_package:__subpackages__",
+    ],
+    deps = [
+        "@pybind11",
+        "@xla//xla/python/profiler/internal:traceme_state",
+        "@xla//xla/python/profiler/internal:traceme_wrapper",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_profiler",
+    srcs = ["profiler_wrapper.cc"],
+    copts = tf_profiler_copts(),
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_profiler.pyi",
+    ],
+    visibility = [
+        "//tensorflow/core/profiler:internal",
+        "//tensorflow/python:__pkg__",
+        "//tensorflow/python/eager:__pkg__",
+        "//tensorflow/python/profiler:__pkg__",
+        "//tensorflow/tools/pip_package:__subpackages__",
+    ],
+    deps = [
+        ":profiler_pywrap_impl",
+        "//tensorflow/core/profiler/rpc:profiler_server_for_pybind",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@com_google_absl//absl/status",
+        "@org_xprof//xprof/convert:repository",
+        "@org_xprof//xprof/convert:tool_options",
+        "@org_xprof//xprof/convert:xplane_to_tools_data",
+        "@pybind11",
+    ],
+)
+
+cc_library(
+    name = "python_hooks",
+    hdrs = ["python_hooks.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tf_profiler_copts() + ["-fexceptions"],
+    features = ["-use_header_modules"],  # Incompatible with -fexceptions.
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/memory",
+        "@pybind11",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+        "@xla//xla/python/profiler/internal:python_hooks",
+    ],
+    alwayslink = True,
+)
+
+cc_library(
+    name = "profiler_pywrap_impl",
+    srcs = ["profiler_pywrap_impl.cc"],
+    hdrs = ["profiler_pywrap_impl.h"],
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core/profiler/lib:profiler_session_for_pybind",
+        "//tensorflow/core/profiler/rpc:profiler_server_for_pybind",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:variant",
+        "@tsl//tsl/profiler/lib:profiler_session",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+        "@xla//xla/tsl/profiler/convert:xplane_to_trace_events",
+        "@xla//xla/tsl/profiler/rpc/client:capture_profile",
+        "@xla//xla/tsl/profiler/utils:session_manager",
+    ],
+)
+
+tsl_pybind_extension(
+    name = "_pywrap_profiler_plugin",
+    srcs = ["pywrap_profiler_plugin.cc"],
+    pytype_srcs = [
+        "_pywrap_profiler_plugin.pyi",
+    ],
+    visibility = internal_visibility([
+        "//tensorflow/core/profiler:internal",
+        "//tensorflow/python/eager:__pkg__",
+        "//tensorflow/python/profiler:__pkg__",
+    ]),
+    deps = [
+        "//tensorflow/core:protos_all_cc_impl",
+        "//tensorflow/core/framework:attr_value_proto_cc_impl",
+        "//tensorflow/core/framework:op",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/python/lib/core:py_exception_registry",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_protobuf//:protobuf",
+        "@org_xprof//xprof/convert:tool_options",
+        "@org_xprof//xprof/pywrap:profiler_plugin_impl",
+        "@pybind11",
+        "@tsl//tsl/profiler/protobuf:profiler_analysis_proto_cc_impl",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc_impl",
+        "@tsl//tsl/profiler/protobuf:profiler_service_monitor_result_proto_cc_impl",
+        "@tsl//tsl/profiler/protobuf:profiler_service_proto_cc_impl",
+        "@xla//xla:autotune_results_proto_cc_impl",
+        "@xla//xla:autotuning_proto_cc_impl",
+        "@xla//xla:xla_data_proto_cc_impl",
+        "@xla//xla:xla_proto_cc_impl",
+        "@xla//xla/pjrt:status_casters",
+        "@xla//xla/service:hlo_proto_cc_impl",
+        "@xla//xla/service:metrics_proto_cc_impl",
+        "@xla//xla/service/gpu:backend_configs_cc_impl",
+        "@xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl",
+        "@xla//xla/stream_executor:device_description_proto_cc_impl",
+        "@xla//xla/stream_executor/cuda:cuda_compute_capability_proto_cc_impl",
+        "@xla//xla/tsl/framework:allocator_registry_impl",
+        "@xla//xla/tsl/lib/io:table",
+        "@xla//xla/tsl/platform:env_impl",
+        "@xla//xla/tsl/platform:types",
+        "@xla//xla/tsl/platform/cloud:gcs_file_system",
+        "@xla//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
+        "@xla//xla/tsl/profiler/rpc:profiler_server_impl",
+        "@xla//xla/tsl/profiler/rpc:profiler_service_impl",
+        "@xla//xla/tsl/profiler/rpc/client:capture_profile",
+        "@xla//xla/tsl/profiler/rpc/client:profiler_client_impl",
+        "@xla//xla/tsl/profiler/utils:session_manager",
+        "@xla//xla/tsl/profiler/utils:time_utils_impl",
+        "@xla//xla/tsl/protobuf:bfc_memory_map_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc_impl",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:rpc_options_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:test_log_proto_cc_impl",
+    ] + if_macos([
+        "@xla//xla/tsl/lib/histogram",
+        "@xla//xla/tsl/lib/io:record_writer",
+        "@xla//xla/tsl/lib/monitoring:collection_registry",
+        "@xla//xla/tsl/lib/monitoring:sampler",
+        "//tensorflow/core:framework_internal_impl",
+        "//tensorflow/core/common_runtime/gpu:gpu_id_impl",
+        "//tensorflow/core/lib/core:arena",
+        "//tensorflow/core/lib/strings:ordered_code",
+        "//tensorflow/core/platform:platform_strings",
+        "@tsl//tsl/platform:random",
+        "@xla//xla/tsl/platform:resource",
+        "@tsl//tsl/profiler/lib:profiler_factory_impl",
+        "@tsl//tsl/profiler/lib:profiler_session_impl",
+    ]),
+)

tensorflow/python/profiler/internal/_pywrap_profiler.pyi

@@ -0,0 +1,22 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class ProfilerSession:
+    def __init__(self) -> None: ...
+    def export_to_tb(self) -> None: ...
+    def start(self, arg0: str, arg1: dict) -> None: ...
+    def stop(self) -> bytes: ...
+
+def start_server(arg0: int) -> None: ...

tensorflow/python/profiler/internal/_pywrap_profiler_plugin.pyi

@@ -0,0 +1,19 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def monitor(arg0: str, arg1: int, arg2: int, arg3: bool) -> str: ...
+def trace(arg0: str, arg1: str, arg2: str, arg3: bool, arg4: int, arg5: int, arg6: dict) -> None: ...
+def xspace_to_tools_data(arg0: list, arg1: str, arg2: dict = ...) -> tuple: ...
+def xspace_to_tools_data_from_byte_string(arg0: list, arg1: list, arg2: str, arg3: dict) -> tuple: ...

tensorflow/python/profiler/internal/_pywrap_traceme.pyi

@@ -0,0 +1,21 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class TraceMe:
+    def __init__(self, arg0: str, **kwargs) -> None: ...
+    def SetMetadata(self, **kwargs) -> None: ...
+    def Stop(self) -> None: ...
+
+def traceme_enabled(*args, **kwargs): ...

tensorflow/python/profiler/internal/flops_registry.py

@@ -0,0 +1,480 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Register flops statistics for various TensorFlow operations.
+"""
+import numpy as np
+
+from tensorflow.python.framework import graph_util
+from tensorflow.python.framework import ops
+
+
+# List of all ops which have implemented flops statistics.
+IMPLEMENTED_OPS = set([
+    # Unary ops
+    "Reciprocal", "Square", "Rsqrt", "Log", "Neg", "AssignSub", "AssignAdd",
+    "L2Loss", "Softmax",
+    # Binary ops
+    "Add", "Sub", "Mul", "RealDiv", "Maximum", "Minimum", "Pow", "RsqrtGrad",
+    "GreaterEqual", "Greater", "LessEqual", "Less", "Equal", "NotEqual",
+    "SquaredDifference", "AddV2",
+    # Reduction ops
+    "Mean", "Sum", "ArgMax", "ArgMin", "BiasAddGrad",
+    # Convolution and pooling
+    "AvgPool", "MaxPool", "AvgPoolGrad", "MaxPoolGrad", "Conv2DBackpropInput",
+    "Conv2DBackpropFilter",
+    # Other ops
+    "AddN", "MatMul",
+    # Ops implemented in core tensorflow:
+    "Conv2D", "DepthwiseConv2dNative", "BiasAdd", "Dilation2D",
+])
+
+
+def _zero_flops(graph, node):
+  """Returns zero flops."""
+  del graph, node  # graph and node are unused
+  return ops.OpStats("flops", 0)
+
+
+def _list_product(lst):
+  """Computes product of element of the list."""
+  result = 1
+  for item in lst:
+    result *= item
+  return result
+
+################################################################################
+# Unary operations
+################################################################################
+
+
+def _unary_op_flops(graph, node, ops_per_element=1):
+  """Common code which compute flops for unary operations."""
+  in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  in_shape.assert_is_fully_defined()
+  return ops.OpStats("flops", in_shape.num_elements() * ops_per_element)
+
+
+@ops.RegisterStatistics("Reciprocal", "flops")
+def _reciprocal_flops(graph, node):
+  """Compute flops for Reciprocal operation."""
+  return _unary_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Square", "flops")
+def _square_flops(graph, node):
+  """Compute flops for Square operation."""
+  return _unary_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Rsqrt", "flops")
+def _rsqrt_flops(graph, node):
+  """Compute flops for Rsqrt operation."""
+  # Rsqrt(x) = 1 / sqrt(x)
+  return _unary_op_flops(graph, node, ops_per_element=2)
+
+
+@ops.RegisterStatistics("Log", "flops")
+def _log_flops(graph, node):
+  """Compute flops for Log operation."""
+  return _unary_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Neg", "flops")
+def _neg_flops(graph, node):
+  """Compute flops for Neg operation."""
+  return _unary_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("AssignSub", "flops")
+def _assign_sub_flops(graph, node):
+  """Compute flops for AssignSub operation."""
+  return _unary_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("AssignAdd", "flops")
+def _assign_add_flops(graph, node):
+  """Compute flops for AssignAdd operation."""
+  return _unary_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("L2Loss", "flops")
+def _l2_loss_flops(graph, node):
+  """Compute flops for L2Loss operation."""
+  in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  in_shape.assert_is_fully_defined()
+  # Tensorflow uses inefficient implementation, with (3*N-1) flops:
+  # Optimal implementation is 2*N flops
+  return ops.OpStats("flops", in_shape.num_elements() * 3 - 1)
+
+
+@ops.RegisterStatistics("Softmax", "flops")
+def _softmax_flops(graph, node):
+  """Compute flops for Softmax operation."""
+  # Softmax implementation:
+  #
+  # Approximate flops breakdown:
+  #   2*n          -- compute shifted logits
+  #   n            -- exp of shifted logits
+  #   2*n          -- compute softmax from exp of shifted logits
+  return _unary_op_flops(graph, node, ops_per_element=5)
+
+################################################################################
+# Binary operations
+################################################################################
+
+
+def _binary_per_element_op_flops(graph, node, ops_per_element=1):
+  """Common code which compute flops for binary operations."""
+  out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  out_shape.assert_is_fully_defined()
+  return ops.OpStats("flops", out_shape.num_elements() * ops_per_element)
+
+
+@ops.RegisterStatistics("Add", "flops")
+@ops.RegisterStatistics("AddV2", "flops")
+def _add_flops(graph, node):
+  """Compute flops for Add operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Sub", "flops")
+def _sub_flops(graph, node):
+  """Compute flops for Sub operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Mul", "flops")
+def _mul_flops(graph, node):
+  """Compute flops for Mul operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("RealDiv", "flops")
+def _real_div_flops(graph, node):
+  """Compute flops for RealDiv operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Maximum", "flops")
+def _maximum_flops(graph, node):
+  """Compute flops for Maximum operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Minimum", "flops")
+def _minimum_flops(graph, node):
+  """Compute flops for Minimum operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Pow", "flops")
+def _pow_flops(graph, node):
+  """Compute flops for Pow operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("RsqrtGrad", "flops")
+def _rsqrt_grad_flops(graph, node):
+  """Compute flops for RsqrtGrad operation."""
+  return _binary_per_element_op_flops(graph, node, ops_per_element=4)
+
+
+@ops.RegisterStatistics("GreaterEqual", "flops")
+def _greater_equal_flops(graph, node):
+  """Compute flops for GreaterEqual operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Greater", "flops")
+def _greater_flops(graph, node):
+  """Compute flops for Greater operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("LessEqual", "flops")
+def _less_equal_flops(graph, node):
+  """Compute flops for LessEqual operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Less", "flops")
+def _less_flops(graph, node):
+  """Compute flops for Less operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("Equal", "flops")
+def _equal_flops(graph, node):
+  """Compute flops for Equal operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("NotEqual", "flops")
+def _not_equal_flops(graph, node):
+  """Compute flops for NotEqual operation."""
+  return _binary_per_element_op_flops(graph, node)
+
+
+@ops.RegisterStatistics("SquaredDifference", "flops")
+def _squared_difference_flops(graph, node):
+  """Compute flops for SquaredDifference operation."""
+  return _binary_per_element_op_flops(graph, node, ops_per_element=2)
+
+################################################################################
+# Reduction ops
+################################################################################
+
+
+def _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0):
+  """Common code which compute flops for reduction operations."""
+  in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  in_shape.assert_is_fully_defined()
+  out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  out_shape.assert_is_fully_defined()
+  num_flops = (in_shape.num_elements() * reduce_flops
+               + out_shape.num_elements() * (finalize_flops - reduce_flops))
+  return ops.OpStats("flops", num_flops)
+
+
+@ops.RegisterStatistics("Mean", "flops")
+def _mean_flops(graph, node):
+  """Compute flops for Mean operation."""
+  # reduction - sum, finalization - divide
+  return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=1)
+
+
+@ops.RegisterStatistics("Sum", "flops")
+def _sum_flops(graph, node):
+  """Compute flops for Sum operation."""
+  # reduction - sum, no finalization
+  return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
+
+
+@ops.RegisterStatistics("ArgMax", "flops")
+def _arg_max_flops(graph, node):
+  """Compute flops for ArgMax operation."""
+  # reduction - comparison, no finalization
+  return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
+
+
+@ops.RegisterStatistics("ArgMin", "flops")
+def _arg_min_flops(graph, node):
+  """Compute flops for ArgMin operation."""
+  # reduction - comparison, no finalization
+  return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
+
+
+@ops.RegisterStatistics("BiasAddGrad", "flops")
+def _bias_add_grad_flops(graph, node):
+  """Compute flops for BiasAddGrad operation."""
+  # Implementation of BiasAddGrad, essentially it's a reduce sum and reshaping:
+  # So computing flops same way as for "Sum"
+  return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
+
+################################################################################
+# Convolution and pooling
+# Note: all flops statistics are implemented only for NHWC data format
+################################################################################
+
+
+def _verify_conv_data_format(node):
+  """Verifies data format for pooling and convolutional operations."""
+  # TODO(xpan): P1: Support NCHW
+  if node.attr["data_format"].s != b"NHWC":
+    raise ValueError("Only NHWC format is supported in flops computations")
+
+
+def _pool_flops(graph, node):
+  """Common code which compute flops for pooling operations."""
+  # compute flops for average and max pooling
+  _verify_conv_data_format(node)
+  #
+  # Pooling declaration:
+  #   Inputs:
+  #     - value
+  #   Outputs:
+  #     - output
+  #   Attributes:
+  #     - ksize
+  #     - strides
+  #     - padding
+  #     - data_format
+  #
+  # Pooling implementation:
+  out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  out_shape.assert_is_fully_defined()
+  kernel_shape = list(node.attr["ksize"].list.i)
+  kernel_area = _list_product(kernel_shape)
+  return ops.OpStats("flops", kernel_area * out_shape.num_elements())
+
+
+@ops.RegisterStatistics("AvgPool", "flops")
+def _avg_pool_flops(graph, node):
+  """Compute flops for AvgPool operation."""
+  return _pool_flops(graph, node)
+
+
+@ops.RegisterStatistics("MaxPool", "flops")
+def _max_pool_flops(graph, node):
+  """Compute flops for MaxPool operation."""
+  return _pool_flops(graph, node)
+
+
+@ops.RegisterStatistics("AvgPoolGrad", "flops")
+def _avg_pool_grad_flops(graph, node):
+  """Compute flops for AvgPoolGrad operation."""
+  _verify_conv_data_format(node)
+  # Pooling gradient implementation:
+  out_backprop_shape = graph_util.tensor_shape_from_node_def_name(graph,
+                                                                  node.input[1])
+  out_backprop_shape.assert_is_fully_defined()
+  kernel_shape = list(node.attr["ksize"].list.i)
+  kernel_area = _list_product(kernel_shape)
+  # TensorFlow multiply each element of pooling window by coefficient,
+  # then sum up all of them, thus we have 2 flops per element:
+  # More optimal implementation - if division is done after.
+  return ops.OpStats("flops",
+                     kernel_area * out_backprop_shape.num_elements() * 2)
+
+
+@ops.RegisterStatistics("MaxPoolGrad", "flops")
+def _max_pool_grad_flops(graph, node):
+  """Compute flops for MaxPoolGrad operation."""
+  _verify_conv_data_format(node)
+  #
+  # MaxPoolGrad declaration:
+  #   Inputs:
+  #     - orig_input  -- original input tensor (of max_pool)
+  #     - orig_output  -- original output tensor (of max_pool)
+  #     - grad --  gradient with respect to output of max_pool
+  #   Outputs:
+  #     - output -- gradient with respect to input of max_pool
+  #   Attributes:
+  #     - ksize
+  #     - strides
+  #     - padding
+  #     - data_format
+  # It computes MaxPool first, then one flop per each element of original output
+  #
+  kernel_shape = list(node.attr["ksize"].list.i)
+  kernel_area = _list_product(kernel_shape)
+  orig_out_shape = graph_util.tensor_shape_from_node_def_name(graph,
+                                                              node.input[1])
+  orig_out_shape.assert_is_fully_defined()
+  max_pool_ops = kernel_area * orig_out_shape.num_elements()
+  return ops.OpStats("flops", max_pool_ops + orig_out_shape.num_elements())
+
+
+@ops.RegisterStatistics("Conv2DBackpropInput", "flops")
+def _conv_2d_backprop_input_flops(graph, node):
+  """Compute flops for Conv2DBackpropInput operation."""
+  # Formula:
+  #  batch_size * image_x_dim * image_y_dim * kernel_x_dim * kernel_y_dim
+  #  * input_depth * output_depth * 2 / (image_x_stride * image_x_stride)
+  #
+  # Where:
+  # image_x_dim, image_y_dim and input_depth --- size of input to source (no
+  #   backprop) convolution, in other words they are sizes of backprop output.
+  # output_depth --- number of filters in the original convolution, thus
+  #   depth of backprop input.
+  # kernel_x_dim and kernel_y_dim --- sizes of filter in spatial dimension
+  # image_x_stride and image_x_stride --- strides of the convolution
+  #
+  _verify_conv_data_format(node)
+  # out_shape = [batch_size, image_y_dim, image_x_dim, input_depth]
+  out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  out_shape.assert_is_fully_defined()
+  # kernel_shape = [kernel_y_dim, kernel_x_dim, input_depth, output_depth]
+  kernel_shape = graph_util.tensor_shape_from_node_def_name(graph,
+                                                            node.input[1])
+  kernel_shape.assert_is_fully_defined()
+  # strides
+  strides_shape = list(node.attr["strides"].list.i)
+  strides_product = strides_shape[1] * strides_shape[2]
+  return ops.OpStats("flops",
+                     (2 * out_shape.num_elements()
+                      * kernel_shape.num_elements()
+                      / (out_shape.dims[-1].value * strides_product)))
+
+
+@ops.RegisterStatistics("Conv2DBackpropFilter", "flops")
+def _conv_2d_backprop_filter_flops(graph, node):
+  """Compute flops for Conv2DBackpropFilter operation."""
+  # Formula same as for Conv2DBackpropInput:
+  #  batch_size * image_x_dim * image_y_dim * kernel_x_dim * kernel_y_dim
+  #  * input_depth * output_depth * 2 / (image_x_stride * image_x_stride)
+  #
+  _verify_conv_data_format(node)
+  # image_shape = [batch_size, image_y_dim, image_x_dim, input_depth]
+  image_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  image_shape.assert_is_fully_defined()
+  # kernel_shape = [kernel_y_dim, kernel_x_dim, input_depth, output_depth]
+  kernel_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  kernel_shape.assert_is_fully_defined()
+  # strides
+  strides_shape = list(node.attr["strides"].list.i)
+  strides_product = strides_shape[1] * strides_shape[2]
+  return ops.OpStats("flops",
+                     (2 * image_shape.num_elements()
+                      * kernel_shape.num_elements()
+                      / (image_shape.dims[-1].value * strides_product)))
+
+################################################################################
+# Other ops
+################################################################################
+
+
+@ops.RegisterStatistics("AddN", "flops")
+def _add_n_flops(graph, node):
+  """Compute flops for AddN operation."""
+  if not node.input:
+    return _zero_flops(graph, node)
+  in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  in_shape.assert_is_fully_defined()
+  return ops.OpStats("flops", in_shape.num_elements() * (len(node.input) - 1))
+
+
+@ops.RegisterStatistics("MatMul", "flops")
+def _calc_mat_mul_flops(graph, node):
+  """Calculates the compute resources needed for MatMul."""
+  transpose_a = node.attr["transpose_a"].b
+  a_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  a_shape.assert_is_fully_defined()
+  if transpose_a:
+    k = int(a_shape[0])
+  else:
+    k = int(a_shape[1])
+  output_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  output_shape.assert_is_fully_defined()
+  output_count = np.prod(output_shape.as_list())
+  return ops.OpStats("flops", (k * output_count * 2))
+
+
+@ops.RegisterStatistics("BatchMatMul", "flops")
+@ops.RegisterStatistics("BatchMatMulV2", "flops")
+@ops.RegisterStatistics("BatchMatMulV3", "flops")
+def _calc_batch_mat_mul_flops(graph, node):
+  """Calculates the compute resources needed for BatchMatMul."""
+  transpose_a = node.attr["transpose_a"].b
+  a_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
+  a_shape.assert_is_fully_defined()
+  if transpose_a:
+    k = int(a_shape[-2])
+  else:
+    k = int(a_shape[-1])
+  output_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
+  output_shape.assert_is_fully_defined()
+  output_count = np.prod(output_shape.as_list())
+  return ops.OpStats("flops", (k * output_count * 2))

tensorflow/python/profiler/internal/flops_registry_test.py

@@ -0,0 +1,51 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.profiler.internal import flops_registry  # pylint: disable=unused-import
+
+
+class FlopsRegistryTest(test.TestCase):
+
+  @test_util.run_v1_only('Test requires a Graph and NodeDef inspection')
+  def testSimpleStatistics(self):
+    a = variables.Variable(random_ops.random_normal([25, 16]))
+    b = variables.Variable(random_ops.random_normal([16, 9]))
+    math_ops.matmul(a, b)
+    g = ops.get_default_graph()
+    for op in g.get_operations():
+      flops = ops.get_stats_for_node_def(g, op.node_def, 'flops').value
+      if op.name == 'MatMul':
+        self.assertEqual(7200, flops)
+
+  @test_util.run_v1_only('Test requires a Graph and NodeDef inspection')
+  def testTransposedStatistics(self):
+    a = variables.Variable(random_ops.random_normal([16, 25]))
+    b = variables.Variable(random_ops.random_normal([16, 9]))
+    math_ops.matmul(a, b, transpose_a=True)
+    g = ops.get_default_graph()
+    for op in g.get_operations():
+      flops = ops.get_stats_for_node_def(g, op.node_def, 'flops').value
+      if op.name == 'MatMul':
+        self.assertEqual(7200, flops)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/profiler/internal/model_analyzer_testlib.py

@@ -0,0 +1,113 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A test lib that defines some models."""
+import contextlib
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import rnn
+from tensorflow.python.ops import rnn_cell
+from tensorflow.python.ops import tensor_array_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.profiler import model_analyzer
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.util import _pywrap_tfprof as print_mdl
+from tensorflow.python.util import compat
+
+
+def BuildSmallModel():
+  """Build a small forward conv model."""
+  image = array_ops.zeros([2, 6, 6, 3])
+  _ = variable_scope.get_variable(
+      'ScalarW', [],
+      dtypes.float32,
+      initializer=init_ops.random_normal_initializer(stddev=0.001))
+  kernel = variable_scope.get_variable(
+      'DW', [3, 3, 3, 6],
+      dtypes.float32,
+      initializer=init_ops.random_normal_initializer(stddev=0.001))
+  x = nn_ops.conv2d(image, kernel, [1, 2, 2, 1], padding='SAME')
+  kernel = variable_scope.get_variable(
+      'DW2', [2, 2, 6, 12],
+      dtypes.float32,
+      initializer=init_ops.random_normal_initializer(stddev=0.001))
+  x = nn_ops.conv2d(x, kernel, [1, 2, 2, 1], padding='SAME')
+  return x
+
+
+def BuildFullModel():
+  """Build the full model with conv,rnn,opt."""
+  seq = []
+  for i in range(4):
+    with variable_scope.variable_scope('inp_%d' % i):
+      seq.append(array_ops.reshape(BuildSmallModel(), [2, 1, -1]))
+
+  cell = rnn_cell.BasicRNNCell(16)
+  out = rnn.dynamic_rnn(
+      cell, array_ops.concat(seq, axis=1), dtype=dtypes.float32)[0]
+
+  target = array_ops.ones_like(out)
+  loss = nn_ops.l2_loss(math_ops.reduce_mean(target - out))
+  sgd_op = gradient_descent.GradientDescentOptimizer(1e-2)
+  return sgd_op.minimize(loss)
+
+
+def BuildSplittableModel():
+  """Build a small model that can be run partially in each step."""
+  image = array_ops.zeros([2, 6, 6, 3])
+
+  kernel1 = variable_scope.get_variable(
+      'DW', [3, 3, 3, 6],
+      dtypes.float32,
+      initializer=init_ops.random_normal_initializer(stddev=0.001))
+  r1 = nn_ops.conv2d(image, kernel1, [1, 2, 2, 1], padding='SAME')
+
+  kernel2 = variable_scope.get_variable(
+      'DW2', [2, 3, 3, 6],
+      dtypes.float32,
+      initializer=init_ops.random_normal_initializer(stddev=0.001))
+  r2 = nn_ops.conv2d(image, kernel2, [1, 2, 2, 1], padding='SAME')
+
+  r3 = r1 + r2
+  return r1, r2, r3
+
+
+def SearchTFProfNode(node, name):
+  """Search a node in the tree."""
+  if node.name == name:
+    return node
+  for c in node.children:
+    r = SearchTFProfNode(c, name)
+    if r: return r
+  return None
+
+
+@contextlib.contextmanager
+def ProfilerFromFile(profile_file):
+  """Initialize a profiler from profile file."""
+  print_mdl.ProfilerFromFile(compat.as_bytes(profile_file))
+  profiler = model_analyzer.Profiler.__new__(model_analyzer.Profiler)
+  yield profiler
+  print_mdl.DeleteProfiler()
+
+
+def CheckAndRemoveDoc(profile):
+  assert 'Doc:' in profile
+  start_pos = profile.find('Profile:')
+  return profile[start_pos + 9:]

tensorflow/python/profiler/internal/print_model_analysis_test.py

@@ -0,0 +1,61 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""print_model_analysis test."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import test
+
+
+# pylint: disable=bad-whitespace
+# pylint: disable=bad-continuation
+TEST_OPTIONS = {
+    'max_depth': 10000,
+    'min_bytes': 0,
+    'min_micros': 0,
+    'min_params': 0,
+    'min_float_ops': 0,
+    'order_by': 'name',
+    'account_type_regexes': ['.*'],
+    'start_name_regexes': ['.*'],
+    'trim_name_regexes': [],
+    'show_name_regexes': ['.*'],
+    'hide_name_regexes': [],
+    'account_displayed_op_only': True,
+    'select': ['params'],
+    'output': 'stdout',
+}
+
+# pylint: enable=bad-whitespace
+# pylint: enable=bad-continuation
+
+
+class PrintModelAnalysisTest(test.TestCase):
+
+  def _BuildSmallModel(self):
+    image = array_ops.zeros([2, 6, 6, 3])
+    kernel = variable_scope.get_variable(
+        'DW', [6, 6, 3, 6],
+        dtypes.float32,
+        initializer=init_ops.random_normal_initializer(stddev=0.001))
+    x = nn_ops.conv2d(image, kernel, [1, 2, 2, 1], padding='SAME')
+    return x
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/profiler/internal/profiler_pywrap_impl.cc

@@ -0,0 +1,73 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/python/profiler/internal/profiler_pywrap_impl.h"
+
+#include <string>
+#include <variant>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/profiler/convert/xplane_to_trace_events.h"
+#include "xla/tsl/profiler/rpc/client/capture_profile.h"
+#include "xla/tsl/profiler/utils/session_manager.h"
+#include "tensorflow/core/platform/types.h"
+#include "tsl/profiler/lib/profiler_session.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+namespace pywrap {
+
+using tsl::profiler::GetRemoteSessionManagerOptionsLocked;
+
+absl::Status ProfilerSessionWrapper::Start(
+    const char* logdir,
+    const absl::flat_hash_map<std::string,
+                              std::variant<bool, int, std::string>>& options) {
+  auto opts = GetRemoteSessionManagerOptionsLocked(logdir, options);
+  session_ = tsl::ProfilerSession::Create(opts.profiler_options());
+  logdir_ = logdir;
+  return session_->Status();
+}
+
+absl::Status ProfilerSessionWrapper::Stop(std::string* result) {
+  if (session_ != nullptr) {
+    tensorflow::profiler::XSpace xspace;
+    absl::Status status = session_->CollectData(&xspace);
+    session_.reset();
+    tsl::profiler::ConvertXSpaceToTraceEventsString(xspace, result);
+    TF_RETURN_IF_ERROR(status);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ProfilerSessionWrapper::ExportToTensorBoard() {
+  if (!session_ || logdir_.empty()) {
+    return absl::OkStatus();
+  }
+  tensorflow::profiler::XSpace xspace;
+  absl::Status status;
+  status = session_->CollectData(&xspace);
+  session_.reset();
+  status = tsl::profiler::ExportToTensorBoard(xspace, logdir_);
+  return status;
+}
+
+}  // namespace pywrap
+}  // namespace profiler
+}  // namespace tensorflow

tensorflow/python/profiler/internal/profiler_pywrap_impl.h

@@ -0,0 +1,50 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_PROFILER_PYWRAP_IMPL_H_
+#define TENSORFLOW_PYTHON_PROFILER_INTERNAL_PROFILER_PYWRAP_IMPL_H_
+
+#include <memory>
+#include <string>
+#include <variant>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/types/variant.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/profiler/lib/profiler_session.h"
+
+namespace tensorflow {
+namespace profiler {
+namespace pywrap {
+
+class ProfilerSessionWrapper {
+ public:
+  absl::Status Start(
+      const char* logdir,
+      const absl::flat_hash_map<std::string,
+                                std::variant<bool, int, std::string>>& options);
+  absl::Status Stop(std::string* result);
+  absl::Status ExportToTensorBoard();
+
+ private:
+  std::unique_ptr<tsl::ProfilerSession> session_;
+  std::string logdir_;
+};
+
+}  // namespace pywrap
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_PROFILER_INTERNAL_PROFILER_PYWRAP_IMPL_H_

tensorflow/python/profiler/internal/profiler_wrapper.cc

@@ -0,0 +1,111 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <variant>
+
+#include "absl/status/status.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/profiler/rpc/profiler_server.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+#include "tensorflow/python/profiler/internal/profiler_pywrap_impl.h"
+#include "xprof/convert/repository.h"  // from @org_xprof
+#include "xprof/convert/tool_options.h"  // from @org_xprof
+#include "xprof/convert/xplane_to_tools_data.h"  // from @org_xprof
+
+namespace py = ::pybind11;
+
+namespace {
+
+using ::tensorflow::profiler::ToolOptions;
+using ::tensorflow::profiler::pywrap::ProfilerSessionWrapper;
+
+// These must be called under GIL because it reads Python objects. Reading
+// Python objects require GIL because the objects can be mutated by other Python
+// threads. In addition, Python objects are reference counted; reading py::dict
+// will increase its reference count.
+ToolOptions ToolOptionsFromPythonDict(const py::dict& dictionary) {
+  ToolOptions map;
+  for (const auto& item : dictionary) {
+    std::variant<bool, int, std::string> value;
+    try {
+      value = item.second.cast<bool>();
+    } catch (...) {
+      try {
+        value = item.second.cast<int>();
+      } catch (...) {
+        try {
+          value = item.second.cast<std::string>();
+        } catch (...) {
+          continue;
+        }
+      }
+    }
+    map.emplace(item.first.cast<std::string>(), value);
+  }
+  return map;
+}
+
+}  // namespace
+
+PYBIND11_MODULE(_pywrap_profiler, m) {
+  py::class_<ProfilerSessionWrapper> profiler_session_class(m,
+                                                            "ProfilerSession");
+  profiler_session_class.def(py::init<>())
+      .def("start",
+           [](ProfilerSessionWrapper& wrapper, const char* logdir,
+              const py::dict& options) {
+             absl::Status status;
+             ToolOptions tool_options = ToolOptionsFromPythonDict(options);
+             {
+               py::gil_scoped_release release;
+               status = wrapper.Start(logdir, tool_options);
+             }
+             // Py_INCREF and Py_DECREF must be called holding the GIL.
+             tensorflow::MaybeRaiseRegisteredFromStatus(status);
+           })
+      .def("stop",
+           [](ProfilerSessionWrapper& wrapper) {
+             std::string content;
+             absl::Status status;
+             {
+               py::gil_scoped_release release;
+               status = wrapper.Stop(&content);
+             }
+             // Py_INCREF and Py_DECREF must be called holding the GIL.
+             tensorflow::MaybeRaiseRegisteredFromStatus(status);
+             // The content is not valid UTF-8. It must be converted to bytes.
+             return py::bytes(content);
+           })
+      .def("export_to_tb", [](ProfilerSessionWrapper& wrapper) {
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = wrapper.ExportToTensorBoard();
+        }
+        // Py_INCREF and Py_DECREF must be called holding the GIL.
+        tensorflow::MaybeRaiseRegisteredFromStatus(status);
+      });
+
+  m.def("start_server", [](int port) {
+    auto profiler_server = std::make_unique<tsl::profiler::ProfilerServer>();
+    profiler_server->StartProfilerServer(port);
+    // Intentionally release profiler server. Should transfer ownership to
+    // caller instead.
+    profiler_server.release();
+  });
+};

tensorflow/python/profiler/internal/python_hooks.h

@@ -0,0 +1,48 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_
+#define TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_
+
+#include <memory>
+#include <stack>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/memory/memory.h"
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "xla/python/profiler/internal/python_hooks.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+
+using xla::profiler::PythonHooksOptions;  // NOLINT
+
+using xla::profiler::PythonTraceEntry;  // NOLINT
+
+using xla::profiler::PerThreadEvents;  // NOLINT
+
+using xla::profiler::PythonHookContext;  // NOLINT
+
+using xla::profiler::PythonHooks;  // NOLINT
+
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_

tensorflow/python/profiler/internal/pywrap_profiler_plugin.cc

@@ -0,0 +1,175 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "xla/pjrt/status_casters.h"
+#include "xla/tsl/platform/types.h"
+#include "xla/tsl/profiler/rpc/client/capture_profile.h"
+#include "xprof/convert/tool_options.h"  // from @org_xprof
+#include "xprof/pywrap/profiler_plugin_impl.h"  // from @org_xprof
+
+namespace py = ::pybind11;
+
+namespace {
+
+using ::tensorflow::profiler::ToolOptions;
+
+// These must be called under GIL because it reads Python objects. Reading
+// Python objects require GIL because the objects can be mutated by other Python
+// threads. In addition, Python objects are reference counted; reading py::dict
+// will increase its reference count.
+ToolOptions ToolOptionsFromPythonDict(const py::dict& dictionary) {
+  ToolOptions map;
+  // Copy keys to avoid iterator invalidation if the dict is mutated
+  // during implicit casts of values which can run Python user code.
+  std::vector<py::object> keys;
+  keys.reserve(dictionary.size());
+  for (const auto& item : dictionary) {
+    keys.push_back(py::reinterpret_borrow<py::object>(item.first));
+  }
+
+  for (const auto& key : keys) {
+    if (!dictionary.contains(key)) continue;
+
+    py::object py_value = dictionary[key];
+    std::variant<bool, int, std::string> value;
+    try {
+      value = py_value.cast<bool>();
+    } catch (...) {
+      try {
+        value = py_value.cast<int>();
+      } catch (...) {
+        try {
+          value = py_value.cast<std::string>();
+        } catch (...) {
+          continue;
+        }
+      }
+    }
+    try {
+      map.emplace(key.cast<std::string>(), value);
+    } catch (...) {
+      // Ignore keys that can't be cast to string
+    }
+  }
+  return map;
+}
+
+}  // namespace
+
+PYBIND11_MODULE(_pywrap_profiler_plugin, m, pybind11::mod_gil_not_used()) {
+  m.def(
+      "trace", [](const char* service_addr, const char* logdir,
+                  const char* worker_list, bool include_dataset_ops,
+                  int duration_ms, int num_tracing_attempts, py::dict options) {
+        absl::Status status;
+        ToolOptions tool_options = ToolOptionsFromPythonDict(options);
+        {
+          py::gil_scoped_release release;
+          status = tsl::profiler::CaptureRemoteTrace(
+              service_addr, logdir, worker_list, include_dataset_ops,
+              duration_ms, num_tracing_attempts, tool_options);
+        }
+        // Py_INCREF and Py_DECREF must be called holding the GIL.
+        xla::ThrowIfError(status);
+      });
+
+  m.def("monitor", [](const char* service_addr, int duration_ms,
+                      int monitoring_level, bool display_timestamp) {
+    std::string content;
+    absl::Status status;
+    {
+      py::gil_scoped_release release;
+      status =
+          xprof::pywrap::Monitor(service_addr, duration_ms, monitoring_level,
+                                 display_timestamp, &content);
+    }
+    // Py_INCREF and Py_DECREF must be called holding the GIL.
+    xla::ThrowIfError(status);
+    return content;
+  });
+
+  m.def(
+      "xspace_to_tools_data",
+      [](const py::list& xspace_path_list, const py::str& py_tool_name,
+         const py::dict options = py::dict()) {
+        std::vector<std::string> xspace_paths;
+        xspace_paths.reserve(xspace_path_list.size());
+        for (py::handle obj : xspace_path_list) {
+          std::string xspace_path = std::string(py::cast<py::str>(obj));
+          xspace_paths.push_back(xspace_path);
+        }
+
+        std::string tool_name = std::string(py_tool_name);
+        ToolOptions tool_options = ToolOptionsFromPythonDict(options);
+        absl::StatusOr<std::pair<std::string, bool>> result;
+
+        {
+          py::gil_scoped_release release;
+          result = xprof::pywrap::XSpaceToToolsData(xspace_paths, tool_name,
+                                                    tool_options);
+        }
+
+        if (!result.ok()) {
+          xla::ThrowIfError(result.status());
+        }
+        return py::make_tuple(py::bytes(result->first),
+                              py::bool_(result->second));
+      },
+      // TODO: consider defaulting `xspace_path_list` to empty list, since
+      // this parameter is only used for two of the tools.
+      py::arg(), py::arg(), py::arg() = py::dict());
+
+  m.def(
+      "xspace_to_tools_data_from_byte_string",
+      [](const py::list& xspace_string_list, const py::list& filenames_list,
+         const py::str& py_tool_name, const py::dict options = py::dict()) {
+        std::vector<std::string> xspace_strings;
+        xspace_strings.reserve(xspace_string_list.size());
+        for (py::handle obj : xspace_string_list) {
+          xspace_strings.push_back(std::string(py::cast<py::bytes>(obj)));
+        }
+
+        std::vector<std::string> xspace_paths;
+        xspace_paths.reserve(filenames_list.size());
+        for (py::handle obj : filenames_list) {
+          xspace_paths.push_back(std::string(py::cast<py::str>(obj)));
+        }
+
+        std::string tool_name = std::string(py_tool_name);
+        ToolOptions tool_options = ToolOptionsFromPythonDict(options);
+
+        absl::StatusOr<std::pair<std::string, bool>> result;
+        {
+          py::gil_scoped_release release;
+          result = xprof::pywrap::XSpaceToToolsDataFromByteString(
+              xspace_strings, xspace_paths, tool_name, tool_options);
+        }
+
+        if (!result.ok()) {
+          xla::ThrowIfError(result.status());
+        }
+        return py::make_tuple(py::bytes(result->first),
+                              py::bool_(result->second));
+      },
+      py::arg(), py::arg(), py::arg(), py::arg() = py::dict());
+};

tensorflow/python/profiler/internal/run_metadata_test.py

@@ -0,0 +1,242 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""test the RunMetadata proto."""
+
+from collections import defaultdict
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.profiler import option_builder
+
+# pylint: disable=g-bad-import-order
+# XXX: this depends on pywrap_tensorflow and must come later
+from tensorflow.python.profiler import model_analyzer
+from tensorflow.python.profiler.internal import model_analyzer_testlib as lib
+
+SIZE = 1300
+builder = option_builder.ProfileOptionBuilder
+
+
+def _extract_node(run_meta, node_name):
+  ret = defaultdict(list)
+  for dev_stat in run_meta.step_stats.dev_stats:
+    dev = dev_stat.device.lower()
+    if dev.find('cpu:') > 0:
+      dev = dev[dev.find('cpu:'):]
+    elif dev.find('gpu:') > 0:
+      dev = dev[dev.find('gpu:'):]
+    elif '/host:cpu' not in dev:
+      assert False, 'Unrecognized device name: %s' % dev
+
+    for node_stat in dev_stat.node_stats:
+      nname = node_stat.node_name
+      if nname.find(':') > 0:
+        nname = nname[:nname.find(':')]
+      if nname == node_name:
+        ret[dev].append(node_stat)
+  return ret
+
+
+def _run_model():
+  x = random_ops.random_normal(shape=[1, SIZE])
+  w = random_ops.random_normal(shape=[SIZE, 2 * SIZE])
+  y = math_ops.matmul(x, w)
+
+  config = config_pb2.ConfigProto()
+  config.graph_options.rewrite_options.arithmetic_optimization = (
+      rewriter_config_pb2.RewriterConfig.OFF)
+  with session.Session(config=config) as sess:
+    run_metadata = config_pb2.RunMetadata()
+    opts = builder.time_and_memory()
+    opts['min_micros'] = 0
+    opts['min_bytes'] = 0
+    opts['order_by'] = 'name'
+    opts['output'] = 'none'
+    _ = sess.run(
+        y,
+        options=config_pb2.RunOptions(
+            trace_level=config_pb2.RunOptions.SOFTWARE_TRACE),
+        run_metadata=run_metadata)
+    tfprof_node = model_analyzer.profile(
+        sess.graph, run_meta=run_metadata, options=opts)
+
+    return tfprof_node, run_metadata
+
+
+def _run_loop_model():
+  config = config_pb2.ConfigProto()
+  # Grappler might fuse MatMul with BiasAdd in remapper optimizer.
+  config.graph_options.rewrite_options.remapping = (
+      rewriter_config_pb2.RewriterConfig.OFF)
+  with session.Session(config=config) as sess:
+    x = lib.BuildFullModel()
+
+    sess.run(variables.global_variables_initializer())
+    run_meta = config_pb2.RunMetadata()
+    _ = sess.run(
+        x,
+        options=config_pb2.RunOptions(
+            trace_level=config_pb2.RunOptions.SOFTWARE_TRACE),
+        run_metadata=run_meta)
+
+    opts = builder.time_and_memory()
+    opts['order_by'] = 'name'
+    opts['output'] = 'none'
+
+    tfprof_node = model_analyzer.profile(sess.graph, run_meta, options=opts)
+    return tfprof_node, run_meta
+
+
+class RunMetadataTest(test.TestCase):
+
+  # This test requires HARDWARE_TRACE or FULL_TRACE to be specified to
+  # work as expected. Since we now run this test with SOFTWARE_TRACE
+  # (see _run_model routine above), this test will / should fail since
+  # GPU device tracers are not enabled
+  @test_util.run_deprecated_v1
+  def testGPU(self):
+    if not test.is_gpu_available(cuda_only=True):
+      return
+
+    gpu_dev = test.gpu_device_name()
+    ops.reset_default_graph()
+    with ops.device(gpu_dev):
+      tfprof_node, run_meta = _run_model()
+      self.assertEqual(tfprof_node.children[0].name, 'MatMul')
+      self.assertGreater(tfprof_node.children[0].exec_micros, 10)
+
+    ret = _extract_node(run_meta, 'MatMul')
+    self.assertEqual(len(ret['gpu:0']), 1)
+
+  @test_util.run_deprecated_v1
+  def testAllocationHistory(self):
+    if not test.is_gpu_available(cuda_only=True):
+      return
+
+    gpu_dev = test.gpu_device_name()
+    ops.reset_default_graph()
+    with ops.device(gpu_dev):
+      _, run_meta = _run_model()
+
+    mm = _extract_node(run_meta, 'MatMul')['gpu:0'][0]
+    mm_allocs = mm.memory[0].allocation_records
+    # has allocation and deallocation.
+    self.assertEqual(len(mm_allocs), 2)
+    # first allocated.
+    self.assertGreater(mm_allocs[1].alloc_micros, mm_allocs[0].alloc_micros)
+    self.assertGreater(mm_allocs[0].alloc_bytes, 0)
+    # Then deallocated.
+    self.assertLess(mm_allocs[1].alloc_bytes, 0)
+    # All memory deallocated.
+    self.assertEqual(mm_allocs[0].alloc_bytes + mm_allocs[1].alloc_bytes, 0)
+
+    rand = _extract_node(run_meta,
+                         'random_normal/RandomStandardNormal')['gpu:0'][0]
+    random_allocs = rand.memory[0].allocation_records
+    # random normal must allocated first since matmul depends on it.
+    self.assertLess(random_allocs[0].alloc_micros, mm.all_start_micros)
+    # deallocates the memory after matmul started.
+    self.assertGreater(random_allocs[1].alloc_micros, mm.all_start_micros)
+
+  @test_util.run_deprecated_v1
+  def testCPU(self):
+    ops.reset_default_graph()
+    with ops.device('/cpu:0'):
+      tfprof_node, run_meta = _run_model()
+      self.assertEqual(tfprof_node.children[0].name, 'MatMul')
+      self.assertGreater(tfprof_node.children[0].exec_micros, 0)
+
+    ret = _extract_node(run_meta, 'MatMul')
+    self.assertEqual(len(ret['cpu:0']), 1)
+
+    ret = _extract_node(run_meta, 'MatMul:MatMul')
+    self.assertEqual(len(ret), 0)
+
+  @test_util.run_v1_only('b/120545219')
+  def testLoopCPU(self):
+    ops.reset_default_graph()
+    with ops.device('/cpu:0'):
+      tfprof_node, run_meta = _run_loop_model()
+      # The while-loop caused a node to appear 4 times in scheduling.
+      ret = _extract_node(run_meta, 'rnn/while/basic_rnn_cell/MatMul')
+      self.assertEqual(len(ret['cpu:0']), 4)
+
+      total_cpu_execs = 0
+      for node in ret['cpu:0']:
+        total_cpu_execs += node.op_end_rel_micros
+
+      mm_node = lib.SearchTFProfNode(tfprof_node,
+                                     'rnn/while/basic_rnn_cell/MatMul')
+
+      self.assertEqual(mm_node.run_count, 4)
+      self.assertEqual(mm_node.cpu_exec_micros, total_cpu_execs)
+      self.assertEqual(mm_node.exec_micros, total_cpu_execs)
+
+  def testGradientGraph(self):
+    # Note: Please don't just adjust the test to make it pass.
+    # The code view logic depends on it.
+    ops.reset_default_graph()
+    _, _ = _run_loop_model()
+    graph = ops.get_default_graph()
+    forward_op = set()
+    backward_op = set()
+    back_to_forward = {}
+    for op in graph.get_operations():
+      if op.name.find('gradients/') > 0 and op.name.find('_grad/') > 0:
+        backward_op.add(op.name)
+        idx1 = op.name.find('gradients/') + 10
+        idx2 = op.name.find('_grad/')
+        back_to_forward[op.name] = op.name[idx1:idx2]
+      else:
+        forward_op.add(op.name)
+
+    for _, f in back_to_forward.items():
+      self.assertTrue(f in forward_op)
+
+  # This test requires HARDWARE_TRACE or FULL_TRACE to be specified to
+  # work as expected. Since we now run this test with SOFTWARE_TRACE
+  # (see _run_model routine above), this test will / should fail since
+  # GPU device tracers are not enabled
+  @test.disable_with_predicate(
+      pred=test.is_built_with_rocm,
+      skip_message='Test fails on ROCm when run without FULL_TRACE')
+  def testLoopGPU(self):
+    if not test.is_gpu_available():
+      return
+
+    ops.reset_default_graph()
+    with ops.device('/device:GPU:0'):
+      _, run_meta = _run_loop_model()
+      # The while-loop caused a node to appear 4 times in scheduling.
+      ret = _extract_node(run_meta, 'rnn/while/basic_rnn_cell/MatMul')
+      self.assertEqual(len(ret['gpu:0']), 4, '%s' % run_meta)
+
+      total_cpu_execs = 0
+      for node in ret['gpu:0']:
+        total_cpu_execs += node.op_end_rel_micros
+
+        self.assertGreaterEqual(
+            len(ret['gpu:0/stream:all']), 4, '%s' % run_meta)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/profiler/internal/traceme_wrapper.cc

@@ -0,0 +1,49 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/python/profiler/internal/traceme_wrapper.h"
+
+#include "pybind11/attr.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "xla/python/profiler/internal/traceme_state.h"
+
+namespace py = ::pybind11;
+
+using ::xla::profiler::TraceMeWrapper;
+
+// Returns true if TraceMe is enabled.
+// This is a low-overhead function that can be called frequently.
+static PyObject* traceme_enabled(PyObject* self, PyObject* args) {
+  if (xla::profiler::traceme_enabled) {
+    Py_RETURN_TRUE;
+  }
+  Py_RETURN_FALSE;
+}
+
+static PyMethodDef traceme_method_def = {"traceme_enabled", traceme_enabled,
+                                         METH_NOARGS,
+                                         "Returns true if TraceMe is enabled."};
+
+PYBIND11_MODULE(_pywrap_traceme, m) {
+  py::class_<TraceMeWrapper>(m, "TraceMe", py::module_local())
+      .def(py::init<const py::str&, const py::kwargs&>())
+      .def("SetMetadata", &TraceMeWrapper::SetMetadata)
+      .def("Stop", &TraceMeWrapper::Stop);
+
+  py::object module_name = m.attr("__name__");
+  m.attr("traceme_enabled") =
+      py::reinterpret_steal<py::object>(PyCFunction_NewEx(
+          &traceme_method_def, /*self=*/nullptr, module_name.ptr()));
+};