PR #22956: vlog device id in while_thunk.

Imported from GitHub PR #22956 Copybara import of the project: -- d462315 by Yunlong Liu <[email protected]>: vlog device id in while_thunk. Merging this change closes #22956 COPYBARA_INTEGRATE_REVIEW=#22956 from yliu120:patch-1 d462315 PiperOrigin-RevId: 730681273
openxla · Feb 25, 2025 · d2aee3b · d2aee3b
1 parent 8d3bdff
commit d2aee3b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/xla/backends/gpu/runtime/while_thunk.cc b/xla/backends/gpu/runtime/while_thunk.cc
@@ -108,7 +108,8 @@ absl::Status WhileThunk::ExecuteOnStream(const ExecuteParams& params) {
   if (trip_count_.has_value()) {
     VLOG(2) << "Executing WhileThunk for " << *trip_count_ << " iterations";
     for (iter = 0; iter < trip_count_; ++iter) {
-      VLOG(3) << "Executing iteration # " << iter;
+      VLOG(3) << "Executing iteration # " << iter
+              << " (Device: " << stream.parent()->device_ordinal() << ")";
       TF_RETURN_IF_ERROR(body_thunk_sequence_->ExecuteOnStream(params));
     }
     return absl::OkStatus();
@@ -142,7 +143,8 @@ absl::Status WhileThunk::ExecuteOnStream(const ExecuteParams& params) {
       break;
     }
 
-    VLOG(3) << "Executing WhileThunk body computation; iter=" << iter;
+    VLOG(3) << "Executing WhileThunk body computation; iter=" << iter
+            << " (Device: " << stream.parent()->device_ordinal() << ")";
     TF_RETURN_IF_ERROR(body_thunk_sequence_->ExecuteOnStream(params));
     ++iter;
   }