Enable profiler tracing support for threadpools

PiperOrigin-RevId: 621967537

third_party/tsl/tsl/platform/default/BUILD

@@ -465,6 +465,7 @@ cc_library(
         "//tsl/platform:strcat",
         "//tsl/platform:stringpiece",
         "//tsl/platform:types",
+        "//tsl/profiler/backends/cpu:threadpool_listener_state",
     ],
     alwayslink = True,
 )

third_party/tsl/tsl/platform/default/tracing_impl.h

@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_TSL_PLATFORM_DEFAULT_TRACING_IMPL_H_
 #define TENSORFLOW_TSL_PLATFORM_DEFAULT_TRACING_IMPL_H_
 
+#include "tsl/profiler/backends/cpu/threadpool_listener_state.h"
+
 // Stub implementations of tracing functionality.
 
 // Definitions that do nothing for platforms that don't have underlying thread
@@ -33,7 +35,9 @@ limitations under the License.
 namespace tsl {
 namespace tracing {
 
-inline bool EventCollector::IsEnabled() { return false; }
+inline bool EventCollector::IsEnabled() {
+  return tsl::profiler::threadpool_listener::IsEnabled();
+}
 
 }  // namespace tracing
 }  // namespace tsl

third_party/tsl/tsl/profiler/backends/cpu/BUILD

@@ -128,3 +128,39 @@ cc_library(
         "@com_google_absl//absl/strings",
     ],
 )
+
+cc_library(
+    name = "threadpool_listener",
+    srcs = ["threadpool_listener.cc"],
+    hdrs = ["threadpool_listener.h"],
+    visibility = internal_visibility([
+        "//tensorflow/python:__pkg__",
+        "//tsl/platform/cloud:__pkg__",
+        "//tsl/profiler:__pkg__",
+        "//tsl/profiler:internal",
+        "//tsl/profiler:xla_internal",
+    ]),
+    deps = [
+        ":threadpool_listener_state",
+        ":traceme_recorder",
+        "//tsl/platform:logging",
+        "//tsl/platform:tracing",
+        "//tsl/platform:types",
+        "//tsl/profiler/lib:context_types_hdrs",
+        "//tsl/profiler/lib:profiler_interface",
+        "//tsl/profiler/lib:traceme_encode",
+        "//tsl/profiler/utils:time_utils",
+        "//tsl/profiler/utils:xplane_schema",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+cc_library(
+    name = "threadpool_listener_state",
+    srcs = ["threadpool_listener_state.cc"],
+    hdrs = ["threadpool_listener_state.h"],
+    visibility = internal_visibility([
+        "//tsl/platform:__subpackages__",
+    ]),
+)

third_party/tsl/tsl/profiler/backends/cpu/threadpool_listener.cc

@@ -0,0 +1,99 @@
+/* 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 "tsl/profiler/backends/cpu/threadpool_listener.h"
+
+#include <cstdint>
+#include <memory>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tsl/platform/logging.h"
+#include "tsl/platform/tracing.h"
+#include "tsl/platform/types.h"
+#include "tsl/profiler/backends/cpu/threadpool_listener_state.h"
+#include "tsl/profiler/backends/cpu/traceme_recorder.h"
+#include "tsl/profiler/lib/context_types.h"
+#include "tsl/profiler/lib/traceme_encode.h"
+#include "tsl/profiler/utils/time_utils.h"
+#include "tsl/profiler/utils/xplane_schema.h"
+
+namespace tsl {
+namespace profiler {
+namespace {
+
+void RegisterThreadpoolEventCollector(ThreadpoolEventCollector* collector) {
+  tracing::SetEventCollector(tracing::EventCategory::kScheduleClosure,
+                             collector);
+  tracing::SetEventCollector(tracing::EventCategory::kRunClosure, collector);
+}
+
+void UnregisterThreadpoolEventCollector() {
+  tracing::SetEventCollector(tracing::EventCategory::kScheduleClosure, nullptr);
+  tracing::SetEventCollector(tracing::EventCategory::kRunClosure, nullptr);
+}
+
+}  // namespace
+
+void ThreadpoolEventCollector::RecordEvent(uint64 arg) const {
+  int64_t now = GetCurrentTimeNanos();
+  TraceMeRecorder::Record(
+      {TraceMeEncode(kThreadpoolListenerRecord,
+                     {{"_pt", ContextType::kThreadpoolEvent}, {"_p", arg}}),
+       now, now});
+}
+void ThreadpoolEventCollector::StartRegion(uint64 arg) const {
+  int64_t now = GetCurrentTimeNanos();
+  TraceMeRecorder::Record(
+      {TraceMeEncode(kThreadpoolListenerStartRegion,
+                     {{"_ct", ContextType::kThreadpoolEvent}, {"_c", arg}}),
+       now, now});
+}
+void ThreadpoolEventCollector::StopRegion() const {
+  int64_t now = GetCurrentTimeNanos();
+  TraceMeRecorder::Record(
+      {TraceMeEncode(kThreadpoolListenerStopRegion, {}), now, now});
+}
+
+absl::Status ThreadpoolProfilerInterface::Start() {
+  if (tracing::EventCollector::IsEnabled()) {
+    LOG(WARNING) << "[ThreadpoolEventCollector] EventCollector is enabled, Not "
+                    "collecting events from ThreadPool.";
+    status_ = absl::FailedPreconditionError(
+        "ThreadpoolEventCollector is enabled, Not collecting events from "
+        "ThreadPool.");
+    return absl::OkStatus();
+  }
+  event_collector_ = std::make_unique<ThreadpoolEventCollector>();
+  RegisterThreadpoolEventCollector(event_collector_.get());
+  threadpool_listener::Activate();
+  return absl::OkStatus();
+}
+
+absl::Status ThreadpoolProfilerInterface::Stop() {
+  threadpool_listener::Deactivate();
+  UnregisterThreadpoolEventCollector();
+  return absl::OkStatus();
+}
+
+absl::Status ThreadpoolProfilerInterface::CollectData(
+    tensorflow::profiler::XSpace* space) {
+  if (!status_.ok()) {
+    *space->add_errors() = status_.ToString();
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace profiler
+}  // namespace tsl

third_party/tsl/tsl/profiler/backends/cpu/threadpool_listener.h

@@ -0,0 +1,58 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_TSL_PROFILER_BACKENDS_CPU_THREADPOOL_LISTENER_H_
+#define TENSORFLOW_TSL_PROFILER_BACKENDS_CPU_THREADPOOL_LISTENER_H_
+
+#include "absl/status/status.h"
+#include "tsl/platform/tracing.h"
+#include "tsl/platform/types.h"
+#include "tsl/profiler/backends/cpu/threadpool_listener_state.h"
+#include "tsl/profiler/lib/profiler_interface.h"
+namespace tsl {
+namespace profiler {
+
+class ThreadpoolEventCollector : public tsl::tracing::EventCollector {
+ public:
+  explicit ThreadpoolEventCollector() = default;
+
+  void RecordEvent(uint64 arg) const override;
+  void StartRegion(uint64 arg) const override;
+  void StopRegion() const override;
+
+  // Annotates the current thread with a name.
+  void SetCurrentThreadName(const char* name) {}
+  // Returns whether event collection is enabled.
+  static bool IsEnabled() { return threadpool_listener::IsEnabled(); }
+};
+
+class ThreadpoolProfilerInterface : public ProfilerInterface {
+ public:
+  explicit ThreadpoolProfilerInterface() = default;
+
+  absl::Status Start() override;
+  absl::Status Stop() override;
+
+  absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
+
+ private:
+  absl::Status status_;
+  std::unique_ptr<ThreadpoolEventCollector> event_collector_;
+};
+
+}  // namespace profiler
+}  // namespace tsl
+
+#endif  // TENSORFLOW_TSL_PROFILER_BACKENDS_CPU_THREADPOOL_LISTENER_H_

third_party/tsl/tsl/profiler/backends/cpu/threadpool_listener_state.cc

@@ -0,0 +1,35 @@
+/* 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 "tsl/profiler/backends/cpu/threadpool_listener_state.h"
+
+#include <atomic>
+
+namespace tsl {
+namespace profiler {
+namespace threadpool_listener {
+namespace {
+std::atomic<bool> enabled;
+}
+
+bool IsEnabled() { return enabled.load(std::memory_order_acquire); }
+
+void Activate() { enabled.store(true, std::memory_order_release); }
+
+void Deactivate() { enabled.store(false, std::memory_order_release); }
+
+}  // namespace threadpool_listener
+}  // namespace profiler
+}  // namespace tsl

third_party/tsl/tsl/profiler/backends/cpu/threadpool_listener_state.h

@@ -0,0 +1,36 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_TSL_PROFILER_BACKENDS_CPU_THREADPOOL_LISTENER_STATE_H_
+#define TENSORFLOW_TSL_PROFILER_BACKENDS_CPU_THREADPOOL_LISTENER_STATE_H_
+
+namespace tsl {
+namespace profiler {
+namespace threadpool_listener {
+
+// Check if the threadpool listener is enabled.
+bool IsEnabled();
+
+// Set global state of threadpool listener to enabled.
+void Activate();
+
+// Set global state of threadpool listener to disabled.
+void Deactivate();
+
+}  // namespace threadpool_listener
+}  // namespace profiler
+}  // namespace tsl
+
+#endif  // TENSORFLOW_TSL_PROFILER_BACKENDS_CPU_THREADPOOL_LISTENER_STATE_H_

third_party/tsl/tsl/profiler/lib/BUILD

@@ -332,6 +332,7 @@ cc_library(
     visibility = internal_visibility([
         "@xla//xla/backends/profiler/plugin:__pkg__",
         "//learning/brain/tfrc/executor/stream_executor:__pkg__",
+        "@xla//xla/backends/profiler/cpu:__pkg__",
     ]),
     deps = [
         ":profiler_interface",

third_party/tsl/tsl/profiler/lib/context_types.cc

@@ -48,6 +48,8 @@ const char* GetContextTypeString(ContextType context_type) {
       return "pathways_exec";
     case ContextType::kPjrtLibraryCall:
       return "pjrt_library_call";
+    case ContextType::kThreadpoolEvent:
+      return "threadpool_event";
   }
 }
 

third_party/tsl/tsl/profiler/lib/context_types.h

@@ -37,6 +37,7 @@ enum class ContextType : int {
   kTpuLaunch,
   kPathwaysExecutor,
   kPjrtLibraryCall,
+  kThreadpoolEvent,
   kLastContextType = ContextType::kTpuLaunch,
 };
 

third_party/tsl/tsl/profiler/utils/BUILD

@@ -426,6 +426,7 @@ cc_library(
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",

third_party/tsl/tsl/profiler/utils/preprocess_xplane.cc

@@ -81,6 +81,7 @@ void MutateXPlane(XPlane* plane,
 std::vector<std::unique_ptr<XplaneEventMutatorFactory>>
 CreateMutatorFactories() {
   std::vector<std::unique_ptr<XplaneEventMutatorFactory>> mutator_factories;
+  mutator_factories.push_back(ThreadpoolLineMutatorFactory::CreateFactory());
   mutator_factories.push_back(XplaneRootEventMutatorFactory::CreateFactory(
       HostEventType::kProcessBatch, 2));
   mutator_factories.push_back(XplaneRootEventMutatorFactory::CreateFactory(