[ExecuTorch][WebGPU] Broadcast add.Tensor (mirror mul) + un-skip 6 models

Pull Request resolved: #21233

`add.Tensor` was elementwise-only: `binary_add.wgsl` computed `output[idx] = input1[idx] + alpha * input2[idx]` over a single flat index up to the output numel, so on a broadcast it read the smaller operand out of bounds (WebGPU robustness clamps/zeros -> silently wrong, not `0x30`). This is the failure mode behind the models the WebGPU flow skips (`resnet50`, `vit_b_16`, `swin_v2_t`, `convnext_small`, `mobilenet_v3_small`, `shufflenet_v2_x1_0`) and the `bcast_first`/`bcast_second` op cases.

Give `add.Tensor` NumPy broadcasting by mirroring the shipped `mul`:
- `binary_add.wgsl` - copy `binary_mul`'s kernel (rank-8 `TensorMeta` layout from the rank-cap fix below it in the stack): keep the identical-shape elementwise fast path (common case stays bit-for-bit `input1[idx] + alpha * input2[idx]`), else relinearize out idx -> per-input coords (clamp size-1 dims) and add. `alpha` is a second pipeline-override constant (read once at build, never rewritten on resize), so no extra UBO.
- `add/BinaryOp.cpp` - port `mul_impl`: 3 `TensorMeta` UBOs via `fill_tensor_meta_broadcast`, rank + fp32 guards, 6-entry bind group, `constantCount = 2` ({wg_size, alpha}), 2D dispatch, and `mul`'s resize hook verbatim (rebuild the 3 metas from `cur_dims`, `set_cur_dims`, rewrite UBOs + dispatch); `own_uniform_buffer` the 3 metas. Drops the old flat `AddParams` path.
- `flows/webgpu.py` - un-skip the 6 broadcast-add models + the `bcast_first`/`bcast_second` op cases (kept the `float16`/`float64` dtype skips and `hardswish`/`lstm_batch_sizes`/`upsample_nearest2d`).

This sits above the rank-cap fix, so the shader uses the rank-8 `array<vec4<u32>, 2>` `TensorMeta` layout. Applied identically to both `xplat/` and `fbcode/` mirrors (byte-identical).

Co-authored-with: Claude Code.
ghstack-source-id: 406388773
@exported-using-ghexport

Differential Revision: [D113319599](https://our.internmc.facebook.com/intern/diff/D113319599/)

Author: JCNTH

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/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/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);
 }