[ExecuTorch][WebGPU] Promote unary op plumbing to a shared 2D-dispatch helper (#21297)

Author: pytorchbot

backends/test/suite/flows/webgpu.py

@@ -16,19 +16,9 @@ def _create_webgpu_flow() -> TestFlow:
         skip_patterns=[
             "float16",
             "float64",  # Not supported in swiftshader
-            # WebGPU add is elementwise-only; broadcasting add.Tensor unsupported.
-            "bcast_first",
-            "bcast_second",
             "hardswish",
             "lstm_batch_sizes",
             "upsample_nearest2d",
-            # torchvision models with broadcasting adds; resnet50 covers wide.
-            "mobilenet_v3_small",
-            "shufflenet_v2_x1_0",
-            "resnet50",
-            "vit_b_16",
-            "swin_v2_t",
-            "convnext_small",
         ],
     )
 

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -164,18 +164,22 @@ Error WebGPUBackend::execute(
     return Error::Internal;
   }
 
-  // Execute the compute graph
-  graph->execute();
-
-  // Copy outputs from GPU staging buffers to EValue tensor data pointers
-  std::vector<std::pair<void*, size_t>> outputs;
-  outputs.reserve(num_outputs);
-  for (size_t i = 0; i < num_outputs; i++) {
-    const size_t arg_idx = num_inputs + i;
-    auto& tensor = args[arg_idx]->toTensor();
-    outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes());
+  // Execute + read back; fail loud as a runtime Error so a throw never crosses
+  // the backend boundary.
+  try {
+    graph->execute();
+    std::vector<std::pair<void*, size_t>> outputs;
+    outputs.reserve(num_outputs);
+    for (size_t i = 0; i < num_outputs; i++) {
+      const size_t arg_idx = num_inputs + i;
+      auto& tensor = args[arg_idx]->toTensor();
+      outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes());
+    }
+    graph->copy_outputs(outputs);
+  } catch (const std::exception& e) {
+    ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what());
+    return Error::Internal;
   }
-  graph->copy_outputs(outputs);
 
   return Error::Ok;
 }

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -508,6 +508,7 @@ void WebGPUGraph::build(
         }
         tensor.elem_size = vk_datatype_size(vk_tensor->datatype());
         tensor.is_int = vk_datatype_is_int(vk_tensor->datatype());
+        tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8;
         tensor.nbytes = numel * tensor.elem_size;
         // Live dims start == max (serialized upper bound); resize_input shrinks
         // them per call. Static graphs keep cur == max forever.
@@ -1902,9 +1903,12 @@ void WebGPUGraph::copy_outputs(std::vector<std::pair<void*, size_t>>& outputs) {
 
   std::vector<MapCallbackData> cb_data(count);
   std::vector<WGPUFuture> map_futures(count, WGPUFuture{});
+  // Map each output's LIVE staging size (an int64 output is int32-backed).
+  std::vector<size_t> map_nbytes(count, 0);
 
   for (size_t i = 0; i < count; i++) {
-    if (outputs[i].second == 0) {
+    map_nbytes[i] = tensors_[output_ids_[i]].cur_nbytes;
+    if (map_nbytes[i] == 0) {
       cb_data[i].status = WGPUMapAsyncStatus_Success;
       continue;
     }
@@ -1916,29 +1920,62 @@ void WebGPUGraph::copy_outputs(std::vector<std::pair<void*, size_t>>& outputs) {
         output_staging_buffers_[i],
         WGPUMapMode_Read,
         0,
-        outputs[i].second,
+        map_nbytes[i],
         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");
-    }
-  }
+  // Tracks which output buffers are currently mapped so a mid-loop throw can
+  // release them before propagating (no dangling mapped buffers).
+  std::vector<bool> is_mapped(count, false);
 
-  for (size_t i = 0; i < count; i++) {
-    if (outputs[i].second == 0) {
-      continue;
+  try {
+    for (size_t i = 0; i < count; i++) {
+      if (map_nbytes[i] == 0) {
+        continue;
+      }
+      if (webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) {
+        throw std::runtime_error("WebGPU: WaitAny failed for output map");
+      }
+      if (cb_data[i].status == WGPUMapAsyncStatus_Success) {
+        is_mapped[i] = true;
+      }
     }
-    if (cb_data[i].status == WGPUMapAsyncStatus_Success) {
+
+    for (size_t i = 0; i < count; i++) {
+      if (map_nbytes[i] == 0) {
+        continue;
+      }
+      if (cb_data[i].status != WGPUMapAsyncStatus_Success) {
+        throw std::runtime_error("WebGPU buffer map failed for output");
+      }
       const void* mapped = wgpuBufferGetConstMappedRange(
-          output_staging_buffers_[i], 0, outputs[i].second);
-      std::memcpy(outputs[i].first, mapped, outputs[i].second);
+          output_staging_buffers_[i], 0, map_nbytes[i]);
+      const size_t dst_nbytes = outputs[i].second;
+      if (dst_nbytes == map_nbytes[i]) {
+        std::memcpy(outputs[i].first, mapped, map_nbytes[i]);
+      } else if (
+          dst_nbytes == 2 * map_nbytes[i] && tensors_[output_ids_[i]].is_int &&
+          tensors_[output_ids_[i]].elem_size == 4) {
+        // int64 host output backed by an int32 GPU buffer: widen (sign-extend).
+        const int32_t* src = static_cast<const int32_t*>(mapped);
+        int64_t* dst = static_cast<int64_t*>(outputs[i].first);
+        const size_t n = map_nbytes[i] / sizeof(int32_t);
+        for (size_t k = 0; k < n; k++) {
+          dst[k] = static_cast<int64_t>(src[k]);
+        }
+      } else {
+        throw std::runtime_error("WebGPU: output buffer size mismatch");
+      }
       wgpuBufferUnmap(output_staging_buffers_[i]);
-    } else {
-      throw std::runtime_error("WebGPU buffer map failed for output");
+      is_mapped[i] = false;
+    }
+  } catch (...) {
+    for (size_t j = 0; j < count; j++) {
+      if (is_mapped[j]) {
+        wgpuBufferUnmap(output_staging_buffers_[j]);
+      }
     }
+    throw;
   }
 }
 

backends/webgpu/runtime/WebGPUGraph.h

@@ -34,6 +34,8 @@ struct WebGPUTensor {
   // Serialized (GPU-side) element type, used to narrow wider host inputs.
   size_t elem_size = 0;
   bool is_int = false;
+  // Exactly int8 (not uint8/bool), so int8-only ops can guard their dtype.
+  bool is_int8 = false;
 };
 
 // Host-side view of one graph input, passed to copy_inputs.

backends/webgpu/runtime/WebGPUUtils.h

@@ -457,6 +457,63 @@ inline ComputePipelineBundle make_compute_pipeline(
   return bundle;
 }
 
+// Builds another pipeline + bind group from resources owned by an earlier
+// bundle. The binding indices and types must match the shared bind-group
+// layout. This preserves shared-shader/layout multi-pipeline construction
+// without transferring ownership of those resources to the returned bundle.
+inline ComputePipelineBundle make_compute_pipeline(
+    WGPUDevice device,
+    const ComputePipelineBundle& shared_resources,
+    const std::vector<BindingSpec>& bindings,
+    const WGPUConstantEntry* constants = nullptr,
+    size_t constant_count = 0,
+    const char* entry_point = "main") {
+  if (shared_resources.shader == nullptr ||
+      shared_resources.bind_group_layout == nullptr ||
+      shared_resources.pipeline_layout == nullptr) {
+    throw std::runtime_error(
+        "make_compute_pipeline: shared resources are not available");
+  }
+
+  ComputePipelineBundle bundle;
+
+  std::vector<WGPUBindGroupEntry> bind_entries(bindings.size());
+  for (size_t i = 0; i < bindings.size(); i++) {
+    bind_entries[i] = {};
+    bind_entries[i].binding = bindings[i].binding;
+    bind_entries[i].buffer = bindings[i].buffer;
+    bind_entries[i].size = bindings[i].size;
+  }
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = shared_resources.pipeline_layout;
+  pipeline_desc.compute.module = shared_resources.shader;
+  pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = constant_count;
+  pipeline_desc.compute.constants = constants;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+  if (pipeline == nullptr) {
+    throw std::runtime_error(
+        "make_compute_pipeline: compute pipeline creation failed");
+  }
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = shared_resources.bind_group_layout;
+  bg_desc.entryCount = bind_entries.size();
+  bg_desc.entries = bind_entries.data();
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+  if (bind_group == nullptr) {
+    wgpuComputePipelineRelease(pipeline);
+    throw std::runtime_error(
+        "make_compute_pipeline: bind group creation failed");
+  }
+
+  bundle.pipeline = pipeline;
+  bundle.bind_group = bind_group;
+  return bundle;
+}
+
 // The {wg_size, stride_x} override-constant pair every 2D-spill dispatch
 // builds from its DispatchGrid; was hand-rolled identically at 7 call sites.
 inline std::array<WGPUConstantEntry, 2> make_grid_constants(

backends/webgpu/runtime/ops/TensorMeta.h

@@ -16,9 +16,9 @@
 
 namespace executorch::backends::webgpu {
 
-constexpr uint32_t kTensorMetaMaxNdim = 4;
+constexpr uint32_t kTensorMetaMaxNdim = 8;
 
-// Per-tensor metadata UBO; mirrors Vulkan BufferMetadata (4-dim NCHW, std140).
+// Per-tensor metadata UBO; mirrors Vulkan BufferMetadata (8-dim NCHW, std140).
 struct TensorMeta {
   uint32_t ndim;
   uint32_t numel;
@@ -28,19 +28,19 @@ struct TensorMeta {
 };
 
 static_assert(
-    sizeof(TensorMeta) == 48,
-    "TensorMeta std140 layout must be 48 bytes to match the WGSL uniform");
+    sizeof(TensorMeta) == 80,
+    "TensorMeta std140 layout must be 80 bytes to match the WGSL uniform");
 // Lock the std140 field offsets the WGSL uniform reads, not just total size.
 static_assert(offsetof(TensorMeta, ndim) == 0);
 static_assert(offsetof(TensorMeta, numel) == 4);
 static_assert(offsetof(TensorMeta, sizes) == 16);
-static_assert(offsetof(TensorMeta, strides) == 32);
+static_assert(offsetof(TensorMeta, strides) == 48);
 
 // Fill TensorMeta from NCHW dims: contiguous strides, padded trailing slots.
 inline void fill_tensor_meta(const WebGPUTensor& t, TensorMeta* m) {
   const uint32_t ndim = static_cast<uint32_t>(t.dims.size());
   if (ndim > kTensorMetaMaxNdim) {
-    throw std::runtime_error("TensorMeta: tensor rank exceeds 4 (MAX_NDIM)");
+    throw std::runtime_error("TensorMeta: tensor rank exceeds 8 (MAX_NDIM)");
   }
   *m = {};
   for (uint32_t d = 0; d < kTensorMetaMaxNdim; d++) {
@@ -67,7 +67,7 @@ inline void fill_tensor_meta_broadcast(
     TensorMeta* m) {
   const uint32_t rank = static_cast<uint32_t>(t.dims.size());
   if (out_ndim > kTensorMetaMaxNdim) {
-    throw std::runtime_error("TensorMeta: out_ndim exceeds 4 (MAX_NDIM)");
+    throw std::runtime_error("TensorMeta: out_ndim exceeds 8 (MAX_NDIM)");
   }
   if (rank > out_ndim) {
     throw std::runtime_error("TensorMeta: operand rank exceeds out_ndim");

backends/webgpu/runtime/ops/adamw/AdamwStep.cpp

@@ -97,78 +97,26 @@ void adamw_step_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       utils::make_uniform(device, &params, sizeof(params));
   graph.add_uniform_buffer_bytes(sizeof(params));
 
-  WGPUShaderSourceWGSL wgsl_desc = {};
-  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
-  wgsl_desc.code = {kAdamwStepWGSL, WGPU_STRLEN};
-  WGPUShaderModuleDescriptor shader_desc = {};
-  shader_desc.nextInChain = &wgsl_desc.chain;
-  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
-
-  WGPUBindGroupLayoutEntry entries[5] = {};
-  for (uint32_t i = 0; i <= 2; i++) {
-    entries[i].binding = i;
-    entries[i].visibility = WGPUShaderStage_Compute;
-    entries[i].buffer.type = WGPUBufferBindingType_Storage;
-  }
-  entries[3].binding = 3;
-  entries[3].visibility = WGPUShaderStage_Compute;
-  entries[3].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
-  entries[4].binding = 4;
-  entries[4].visibility = WGPUShaderStage_Compute;
-  entries[4].buffer.type = WGPUBufferBindingType_Uniform;
-
-  WGPUBindGroupLayoutDescriptor bgl_desc = {};
-  bgl_desc.entryCount = 5;
-  bgl_desc.entries = entries;
-  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
-
-  WGPUPipelineLayoutDescriptor pl_desc = {};
-  pl_desc.bindGroupLayoutCount = 1;
-  pl_desc.bindGroupLayouts = &bgl;
-  WGPUPipelineLayout pipeline_layout =
-      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
-
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
 
-  WGPUComputePipelineDescriptor pipeline_desc = {};
-  pipeline_desc.layout = pipeline_layout;
-  pipeline_desc.compute.module = shader;
-  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
-  pipeline_desc.compute.constantCount = 1;
-  pipeline_desc.compute.constants = &wg_size_constant;
-  WGPUComputePipeline pipeline =
-      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
-
-  WGPUBindGroupEntry bg_entries[5] = {};
-  bg_entries[0].binding = 0;
-  bg_entries[0].buffer = param.buffer;
-  bg_entries[0].size = param.nbytes;
-  bg_entries[1].binding = 1;
-  bg_entries[1].buffer = m.buffer;
-  bg_entries[1].size = m.nbytes;
-  bg_entries[2].binding = 2;
-  bg_entries[2].buffer = v.buffer;
-  bg_entries[2].size = v.nbytes;
-  bg_entries[3].binding = 3;
-  bg_entries[3].buffer = grad.buffer;
-  bg_entries[3].size = grad.nbytes;
-  bg_entries[4].binding = 4;
-  bg_entries[4].buffer = uniform_buffer;
-  bg_entries[4].size = sizeof(params);
-
-  WGPUBindGroupDescriptor bg_desc = {};
-  bg_desc.layout = bgl;
-  bg_desc.entryCount = 5;
-  bg_desc.entries = bg_entries;
-  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
-
-  graph.add_dispatch({pipeline, bind_group, workgroup_count, "adamw_step"});
-
-  wgpuShaderModuleRelease(shader);
-  wgpuBindGroupLayoutRelease(bgl);
-  wgpuPipelineLayoutRelease(pipeline_layout);
+  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+      device,
+      kAdamwStepWGSL,
+      {
+          {0, WGPUBufferBindingType_Storage, param.buffer, param.nbytes},
+          {1, WGPUBufferBindingType_Storage, m.buffer, m.nbytes},
+          {2, WGPUBufferBindingType_Storage, v.buffer, v.nbytes},
+          {3, WGPUBufferBindingType_ReadOnlyStorage, grad.buffer, grad.nbytes},
+          {4, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(params)},
+      },
+      &wg_size_constant,
+      1);
+
+  graph.add_dispatch(
+      {bundle.pipeline, bundle.bind_group, workgroup_count, "adamw_step"});
+
   graph.own_uniform_buffer(uniform_buffer);
 }