Move tsl/profiler/utils to xla/tsl/profiler/utils

PiperOrigin-RevId: 677907959
openxla · Sep 27, 2024 · 73e7685 · 73e7685
1 parent b7dbeb6
commit 73e7685
Show file tree

Hide file tree

Showing 104 changed files with 447 additions and 447 deletions.
diff --git a/third_party/tsl/tsl/platform/cloud/BUILD b/third_party/tsl/tsl/platform/cloud/BUILD
@@ -396,9 +396,9 @@ tsl_cc_test(
         "//tsl/platform:strcat",
         "//tsl/platform:test",
         "//tsl/platform:test_main",
-        "//tsl/profiler/utils:time_utils_impl",
         "@xla//xla/tsl/lib/core:status_test_util",
         "@xla//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
+        "@xla//xla/tsl/profiler/utils:time_utils_impl",
     ],
 )
 

diff --git a/third_party/tsl/tsl/profiler/lib/BUILD b/third_party/tsl/tsl/profiler/lib/BUILD
@@ -38,7 +38,7 @@ filegroup(
         "scoped_memory_debug_annotation.h",
         "traceme.h",
         "traceme_encode.h",
-        "//tsl/profiler/utils:mobile_srcs_no_runtime",
+        "@xla//xla/tsl/profiler/utils:mobile_srcs_no_runtime",
     ],
     compatible_with = get_compatible_with_portable(),
     visibility = ["//visibility:public"],
@@ -221,7 +221,7 @@ cc_library(
         "//tsl/platform:platform_port",
         "//tsl/platform:status",
         "@xla//xla/tsl/profiler/convert:post_process_single_host_xplane",
-        "//tsl/profiler/utils:time_utils",
+        "@xla//xla/tsl/profiler/utils:time_utils",
     ]),
     alwayslink = True,
 )
@@ -268,11 +268,11 @@ cc_library(
         "//tsl/platform:logging",
         "//tsl/platform:macros",
         "//tsl/platform:types",
-        "//tsl/profiler/utils:no_init",
         "@com_google_absl//absl/strings",
+        "@xla//xla/tsl/profiler/utils:no_init",
     ] + if_not_android([
         "@xla//xla/tsl/profiler/backends/cpu:traceme_recorder",
-        "//tsl/profiler/utils:time_utils",
+        "@xla//xla/tsl/profiler/utils:time_utils",
     ]),
 )
 

diff --git a/third_party/tsl/tsl/profiler/lib/profiler_session.cc b/third_party/tsl/tsl/profiler/lib/profiler_session.cc
@@ -28,12 +28,12 @@ limitations under the License.
 
 #if !defined(IS_MOBILE_PLATFORM)
 #include "xla/tsl/profiler/convert/post_process_single_host_xplane.h"
+#include "xla/tsl/profiler/utils/time_utils.h"
 #include "tsl/platform/host_info.h"
 #include "tsl/profiler/lib/profiler_collection.h"
 #include "tsl/profiler/lib/profiler_factory.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/lib/profiler_lock.h"
-#include "tsl/profiler/utils/time_utils.h"
 #endif
 
 namespace tsl {

diff --git a/third_party/tsl/tsl/profiler/lib/traceme.h b/third_party/tsl/tsl/profiler/lib/traceme.h
@@ -21,14 +21,14 @@ limitations under the License.
 #include <utility>
 
 #include "absl/strings/string_view.h"
+#include "xla/tsl/profiler/utils/no_init.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/macros.h"
 #include "tsl/profiler/lib/traceme_encode.h"  // IWYU pragma: export
-#include "tsl/profiler/utils/no_init.h"
 
 #if !defined(IS_MOBILE_PLATFORM)
 #include "xla/tsl/profiler/backends/cpu/traceme_recorder.h"
-#include "tsl/profiler/utils/time_utils.h"
+#include "xla/tsl/profiler/utils/time_utils.h"
 #endif
 
 namespace tsl {

diff --git a/xla/backends/profiler/cpu/BUILD b/xla/backends/profiler/cpu/BUILD
@@ -36,15 +36,15 @@ cc_library(
         "//xla/tsl/profiler/backends/cpu:host_tracer_utils",
         "//xla/tsl/profiler/backends/cpu:threadpool_listener",
         "//xla/tsl/profiler/backends/cpu:traceme_recorder",
+        "//xla/tsl/profiler/utils:time_utils",
+        "//xla/tsl/profiler/utils:xplane_schema",
+        "//xla/tsl/profiler/utils:xplane_utils",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/profiler/lib:profiler_collection",
         "@tsl//tsl/profiler/lib:profiler_interface",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@tsl//tsl/profiler/utils:time_utils",
-        "@tsl//tsl/profiler/utils:xplane_schema",
-        "@tsl//tsl/profiler/utils:xplane_utils",
     ],
 )
 
@@ -96,14 +96,14 @@ cc_library(
         ":metadata_utils",
         "//xla/service:hlo_proto_cc",
         "//xla/service:xla_debug_info_manager",
+        "//xla/tsl/profiler/utils:xplane_builder",
+        "//xla/tsl/profiler/utils:xplane_schema",
+        "//xla/tsl/profiler/utils:xplane_utils",
         "@com_google_absl//absl/status",
         "@tsl//tsl/profiler/lib:profiler_factory",
         "@tsl//tsl/profiler/lib:profiler_interface",
         "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@tsl//tsl/profiler/utils:xplane_builder",
-        "@tsl//tsl/profiler/utils:xplane_schema",
-        "@tsl//tsl/profiler/utils:xplane_utils",
     ],
     alwayslink = True,
 )
@@ -117,9 +117,9 @@ cc_library(
     deps = [
         "//xla/service:hlo_proto_cc",
         "//xla/tsl/profiler/convert:xla_op_utils",
+        "//xla/tsl/profiler/utils:xplane_builder",
+        "//xla/tsl/profiler/utils:xplane_schema",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@tsl//tsl/profiler/utils:xplane_builder",
-        "@tsl//tsl/profiler/utils:xplane_schema",
     ],
 )
 
@@ -129,6 +129,10 @@ xla_cc_test(
     deps = [
         ":host_tracer_impl",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/profiler/utils:tf_xplane_visitor",
+        "//xla/tsl/profiler/utils:timespan",
+        "//xla/tsl/profiler/utils:xplane_schema",
+        "//xla/tsl/profiler/utils:xplane_visitor",
         "@com_google_absl//absl/types:optional",
         "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:blocking_counter",
@@ -138,9 +142,5 @@ xla_cc_test(
         "@tsl//tsl/profiler/lib:profiler_interface",
         "@tsl//tsl/profiler/lib:traceme",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@tsl//tsl/profiler/utils:tf_xplane_visitor",
-        "@tsl//tsl/profiler/utils:timespan",
-        "@tsl//tsl/profiler/utils:xplane_schema",
-        "@tsl//tsl/profiler/utils:xplane_visitor",
     ],
 )
diff --git a/xla/backends/profiler/cpu/host_tracer.cc b/xla/backends/profiler/cpu/host_tracer.cc
@@ -24,13 +24,13 @@ limitations under the License.
 #include "xla/tsl/profiler/backends/cpu/host_tracer_utils.h"
 #include "xla/tsl/profiler/backends/cpu/threadpool_listener.h"
 #include "xla/tsl/profiler/backends/cpu/traceme_recorder.h"
+#include "xla/tsl/profiler/utils/time_utils.h"
+#include "xla/tsl/profiler/utils/xplane_schema.h"
+#include "xla/tsl/profiler/utils/xplane_utils.h"
 #include "tsl/platform/errors.h"
 #include "tsl/profiler/lib/profiler_collection.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
-#include "tsl/profiler/utils/time_utils.h"
-#include "tsl/profiler/utils/xplane_schema.h"
-#include "tsl/profiler/utils/xplane_utils.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/cpu/host_tracer_test.cc b/xla/backends/profiler/cpu/host_tracer_test.cc
@@ -23,6 +23,10 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/types/optional.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/profiler/utils/tf_xplane_visitor.h"
+#include "xla/tsl/profiler/utils/timespan.h"
+#include "xla/tsl/profiler/utils/xplane_schema.h"
+#include "xla/tsl/profiler/utils/xplane_visitor.h"
 #include "tsl/platform/blocking_counter.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/test.h"
@@ -31,10 +35,6 @@ limitations under the License.
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/lib/traceme.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
-#include "tsl/profiler/utils/tf_xplane_visitor.h"
-#include "tsl/profiler/utils/timespan.h"
-#include "tsl/profiler/utils/xplane_schema.h"
-#include "tsl/profiler/utils/xplane_visitor.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/cpu/metadata_collector.cc b/xla/backends/profiler/cpu/metadata_collector.cc
@@ -22,13 +22,13 @@ limitations under the License.
 #include "xla/backends/profiler/cpu/metadata_utils.h"
 #include "xla/service/hlo.pb.h"
 #include "xla/service/xla_debug_info_manager.h"
+#include "xla/tsl/profiler/utils/xplane_builder.h"
+#include "xla/tsl/profiler/utils/xplane_schema.h"
+#include "xla/tsl/profiler/utils/xplane_utils.h"
 #include "tsl/profiler/lib/profiler_factory.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/profiler_options.pb.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
-#include "tsl/profiler/utils/xplane_builder.h"
-#include "tsl/profiler/utils/xplane_schema.h"
-#include "tsl/profiler/utils/xplane_utils.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/cpu/metadata_utils.h b/xla/backends/profiler/cpu/metadata_utils.h
@@ -18,9 +18,9 @@ limitations under the License.
 
 #include "xla/service/hlo.pb.h"
 #include "xla/tsl/profiler/convert/xla_op_utils.h"
+#include "xla/tsl/profiler/utils/xplane_builder.h"
+#include "xla/tsl/profiler/utils/xplane_schema.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
-#include "tsl/profiler/utils/xplane_builder.h"
-#include "tsl/profiler/utils/xplane_schema.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/gpu/BUILD b/xla/backends/profiler/gpu/BUILD
@@ -29,6 +29,7 @@ tsl_gpu_library(
     copts = tf_profiler_copts() + tsl_copts(),
     deps = [
         ":cupti_utils",
+        "//xla/tsl/profiler/utils:time_utils",
         "//xla/tsl/util:env_var",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -39,7 +40,6 @@ tsl_gpu_library(
         "@tsl//tsl/profiler/lib:profiler_factory",
         "@tsl//tsl/profiler/lib:profiler_interface",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@tsl//tsl/profiler/utils:time_utils",
     ] + if_cuda([
         # keep sorted
         ":cupti_buffer_events",
@@ -117,7 +117,7 @@ xla_test(
         ":cupti_wrapper",
         ":mock_cupti",
         "@com_google_absl//absl/memory",
-        "@tsl//tsl/profiler/utils:time_utils",
+        "//xla/tsl/profiler/utils:time_utils",
     ]),
 )
 
@@ -175,6 +175,8 @@ tsl_gpu_library(
         ":cupti_utils",
         ":nvtx_utils",
         "//xla/tsl/profiler/backends/cpu:annotation_stack",
+        "//xla/tsl/profiler/utils:lock_free_queue",
+        "//xla/tsl/profiler/utils:per_thread",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
@@ -186,8 +188,6 @@ tsl_gpu_library(
         "@tsl//tsl/platform:macros",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/platform:types",
-        "@tsl//tsl/profiler/utils:lock_free_queue",
-        "@tsl//tsl/profiler/utils:per_thread",
     ],
 )
 
@@ -232,6 +232,10 @@ tsl_gpu_library(
     deps = [
         "//xla/stream_executor/rocm:roctracer_wrapper",
         "//xla/tsl/profiler/backends/cpu:annotation_stack",
+        "//xla/tsl/profiler/utils:parse_annotation",
+        "//xla/tsl/profiler/utils:xplane_builder",
+        "//xla/tsl/profiler/utils:xplane_schema",
+        "//xla/tsl/profiler/utils:xplane_utils",
         "//xla/tsl/util:env_var",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -251,10 +255,6 @@ tsl_gpu_library(
         "@tsl//tsl/platform:types",
         "@tsl//tsl/profiler/lib:profiler_factory",
         "@tsl//tsl/profiler/lib:profiler_interface",
-        "@tsl//tsl/profiler/utils:parse_annotation",
-        "@tsl//tsl/profiler/utils:xplane_builder",
-        "@tsl//tsl/profiler/utils:xplane_schema",
-        "@tsl//tsl/profiler/utils:xplane_utils",
     ],
 )
 
@@ -275,6 +275,7 @@ tsl_gpu_library(
         ":rocm_collector",
         "//xla/stream_executor/rocm:roctracer_wrapper",
         "//xla/tsl/profiler/backends/cpu:annotation_stack",
+        "//xla/tsl/profiler/utils:time_utils",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -288,7 +289,6 @@ tsl_gpu_library(
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/platform:status",
         "@tsl//tsl/platform:types",
-        "@tsl//tsl/profiler/utils:time_utils",
     ],
 )
 
@@ -313,6 +313,12 @@ tsl_gpu_library(
     deps = [
         ":cupti_buffer_events",
         ":cupti_interface",
+        "//xla/tsl/profiler/utils:lock_free_queue",
+        "//xla/tsl/profiler/utils:parse_annotation",
+        "//xla/tsl/profiler/utils:trace_utils",
+        "//xla/tsl/profiler/utils:xplane_builder",
+        "//xla/tsl/profiler/utils:xplane_schema",
+        "//xla/tsl/profiler/utils:xplane_utils",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/hash",
@@ -321,12 +327,6 @@ tsl_gpu_library(
         "@tsl//tsl/platform:mutex",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@tsl//tsl/profiler/utils:lock_free_queue",
-        "@tsl//tsl/profiler/utils:parse_annotation",
-        "@tsl//tsl/profiler/utils:trace_utils",
-        "@tsl//tsl/profiler/utils:xplane_builder",
-        "@tsl//tsl/profiler/utils:xplane_schema",
-        "@tsl//tsl/profiler/utils:xplane_utils",
     ] + if_cuda([
         "//xla/tsl/cuda:cupti",
         "//xla/tsl/cuda",
@@ -341,6 +341,8 @@ tsl_gpu_library(
     visibility = ["//visibility:public"],
     deps = [
         ":cupti_interface",
+        "//xla/tsl/profiler/utils:buffer_pool",
+        "//xla/tsl/profiler/utils:lock_free_queue",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -351,8 +353,6 @@ tsl_gpu_library(
         "@tsl//tsl/platform:mutex",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/platform:thread_annotations",
-        "@tsl//tsl/profiler/utils:buffer_pool",
-        "@tsl//tsl/profiler/utils:lock_free_queue",
     ] + if_cuda(["//xla/tsl/cuda:cupti"]),
 )
 

diff --git a/xla/backends/profiler/gpu/cupti_buffer_events.h b/xla/backends/profiler/gpu/cupti_buffer_events.h
@@ -31,10 +31,10 @@ limitations under the License.
 #include "absl/container/node_hash_set.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "xla/tsl/profiler/utils/buffer_pool.h"
+#include "xla/tsl/profiler/utils/lock_free_queue.h"
 #include "tsl/platform/mutex.h"
 #include "tsl/platform/thread_annotations.h"
-#include "tsl/profiler/utils/buffer_pool.h"
-#include "tsl/profiler/utils/lock_free_queue.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/gpu/cupti_collector.cc b/xla/backends/profiler/gpu/cupti_collector.cc
@@ -28,15 +28,15 @@ limitations under the License.
 #include "third_party/gpus/cuda/extras/CUPTI/include/cupti_activity.h"
 #include "third_party/gpus/cuda/include/cuda.h"
 #include "third_party/gpus/cuda/include/cuda_occupancy.h"
+#include "xla/tsl/profiler/utils/parse_annotation.h"
+#include "xla/tsl/profiler/utils/trace_utils.h"
+#include "xla/tsl/profiler/utils/xplane_builder.h"
+#include "xla/tsl/profiler/utils/xplane_schema.h"
+#include "xla/tsl/profiler/utils/xplane_utils.h"
 #include "tsl/platform/abi.h"
 #include "tsl/platform/host_info.h"
 #include "tsl/platform/mem.h"
 #include "tsl/platform/mutex.h"
-#include "tsl/profiler/utils/parse_annotation.h"
-#include "tsl/profiler/utils/trace_utils.h"
-#include "tsl/profiler/utils/xplane_builder.h"
-#include "tsl/profiler/utils/xplane_schema.h"
-#include "tsl/profiler/utils/xplane_utils.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/gpu/cupti_error_manager_test.cc b/xla/backends/profiler/gpu/cupti_error_manager_test.cc
@@ -27,8 +27,8 @@ limitations under the License.
 #include "xla/backends/profiler/gpu/cupti_tracer.h"
 #include "xla/backends/profiler/gpu/cupti_wrapper.h"
 #include "xla/backends/profiler/gpu/mock_cupti.h"
+#include "xla/tsl/profiler/utils/time_utils.h"
 #include "tsl/platform/test.h"
-#include "tsl/profiler/utils/time_utils.h"
 
 namespace xla {
 namespace profiler {

diff --git a/xla/backends/profiler/gpu/cupti_tracer.cc b/xla/backends/profiler/gpu/cupti_tracer.cc
@@ -31,11 +31,11 @@ limitations under the License.
 #include "xla/backends/profiler/gpu/cupti_interface.h"
 #include "xla/backends/profiler/gpu/nvtx_utils.h"
 #include "xla/tsl/profiler/backends/cpu/annotation_stack.h"
+#include "xla/tsl/profiler/utils/per_thread.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/host_info.h"
 #include "tsl/platform/logging.h"
-#include "tsl/profiler/utils/per_thread.h"
 
 namespace xla {
 namespace profiler {