Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/runtime/WebGPUCompat.h

@@ -12,6 +12,13 @@
 
 #include <cstdint>
 
+#if defined(__EMSCRIPTEN__)
+#include <emscripten/emscripten.h>
+#else
+#include <chrono>
+#include <thread>
+#endif
+
 namespace executorch::backends::webgpu {
 
 // Caller's instance must enable TimedWaitAny; returns the WaitAny status.
@@ -21,4 +28,15 @@ inline WGPUWaitStatus webgpu_wait(WGPUInstance instance, WGPUFuture future) {
   return wgpuInstanceWaitAny(instance, 1, &info, UINT64_MAX);
 }
 
+// Blocking event-loop pump for the output-map readback (browser: JS yield).
+inline void webgpu_poll(WGPUInstance instance) {
+#if defined(__EMSCRIPTEN__)
+  (void)instance;
+  emscripten_sleep(0);
+#else
+  wgpuInstanceProcessEvents(instance);
+  std::this_thread::sleep_for(std::chrono::microseconds(50));
+#endif
+}
+
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -580,6 +580,7 @@ void WebGPUGraph::execute() {
 namespace {
 
 struct MapCallbackData {
+  bool done = false;
   WGPUMapAsyncStatus status = WGPUMapAsyncStatus_Error;
 };
 
@@ -590,6 +591,7 @@ void buffer_map_callback(
     void* /*userdata2*/) {
   auto* data = static_cast<MapCallbackData*>(userdata1);
   data->status = status;
+  data->done = true;
 }
 
 } // namespace
@@ -598,29 +600,35 @@ void WebGPUGraph::copy_outputs(std::vector<std::pair<void*, size_t>>& outputs) {
   const size_t count = std::min(outputs.size(), output_staging_buffers_.size());
 
   std::vector<MapCallbackData> cb_data(count);
-  std::vector<WGPUFuture> map_futures(count, WGPUFuture{});
 
   for (size_t i = 0; i < count; i++) {
     if (outputs[i].second == 0) {
+      cb_data[i].done = true;
       cb_data[i].status = WGPUMapAsyncStatus_Success;
       continue;
     }
     WGPUBufferMapCallbackInfo cb_info = {};
-    cb_info.mode = WGPUCallbackMode_WaitAnyOnly;
+    cb_info.mode = WGPUCallbackMode_AllowSpontaneous;
     cb_info.callback = buffer_map_callback;
     cb_info.userdata1 = &cb_data[i];
-    map_futures[i] = wgpuBufferMapAsync(
+    wgpuBufferMapAsync(
         output_staging_buffers_[i],
         WGPUMapMode_Read,
         0,
         outputs[i].second,
         cb_info);
   }
 
-  for (size_t i = 0; i < count; i++) {
-    if (outputs[i].second != 0 &&
-        webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) {
-      throw std::runtime_error("WebGPU: WaitAny failed for output map");
+  // WaitAny returns pre-queue on CI Dawn+SwiftShader here; pump until mapped.
+  bool all_mapped = false;
+  while (!all_mapped) {
+    webgpu_poll(instance_);
+    all_mapped = true;
+    for (size_t i = 0; i < count; i++) {
+      if (outputs[i].second != 0 && !cb_data[i].done) {
+        all_mapped = false;
+        break;
+      }
     }
   }