Merge pull request #17 from ROCm/rocm-jaxlib-v0.4.28-qa-fix-async-wait

Let the other stream wait for the main stream before issuing memcpy d2h

xla/service/gpu/runtime/copy_thunk.cc

@@ -126,12 +126,13 @@ absl::Status DeviceToHostCopyThunk::ExecuteOnStream(
   TF_ASSIGN_OR_RETURN(
       se::Stream * stream,
       GetStreamForExecution(Thunk::execution_stream_id(), params));
+  TF_RETURN_IF_ERROR(stream->WaitFor(params.stream));
   TF_RETURN_IF_ERROR(stream->Memcpy(cpu_dst, source_data, size_bytes()));
   if (stream == params.stream) {
     VLOG(2) << "Memcpy D2H from the main stream";
     return absl::OkStatus();
   }
-  VLOG(2) << "Memcpy D2H from the other stream";
+  VLOG(2) << "Memcpy D2H from the stream " << Thunk::execution_stream_id();
   se::StreamExecutor* executor = params.stream->parent();
   se::Event event(executor);
   if (!event.Init()) {