From 94fe20bf68cee874695396461a095ddda9d44a1d Mon Sep 17 00:00:00 2001 From: Eugene Zhulenev Date: Fri, 3 Apr 2026 18:19:46 +0000 Subject: [PATCH] [xla:gpu] TraceMe for while loop iterations --- xla/backends/gpu/runtime/while_thunk.cc | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/xla/backends/gpu/runtime/while_thunk.cc b/xla/backends/gpu/runtime/while_thunk.cc index 75febd8cd7f86..bb2df309a22e8 100644 --- a/xla/backends/gpu/runtime/while_thunk.cc +++ b/xla/backends/gpu/runtime/while_thunk.cc @@ -36,10 +36,10 @@ limitations under the License. #include "xla/service/buffer_assignment.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/stream.h" +#include "xla/tsl/platform/status_macros.h" #include "xla/util.h" #include "xla/xla_data.pb.h" #include "tsl/profiler/lib/traceme.h" -#include "xla/tsl/platform/status_macros.h" namespace xla::gpu { @@ -86,6 +86,8 @@ absl::Status WhileThunk::ExecuteOnStream(const ExecuteParams& params) { XLA_VLOG_DEVICE(2, device_ordinal) << "Executing WhileThunk for " << *trip_count_ << " iterations"; for (size_t i = 0; i < trip_count_; loop.IncLoopIteration(), ++i) { + TraceMe trace( + [&] { return absl::StrFormat("[iter=%d]", loop.loop_iteration()); }); XLA_VLOG_DEVICE(3, device_ordinal) << "Executing iteration # " << i << " (Device: " << stream.parent()->device_ordinal() << ")"; @@ -106,9 +108,8 @@ absl::Status WhileThunk::ExecuteOnStream(const ExecuteParams& params) { condition_result_buffer_index_); for (;; loop.IncLoopIteration()) { - TraceMe trace([&] { - return TraceMeEncode("While", {{"iteration:", loop.loop_iteration()}}); - }); + TraceMe trace( + [&] { return absl::StrFormat("[iter=%d]", loop.loop_iteration()); }); XLA_VLOG_DEVICE(3, device_ordinal) << "Executing WhileThunk condition computation; iter=" @@ -120,9 +121,9 @@ absl::Status WhileThunk::ExecuteOnStream(const ExecuteParams& params) { stream.Memcpy(condition_result, condition_result_data, sizeof(bool))); if (absl::Status blocked = stream.BlockHostUntilDone(); !blocked.ok()) { - return absl::InternalError(absl::StrFormat( + return Internal( "Failed to complete all kernels launched on stream %p: %s", &stream, - blocked.message())); + blocked.message()); } XLA_VLOG_DEVICE(3, device_ordinal)