add stream sync

openxla · Aug 5, 2024 · 37c3bb7 · 37c3bb7
1 parent 6db4cec
commit 37c3bb7
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/xla/service/gpu/runtime/nccl_all_to_all_thunk.cc b/xla/service/gpu/runtime/nccl_all_to_all_thunk.cc
@@ -126,7 +126,9 @@ absl::Status RunAllToAll(NcclApi* nccl_api, bool has_split_dimension,
 
   TF_ASSIGN_OR_RETURN(int32_t num_participants, nccl_api->CommCount(comm));
 
-  if (!use_memcpy) {
+  if (use_memcpy) {
+    TF_RETURN_IF_ERROR(stream.BlockHostUntilDone());
+  } else {
     TF_RETURN_IF_ERROR(nccl_api->GroupStart());
   }