PR #38110: [xla:gpu] Improve host tracing debuggability

Imported from GitHub PR #38110

Sprinkle `TraceMe` annotations with `XLA_LOG_DEVICE`-compatible format in `CommonPjRtClient` and in XLA:GPU-specific `StreamExecutor` client. This improves debuggability of host runtime by including important metadata into each trace.
Copybara import of the project:

--
00bda25 by Eugene Zhulenev <ezhulenev@openxla.org>:

[xla:gpu] Improve host tracing debuggability

Merging this change closes #38110

FUTURE_COPYBARA_INTEGRATE_REVIEW=#38110 from ezhulenev:host-tracing-0 00bda25
PiperOrigin-RevId: 874225232

Author: ezhulenev

xla/pjrt/common_pjrt_client.cc

@@ -19,6 +19,7 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <functional>
+#include <iterator>
 #include <memory>
 #include <optional>
 #include <string>
@@ -1108,8 +1109,14 @@ CommonPjRtLoadedExecutable::Execute(
     absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
     const ExecuteOptions& options,
     std::optional<std::vector<tsl::Future<void>>>& returned_futures) const {
-  tsl::profiler::TraceMe traceme("CommonPjRtLoadedExecutable::Execute");
-  VLOG(1) << "CommonPjRtLoadedExecutable::Execute";
+  RunId run_id = options.launch_id != 0 ? RunId(options.launch_id)
+                                        : RunId::CreateUniqueId();
+  int num_addressable_devices = addressable_devices_.size();
+
+  VLOG(1) << absl::StreamFormat(
+      "CommonPjRtLoadedExecutable::Execute: run_id=%d, execution_mode=%v",
+      run_id.ToInt(), options.execution_mode);
+
   if (!client()->allows_execute_recursion() &&
       ThisThreadIsInsideHostCallback()) {
     // Because TPU is single threaded, and the host callback currently blocking
@@ -1118,13 +1125,18 @@ CommonPjRtLoadedExecutable::Execute(
     return InvalidArgument("Execute() called from inside host callback.");
   }
 
-  RunId run_id = options.launch_id != 0 ? RunId(options.launch_id)
-                                        : RunId::CreateUniqueId();
-  tsl::profiler::TraceMeProducer producer("CommonPjRtLoadedExecutable::Execute",
-                                          tsl::profiler::ContextType::kPjRt,
-                                          run_id.ToInt());
-
-  const int num_addressable_devices = addressable_devices_.size();
+  tsl::profiler::TraceMeProducer producer(
+      [&] {
+        return tsl::profiler::TraceMeEncode(
+            absl::StrFormat("CommonPjRtLoadedExecutable::Execute (%s)", name()),
+            {{"run_id", run_id.ToInt()},
+             {"execution_mode", absl::StrCat(options.execution_mode)},
+             {"name", name()},
+             {"num_replicas", num_replicas()},
+             {"num_partitions", num_partitions()},
+             {"num_addressable_devices", num_addressable_devices}});
+      },
+      tsl::profiler::ContextType::kPjRt, run_id.ToInt());
 
   if (argument_handles.size() != num_addressable_devices) {
     return InvalidArgument(
@@ -1170,9 +1182,18 @@ CommonPjRtLoadedExecutable::Execute(
         const int replica = addressable_device_logical_ids_[i].replica;
         const int partition = addressable_device_logical_ids_[i].partition;
         PjRtDevice* device = addressable_devices_[i];
-        LaunchOnDevice(device, [&, replica, partition, i, context_id] {
+        LaunchOnDevice(device, [&, context_id, i, replica, partition, device] {
           tsl::profiler::TraceMeConsumer consumer(
-              "Scheduled CommonPjRtLoadedExecutable::Execute",
+              [&] {
+                return tsl::profiler::TraceMeEncode(
+                    absl::StrFormat(
+                        "[%d] CommonPjRtLoadedExecutable::Execute (%s)", i,
+                        name()),
+                    {{"name", name()},
+                     {"replica", replica},
+                     {"partition", partition},
+                     {"global_device_id", device->global_device_id()}});
+              },
               tsl::profiler::ContextType::kPjRt, context_id);
 
           // Two phase launch. Phase 1: Prepare on all cores. Abort
@@ -1216,6 +1237,7 @@ CommonPjRtLoadedExecutable::Execute(
     }
 
     // Wait until we either fail Phase 1 or completes two phases.
+    tsl::profiler::TraceMe trace_wait("Wait for LaunchOnDevice completion");
     auto done = [&]() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu) {
       return launching == 0;
     };

xla/pjrt/gpu/se_gpu_pjrt_client.cc

@@ -549,6 +549,12 @@ absl::StatusOr<PreparedSend> PrepareSend(
   GlobalDeviceId src_device(buffer->device()->global_device_id().value());
   GlobalDeviceId dst_device(dst_global_device_id.value());
 
+  tsl::profiler::TraceMe trace([&] {
+    return tsl::profiler::TraceMeEncode(
+        absl::StrFormat("PrepareSend: src=%v dst=%v", src_device, dst_device),
+        {{"transfer_key", transfer_key}});
+  });
+
   // Form the GPU clique key.
   // TODO(asrao, mwhittaker): Supply correct incarnations when creating the
   // clique key.
@@ -601,6 +607,13 @@ absl::StatusOr<PreparedReceive> PrepareReceive(
   GlobalDeviceId src_device(src_global_device_id.value());
   GlobalDeviceId dst_device(device->global_device_id().value());
 
+  tsl::profiler::TraceMe trace([&] {
+    return tsl::profiler::TraceMeEncode(
+        absl::StrFormat("PrepareReceive: src=%v dst=%v", src_device,
+                        dst_device),
+        {{"transfer_key", transfer_key}});
+  });
+
   // Form the GPU clique key.
   // TODO(asrao, mwhittaker): Supply correct incarnations when creating the
   // clique key.
@@ -776,6 +789,14 @@ void StreamExecutorGpuClient::ScheduleSendsOnLocalDevice(
   prepared_sends.reserve(buffers.size());
   tsl::RCReference<PjRtStreamExecutorDeviceEvent> usage_event;
 
+  tsl::profiler::TraceMe trace([&] {
+    return tsl::profiler::TraceMeEncode(
+        absl::StrFormat(
+            "[%v] StreamExecutorGpuClient::ScheduleSendsOnLocalDevice",
+            device->local_device_id()),
+        {{"num_buffers", buffers.size()}});
+  });
+
   auto setup_sends = [&]() -> absl::Status {
     TF_ASSIGN_OR_RETURN(local_device_state, GetLocalDeviceState(device));
     stream = local_device_state->GetDeviceToDeviceStream();
@@ -853,6 +874,10 @@ void StreamExecutorGpuClient::ScheduleSendsOnLocalDevice(
     group_futures.reserve(grouped_sends.size());
 
     for (auto& [clique_key, curr_sends] : grouped_sends) {
+      tsl::profiler::TraceMe trace([&k = clique_key] {
+        return tsl::profiler::TraceMeEncode("LaunchSend", {{"clique", k}});
+      });
+
       // Get the communicator on which we will execute this group of
       // transfers. We assume each clique key is associated with a unique
       // communicator, so we just take the communicator of the first
@@ -994,6 +1019,13 @@ StreamExecutorGpuClient::CrossHostReceiveBuffers(
   prepared_receives.reserve(shapes.size());
   tsl::RCReference<PjRtStreamExecutorDeviceEvent> definition_event;
 
+  tsl::profiler::TraceMe trace([&] {
+    return tsl::profiler::TraceMeEncode(
+        absl::StrFormat("[%v] StreamExecutorGpuClient::CrossHostReceiveBuffers",
+                        device->local_device_id()),
+        {{"num_shapes", shapes.size()}});
+  });
+
   auto setup_receives = [&]() -> absl::Status {
     TF_ASSIGN_OR_RETURN(local_device_state, GetLocalDeviceState(device));
     stream = local_device_state->GetDeviceToDeviceStream();
@@ -1064,6 +1096,10 @@ StreamExecutorGpuClient::CrossHostReceiveBuffers(
         group_futures.reserve(grouped_receives.size());
 
         for (auto& [clique_key, curr_receives] : grouped_receives) {
+          tsl::profiler::TraceMe trace([&k = clique_key] {
+            return tsl::profiler::TraceMeEncode("LaunchRecv", {{"clique", k}});
+          });
+
           // Get the communicator on which we will execute this group of
           // transfers. We assume each clique key is associated with a unique
           // communicator, so we just take the communicator of the first

xla/pjrt/pjrt_executable.h

@@ -301,6 +301,21 @@ struct ExecuteOptions {
   absl::StatusOr<ExecuteOptionsProto> ToProto() const;
   static absl::StatusOr<ExecuteOptions> FromProto(
       const ExecuteOptionsProto& proto);
+
+  // Pretty-printing for ExecutionMode enum.
+  template <typename Sink>
+  friend void AbslStringify(Sink& sink, const ExecutionMode& mode) {
+    absl::Format(&sink, "%s", [&] {
+      switch (mode) {
+        case ExecutionMode::kDefault:
+          return "default";
+        case ExecutionMode::kSynchronous:
+          return "synchronous";
+        case ExecutionMode::kAsynchronous:
+          return "asynchronous";
+      }
+    }());
+  }
 };
 
 // Static memory usage for a compiled program.

xla/pjrt/pjrt_stream_executor_client.cc

@@ -1685,9 +1685,14 @@ PjRtStreamExecutorRawLoadedExecutable::Execute(
                            ->local_device_id()
                            .value();
   LocalDeviceState* device_state = &(client_->device_state(device_ordinal));
-  tsl::profiler::TraceMeConsumer activity(
-      "PjRtStreamExecutorLoadedExecutable::EnqueueExecution",
-      tsl::profiler::ContextType::kPjRt, run_id_.ToInt());
+
+  tsl::profiler::TraceMe trace([&] {
+    return tsl::profiler::TraceMeEncode(
+        absl::StrFormat("[%d] PjRtStreamExecutorRawLoadedExecutable::Execute",
+                        device_ordinal),
+        {{"replica", replica_}, {"partition", partition_}});
+  });
+
   VLOG(3) << "Replica " << replica_ << ", partition " << partition_
           << " mapped to device ordinal for execution: " << device_ordinal;
 
@@ -1762,9 +1767,14 @@ PjRtStreamExecutorRawLoadedExecutable::Execute(
     // launch is delayed.
     std::shared_ptr<Semaphore::ScopedReservation> compute_reservation;
     {
-      tsl::profiler::TraceMe traceme("ComputeSemaphoreAcquire");
+      Semaphore& compute_semaphore = device_state->compute_semaphore();
+      tsl::profiler::TraceMe traceme([&] {
+        return absl::StrFormat(
+            "ComputeSemaphoreAcquire [capacity=%d, value=%d]",
+            compute_semaphore.capacity(), compute_semaphore.value());
+      });
       compute_reservation = std::make_shared<Semaphore::ScopedReservation>(
-          device_state->compute_semaphore().ScopedAcquire(1));
+          compute_semaphore.ScopedAcquire(1));
     }
 
     absl::Status predetermined_error;

xla/pjrt/semaphore.h

@@ -41,6 +41,12 @@ class Semaphore {
   // Returns the capacity of the semaphore.
   int64_t capacity() const { return max_capacity_; }
 
+  // Returns the current value of the semaphore.
+  int64_t value() const {
+    absl::MutexLock lock(mu_);
+    return value_;
+  }
+
   class ScopedReservation {
    public:
     ScopedReservation(Semaphore* semaphore, int64_t amount)
@@ -69,7 +75,7 @@ class Semaphore {
   static bool CanAcquire(CanAcquireArgs* args)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(args->semaphore->mu_);
 
-  absl::Mutex mu_;
+  mutable absl::Mutex mu_;
   int64_t value_ ABSL_GUARDED_BY(mu_);
   const int64_t max_capacity_;
 };

xla/service/gpu/gpu_executable.cc

@@ -1200,8 +1200,10 @@ absl::Status GpuExecutable::ExecuteThunks(
     const BufferAllocations& buffer_allocations,
     const ServiceExecutableRunOptions* run_options) {
   tsl::profiler::TraceMe trace([&] {
-    return tsl::profiler::TraceMeEncode("GpuExecutable::ExecuteThunks",
-                                        {{"module_name", module_name_}});
+    return tsl::profiler::TraceMeEncode(
+        absl::StrFormat("[%d] GpuExecutable::ExecuteThunks",
+                        run_options->device_ordinal()),
+        {{"module_name", module_name_}});
   });
 
   if (VLOG_IS_ON(5)) {