[XLA:GPU] Add debug messages to track which thunk hangs.

PiperOrigin-RevId: 736468039

Author: Google-ML-Automation

xla/backends/gpu/runtime/BUILD

@@ -1100,10 +1100,11 @@ cc_library(
     deps = [
         ":annotation",
         ":thunk",
+        "//xla/tsl/platform:errors",
         "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
-        "@tsl//tsl/platform:errors",
         "@tsl//tsl/profiler/lib:scoped_annotation",
     ],
 )

xla/backends/gpu/runtime/sequential_thunk.cc

@@ -21,11 +21,12 @@ limitations under the License.
 #include <utility>
 
 #include "absl/functional/function_ref.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "xla/backends/gpu/runtime/annotation.h"
 #include "xla/backends/gpu/runtime/thunk.h"
-#include "tsl/platform/errors.h"
+#include "xla/tsl/platform/errors.h"
 #include "tsl/profiler/lib/scoped_annotation.h"
 
 namespace xla {
@@ -83,7 +84,13 @@ absl::Status SequentialThunk::ExecuteOnStream(const ExecuteParams& params) {
     if (params.mock_collectives && thunk->IsCollective()) {
       continue;
     }
+    VLOG(1) << "[" << params.stream->parent()->device_ordinal() << "] "
+            << "Start SequentialThunk::ExecuteOnStream: "
+            << thunk->profile_annotation();
     TF_RETURN_IF_ERROR(thunk->ExecuteOnStream(params));
+    VLOG(1) << "[" << params.stream->parent()->device_ordinal() << "] "
+            << "End SequentialThunk::ExecuteOnStream: "
+            << thunk->profile_annotation();
   }
   return absl::OkStatus();
 }

xla/service/gpu/gpu_executable.cc

@@ -345,7 +345,11 @@ absl::Status ExecuteThunksImpl(
       command_buffer_trace_stream, &collective_params, &collective_cliques,
       std::move(additional_execution_streams));
 
+  VLOG(1) << "[" << run_options->device_ordinal() << "] "
+          << "Start GpuExecutable::ExecuteOnStream module: " << module_name;
   TF_RETURN_IF_ERROR(thunk_sequence.ExecuteOnStream(execute_params));
+  VLOG(1) << "[" << run_options->device_ordinal() << "] "
+          << "End GpuExecutable::ExecuteOnStream module: " << module_name;
 
   return MaybeSyncAndProfile(run_options, execution_timer.get(),
                              block_host_until_done ? main_stream : nullptr);