From 84c67df666d669816f347ac41b5f02ddbc3d65ab Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:21 -0700 Subject: [PATCH 01/22] [ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20863 **Adds `relu` to the WebGPU backend via a shared elementwise-unary handler.** ReLU is on the SAM2/SAM3 mask-decoder MLP path, so it is needed to delegate those decoders. **Problem** — The backend had no `aten.relu.default` kernel, and `sigmoid` (the only prior unary op) built its compute pipeline inline rather than through the shared helper — duplicating the bind-group/dispatch boilerplate that a second unary op would repeat. **Solution** - Before: `sigmoid` was implemented with a bespoke inline pipeline; there was no `relu`. - After: a single generic `add_unary_op` helper (in `runtime/ops/sigmoid/UnaryOp.cpp`) builds the input/output/params binding and dispatch for any elementwise-unary WGSL; `sigmoid_impl` and the new `relu_impl` are thin wrappers over it, so `sigmoid` now goes through the same `utils::make_compute_pipeline` path as `relu`. `relu.wgsl` is a one-element-per-thread `output[idx] = max(input[idx], 0.0)`. **Implementation** - `add_unary_op(graph, in, out, wgsl_source, wg_size_x, op_name)` centralizes: the fp32/4-byte-alignment and same-size guards, `utils::clamp_workgroup_size` + `utils::compute_1d_workgroup_count` for the 1D dispatch, the `wg_size` override constant, the uniform (`num_elements`) via `utils::make_uniform`, and the three-binding pipeline via `utils::make_compute_pipeline`. - Dynamic shapes are supported: a `graph.add_tensor_resize_hook` recomputes `num_elements`, rewrites the uniform via `wgpuQueueWriteBuffer`, and updates the dispatch's workgroup count for the live shape; the graph owns the uniform buffer so the hook can rewrite it. - Both ops self-register: `aten.sigmoid.default -> sigmoid_impl` and `aten.relu.default -> relu_impl`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp` (`add_unary_op_node`); Vulkan expresses `relu` as `clamp(0, inf)`, whereas this kernel uses a direct `max(x, 0.0)`. **Constraints** — fp32 only (both operands 4-byte aligned); input and output must have identical byte size (same-shape elementwise); 1D dispatch only (throws past the 65535 workgroup cap). Co-authored-with: Claude Code. ghstack-source-id: 405949742 @exported-using-ghexport Differential Revision: [D110836664](https://our.internmc.facebook.com/intern/diff/D110836664/) --- backends/webgpu/runtime/ops/relu/relu.wgsl | 20 +++ backends/webgpu/runtime/ops/relu/relu_wgsl.h | 44 ++++++ .../webgpu/runtime/ops/sigmoid/UnaryOp.cpp | 135 ++++++------------ .../webgpu/runtime/ops/sigmoid/sigmoid.wgsl | 10 +- .../webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h | 14 +- 5 files changed, 120 insertions(+), 103 deletions(-) create mode 100644 backends/webgpu/runtime/ops/relu/relu.wgsl create mode 100644 backends/webgpu/runtime/ops/relu/relu_wgsl.h diff --git a/backends/webgpu/runtime/ops/relu/relu.wgsl b/backends/webgpu/runtime/ops/relu/relu.wgsl new file mode 100644 index 00000000000..0d9885421f9 --- /dev/null +++ b/backends/webgpu/runtime/ops/relu/relu.wgsl @@ -0,0 +1,20 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = max(input[idx], 0.0); +} diff --git a/backends/webgpu/runtime/ops/relu/relu_wgsl.h b/backends/webgpu/runtime/ops/relu/relu_wgsl.h new file mode 100644 index 00000000000..ae74266ce8a --- /dev/null +++ b/backends/webgpu/runtime/ops/relu/relu_wgsl.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from relu.wgsl - DO NOT EDIT. +// wgsl-sha256: 2f7fac19d55cb7e55749f2bc1856278c1f2afaf0bc7ff8e663fb9e14b4188199 +inline constexpr const char* kReluWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = max(input[idx], 0.0); +} +)"; + +inline constexpr uint32_t kReluWorkgroupSizeX = 256; +inline constexpr uint32_t kReluWorkgroupSizeY = 1; +inline constexpr uint32_t kReluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp index 7c8b8c0e9ce..fc228853d93 100644 --- a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp +++ b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -39,31 +40,20 @@ void add_unary_op( const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); - if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { - throw std::runtime_error(std::string(op_name) + ": null buffer binding"); - } - // 4-byte (fp32) alignment guard on both operands; also the dtype guard. - if (in_tensor.nbytes % sizeof(float) != 0 || - out_tensor.nbytes % sizeof(float) != 0) { - throw std::runtime_error( - std::string(op_name) + ": operand not 4-byte aligned"); - } - if (in_tensor.nbytes != out_tensor.nbytes) { - throw std::runtime_error( - std::string(op_name) + ": input/output size mismatch"); - } + utils::check_elementwise_fp32_io(in_tensor, out_tensor, op_name); uint32_t num_elements = static_cast(out_tensor.nbytes / sizeof(float)); + // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 + // per-dim ceiling. The shader decodes idx via num_workgroups.x, so the live + // count_x sets the stride at runtime (resize-safe, no override to re-bake). uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); - uint32_t workgroup_count = - utils::compute_1d_workgroup_count(device, num_elements, wg_size, op_name); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_elements, wg_size, op_name); - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); UnaryParams params = {}; params.num_elements = num_elements; @@ -72,71 +62,29 @@ void add_unary_op( utils::make_uniform(device, ¶ms, sizeof(UnaryParams)); graph.add_uniform_buffer_bytes(sizeof(UnaryParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: input (read storage) + output (storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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); - - 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[3] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(UnaryParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + // input (read storage) + output (storage) + params. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl_source, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(UnaryParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); // Dynamic shapes: recompute num_elements/dispatch for the live shape. WGPUBuffer params_buf = uniform_buffer; @@ -149,18 +97,12 @@ void add_unary_op( UnaryParams p = {}; p.num_elements = static_cast(numel); wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - g.device(), - static_cast(numel), - wg_size, - "unary(resize)"); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), static_cast(numel), wg_size, "unary(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - // Release intermediates (pipeline + bind_group are kept by dispatch). - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } @@ -176,10 +118,17 @@ void sigmoid_impl(WebGPUGraph& graph, const std::vector& args) { "sigmoid"); } +void relu_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.relu.default args: [in, out] + add_unary_op( + graph, args.at(0), args.at(1), kReluWGSL, kReluWorkgroupSizeX, "relu"); +} + } // namespace WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.sigmoid.default, sigmoid_impl); + WEBGPU_REGISTER_OP(aten.relu.default, relu_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl b/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl index 09b3e5457b3..a57dd46f150 100644 --- a/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl +++ b/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl @@ -6,11 +6,13 @@ struct Params { } @group(0) @binding(2) var params: Params; -override wg_size: u32 = 64u; +override wg_size: u32 = 256; -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h b/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h index 48e6efc607a..fc87b981d34 100644 --- a/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h +++ b/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sigmoid.wgsl - DO NOT EDIT. -// wgsl-sha256: 70395dbb107b8b95ae13c0a6fb12a8415c561c645da0347294c92904314ae84c +// wgsl-sha256: 79a1554e9d957e10c0f4379b5e0240191743952198e9ed9c0342402bacd356de inline constexpr const char* kSigmoidWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -23,11 +23,13 @@ struct Params { } @group(0) @binding(2) var params: Params; -override wg_size: u32 = 64u; +override wg_size: u32 = 256; -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } @@ -35,7 +37,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } )"; -inline constexpr uint32_t kSigmoidWorkgroupSizeX = 64; +inline constexpr uint32_t kSigmoidWorkgroupSizeX = 256; inline constexpr uint32_t kSigmoidWorkgroupSizeY = 1; inline constexpr uint32_t kSigmoidWorkgroupSizeZ = 1; From f51bdd8a0a68e706a93074c947da8051e8c019fe Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:22 -0700 Subject: [PATCH 02/22] [ExecuTorch][WebGPU] relu op tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20864 Splits the `relu` op tests into their own diff, stacked directly above the `relu` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_relu.py` and registers the `relu` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949743 @exported-using-ghexport Differential Revision: [D111072725](https://our.internmc.facebook.com/intern/diff/D111072725/) --- backends/webgpu/test/op_tests/cases.py | 28 ++++++++++++++++++++++++++ backends/webgpu/test/ops/test_relu.py | 27 +++++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 backends/webgpu/test/ops/test_relu.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 4ca730428ce..f37e157866d 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -719,3 +719,31 @@ def _max_pool2d_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_relu import ( + _det_input as _relu_det_input, + ReluModule, +) + + +@register_op_test("relu") +def _relu_suite() -> WebGPUTestSuite: + # SAM2/SAM3 mask-decoder MLP path; tiny + a decoder_mlp-shaped case. + return WebGPUTestSuite( + module_factory=lambda: ReluModule(), + cases=[ + Case( + name="tiny", + construct={}, + inputs=(InputSpec(shape=(1, 4, 8), gen=_relu_det_input),), + ), + Case( + name="decoder_mlp", + construct={}, + inputs=(InputSpec(shape=(1, 576, 768), gen=_relu_det_input),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_relu.py b/backends/webgpu/test/ops/test_relu.py new file mode 100644 index 00000000000..f8303fcde58 --- /dev/null +++ b/backends/webgpu/test/ops/test_relu.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.relu.default` module for the WebGPU op-test framework. + +ReLU is on the SAM2/SAM3 mask-decoder MLP path. +""" + +import torch + + +class ReluModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x) + + +def _det_input(shape) -> torch.Tensor: + # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives + # so the clamp is exercised, not just an identity pass-through. + n = 1 + for s in shape: + n *= s + idx = torch.arange(n, dtype=torch.int64) + return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape) From 2fb2297e5f55a5243d04c6d68a141e377ad1c71b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:22 -0700 Subject: [PATCH 03/22] [ExecuTorch][WebGPU] slice_copy Double-arg scalar handling + view_copy/alias glue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20865 **Hardens `slice_copy` scalar-arg decoding against edge-dialect `Double`-serialized indices and adds the `view_copy` / `alias` / `clone` reshape pass-throughs needed for graph glue.** **Problem** — two graph-glue gaps: (1) the edge dialect sometimes serializes an integer `slice` index (`dim` / `start` / `end` / `step`) as a floating-point `Double` (e.g. a `0` start), which the `slice` handler rejected as unsupported; (2) contiguous reshape / aliasing ops (`view_copy`, `alias_copy`, `clone`, `_clone_dim_order`) had no handler, breaking otherwise-delegatable subgraphs. **Solution** — Before — a `Double`-typed slice index threw, and reshape/alias ops had no handler. After — `slice_copy` scalar reads accept an integral `Double` (truncating to the int index) and reject only a genuinely fractional one, while `SymInt` (dynamic start/end) and `Null` (default) still resolve as before; and `view_copy` / `alias_copy` / `clone` / `_clone_dim_order` all lower to a single contiguous flat copy (or an in-place no-op when input and output alias the same buffer). **Implementation**: - `read_scalar` (`dim` / `step`) and `read_index` (`start` / `end`) switch on the value type: `Int` (`INT64_MAX` -> default), `Double` -> truncated int iff it round-trips (`static_cast(d)` back to `d`) else throw `"non-integral ..."` (NaN and out-of-`int64`-range doubles are rejected before the cast, since casting them is UB), `Null` -> default; `read_index` additionally resolves a `SymInt` via `read_symint`. - The slice kernel is an index gather: `out_bufi -> in_bufi` by walking per-dim strides, with the sliced dim's input coord `= start + coord*step`; dynamic `start` / `end` / `SymInt` are handled by a resize hook that recomputes the live `out[dim]` length (ceiling division) and rewrites the meta/params uniforms plus the dispatch count (mirrors Vulkan `resize_slice_copy_node`). - `add_flat_copy` (shared by all the reshape/alias ops) type-checks both args are tensors, guards 4-byte alignment and equal `nbytes` (a view preserves `numel`, so this also prevents an OOB copy), then either skips the copy when `in.buffer == out.buffer` (aliased, already in place; `CopyBufferToBuffer` rejects `src == dst`) or emits a buffer-to-buffer copy; a resize hook keeps the live output shape and copy byte-count in sync under dynamic shapes. - `_clone_dim_order` ignores its `dim_order` arg (the AOT pass elides it via shape and dtype). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Slice.cpp` (`normalize_idx` / `INT64_MAX` default and the ceiling-division length) and `backends/vulkan/runtime/graph/ops/impl/View.cpp` (the `view_buffer` no-remap contiguous reshape). **Constraints** — fp32 (4-byte-aligned) operands; `slice` requires `step >= 1` and an in-range `dim`; a fractional `Double` index is a hard error, not truncated; `view` / `alias` / `clone` require equal input/output `numel` (contiguous reshape only, no layout remap). Co-authored-with: Claude Code. ghstack-source-id: 405949746 @exported-using-ghexport Differential Revision: [D110836670](https://our.internmc.facebook.com/intern/diff/D110836670/) --- backends/webgpu/runtime/ops/slice/Slice.cpp | 37 ++++++++++++++++++- .../webgpu/runtime/ops/view_copy/ViewCopy.cpp | 3 ++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/backends/webgpu/runtime/ops/slice/Slice.cpp b/backends/webgpu/runtime/ops/slice/Slice.cpp index 261020e05d5..4fbbfccc5fe 100644 --- a/backends/webgpu/runtime/ops/slice/Slice.cpp +++ b/backends/webgpu/runtime/ops/slice/Slice.cpp @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -30,7 +31,10 @@ struct SliceParams { uint32_t _pad; }; -// Read scalar arg: Int->value (INT64_MAX->default), Null->default, else throw. +// Read scalar arg: Int->value (INT64_MAX->default), Double->truncated int if +// integral (the edge dialect may serialize an integer index as a float, e.g. +// a 0 start; a fractional Double throws, it is not a valid index), +// Null->default, else throw. int64_t read_scalar(WebGPUGraph& graph, int id, int64_t dflt, const char* what) { switch (graph.get_value_type(id)) { @@ -38,6 +42,20 @@ read_scalar(WebGPUGraph& graph, int id, int64_t dflt, const char* what) { const int64_t v = graph.get_int(id); return v == INT64_MAX ? dflt : v; } + case WebGPUGraph::ValueType::Double: { + const double d = graph.get_double(id); + // Casting a NaN or out-of-int64-range double is undefined behavior; + // reject before the cast, not after. + if (std::isnan(d) || d < -9223372036854775808.0 || + d >= 9223372036854775808.0) { + throw std::runtime_error(std::string("slice: non-integral ") + what); + } + const int64_t v = static_cast(d); + if (static_cast(v) != d) { + throw std::runtime_error(std::string("slice: non-integral ") + what); + } + return v; + } case WebGPUGraph::ValueType::Null: return dflt; default: @@ -46,7 +64,8 @@ read_scalar(WebGPUGraph& graph, int id, int64_t dflt, const char* what) { } } -// Read a slice index (start/end) that MAY be a dynamic SymInt; else Int/Null. +// Read a slice index (start/end) that MAY be a dynamic SymInt; else Int/Double +// (truncated int if integral, mirrors read_scalar)/Null. int64_t read_index(WebGPUGraph& graph, int id, int64_t dflt) { switch (graph.get_value_type(id)) { case WebGPUGraph::ValueType::SymInt: @@ -55,6 +74,20 @@ int64_t read_index(WebGPUGraph& graph, int id, int64_t dflt) { const int64_t v = graph.get_int(id); return v == INT64_MAX ? dflt : v; } + case WebGPUGraph::ValueType::Double: { + const double d = graph.get_double(id); + // Casting a NaN or out-of-int64-range double is undefined behavior; + // reject before the cast, not after. + if (std::isnan(d) || d < -9223372036854775808.0 || + d >= 9223372036854775808.0) { + throw std::runtime_error("slice: non-integral start/end index"); + } + const int64_t v = static_cast(d); + if (static_cast(v) != d) { + throw std::runtime_error("slice: non-integral start/end index"); + } + return v; + } case WebGPUGraph::ValueType::Null: return dflt; default: diff --git a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp index eb274481893..7b3b44eb86f 100644 --- a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp +++ b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp @@ -103,6 +103,9 @@ void clone_impl(WebGPUGraph& graph, const std::vector& args) { WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.view_copy.default, view_copy_impl); WEBGPU_REGISTER_OP(aten.clone.default, clone_impl); + WEBGPU_REGISTER_OP(aten.alias_copy.default, clone_impl); + // _clone_dim_order ignores dim_order (AOT pass elides via shape+dtype). + WEBGPU_REGISTER_OP(dim_order_ops._clone_dim_order.default, clone_impl); } } // namespace executorch::backends::webgpu From 88214d17dce37eea6b90bcbc9547a48add6ce0b1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:22 -0700 Subject: [PATCH 04/22] [ExecuTorch][WebGPU] slice_copy native golden test Pull Request resolved: https://github.com/pytorch/executorch/pull/20866 Splits the `SliceDoubleStart` native golden test into its own diff, stacked directly above the `slice_copy` / `view_copy` glue op (op below, tests above). Adds the double-start slice regression case to `test/test_webgpu_native.cpp`, covering an edge-dialect-serialized Double-typed slice `start` argument. Co-authored-with: Claude Code. ghstack-source-id: 405949754 @exported-using-ghexport Differential Revision: [D111072708](https://our.internmc.facebook.com/intern/diff/D111072708/) --- backends/webgpu/test/test_webgpu_native.cpp | 203 ++++++++++++++++++++ 1 file changed, 203 insertions(+) diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index bae1e4e8863..0a16b04a356 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -16,6 +16,7 @@ #include #include +#include #include #include @@ -25,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -1626,6 +1628,202 @@ const EmbConfig kEmbConfigs[] = { 2048}, }; +// Regression: an edge-dialect-serialized integer slice `start` can arrive as a +// Double (e.g. Florence-2 DaViT serialized start=0 as Double 0.0), which once +// threw "slice: dynamic/unsupported start". The Python op-tests can only emit +// an Int start (the serializer keys on the Python runtime type), so this case +// is unreachable from a .pte export -- it must be built natively. Here we +// hand-author a VkGraph flatbuffer whose slice `start` is a Double value, run +// it on the device, and assert the gather matches in[start + i*step]. Two +// cases: a Double 0.0 (identity) and a Double 2.0 (non-zero offset). +static bool test_slice_double_start_case(double start_d, int out_len) { + namespace vk = vkgraph; + // Slice x[1, kInLen] along dim 1 with `start` as a Double; out is + // [1,out_len]. + constexpr int kInLen = 6; + printf( + "\n--- Test: slice Double start (start=%.1f -> out[1,%d]) ---\n", + start_d, + out_len); + + // Value ids: 0=in tensor, 1=dim(Int), 2=start(Double), 3=end(Int), + // 4=step(Int), 5=out tensor. dims are uint vectors; tensors take distinct + // mem_obj_ids so build() allocates a real Storage buffer for each. + ::flatbuffers::FlatBufferBuilder fbb; + + std::vector in_dims = {1u, static_cast(kInLen)}; + std::vector out_dims = {1u, static_cast(out_len)}; + + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &in_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + // The value under test: `start` serialized as a Double, not an Int. + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/start_d).Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Int, + vk::CreateInt(fbb, /*int_val=*/kInLen).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &out_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/1) + .Union())); + + std::vector args = {0, 1, 2, 3, 4, 5}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.slice_copy.Tensor", &args)); + + std::vector input_ids = {0}; + std::vector output_ids = {5}; + auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + } catch (const std::exception& e) { + printf("FAIL: graph build threw: %s\n", e.what()); + return false; + } + + std::vector in(kInLen); + for (int i = 0; i < kInLen; i++) { + in[i] = static_cast(i) + 0.5f; + } + std::vector inputs(1); + inputs[0] = {in.data(), in.size() * sizeof(float), false}; + std::vector out(out_len, -1.0f); + std::vector> outputs(1); + outputs[0] = {out.data(), out.size() * sizeof(float)}; + try { + graph.copy_inputs(inputs); + graph.execute(); + graph.copy_outputs(outputs); + } catch (const std::exception& e) { + printf("FAIL: slice execute threw: %s\n", e.what()); + return false; + } + + const int start = static_cast(start_d); + float max_abs_err = 0.0f; + for (int i = 0; i < out_len; i++) { + const float expected = in[start + i]; // step == 1 + max_abs_err = std::max(max_abs_err, std::abs(out[i] - expected)); + } + printf("Max abs error: %e (checked %d elements)\n", max_abs_err, out_len); + if (max_abs_err != 0.0f) { // pure gather: must be bit-exact + printf("FAIL: slice Double-start gather mismatch\n"); + return false; + } + printf("PASS: slice Double start (start=%.1f)\n", start_d); + return true; +} + +// Negative control: a Double `start` that is fractional, NaN, or outside the +// int64 range must throw (never silently truncate, never invoke UB via the +// int64_t cast). +static bool test_slice_double_start_rejects(double bad_start) { + namespace vk = vkgraph; + constexpr int kInLen = 6; + printf("\n--- Test: slice Double start REJECTS (start=%g) ---\n", bad_start); + + ::flatbuffers::FlatBufferBuilder fbb; + std::vector in_dims = {1u, static_cast(kInLen)}; + std::vector out_dims = {1u, static_cast(kInLen)}; + + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &in_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/bad_start).Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Int, + vk::CreateInt(fbb, /*int_val=*/kInLen).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &out_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/1) + .Union())); + + std::vector args = {0, 1, 2, 3, 4, 5}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.slice_copy.Tensor", &args)); + std::vector input_ids = {0}; + std::vector output_ids = {5}; + auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + } catch (const std::exception& e) { + printf("PASS: rejected as expected: %s\n", e.what()); + return true; + } + printf( + "FAIL: expected a throw for start=%g, graph.build() succeeded\n", + bad_start); + return false; +} + +static bool test_slice_double_start() { + // start=0.0 (identity copy) + start=2.0 (non-zero gather offset). + bool ok = true; + ok = test_slice_double_start_case(/*start_d=*/0.0, /*out_len=*/6) && ok; + ok = test_slice_double_start_case(/*start_d=*/2.0, /*out_len=*/4) && ok; + // Reject: fractional, NaN, and out-of-int64-range Doubles. + ok = test_slice_double_start_rejects(/*bad_start=*/0.5) && ok; + ok = test_slice_double_start_rejects( + /*bad_start=*/std::numeric_limits::quiet_NaN()) && + ok; + ok = test_slice_double_start_rejects(/*bad_start=*/1e300) && ok; + return ok; +} + // apply_rotary_emb on-GPU configs: multi + decode (env-gated, run-if-present). struct RopeConfig { const char* name; @@ -1806,6 +2004,11 @@ TEST(WebGPUNative, Prepack) { test_prepack(g_prepack_model_path, g_prepack_golden_path); } +TEST(WebGPUNative, SliceDoubleStart) { + EXPECT_TRUE(test_slice_double_start()) + << "slice Double-start gather/reject checks failed"; +} + TEST(WebGPUNative, Prepack2) { if (g_prepack2_model_path.empty() || g_prepack2_golden_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_PREPACK2_MODEL/GOLDEN not set"; From aec6573adedaca94f7e6f5352ef999badb14b0ad Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:23 -0700 Subject: [PATCH 05/22] [ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20871 **Problem:** the fused `et_vk.sdpa` QK kernel runs one thread per (b,h,s) row with vec4 loads — ideal for standard attention, but on channel attention (DaViT/Florence, where `S_q = head_dim ~= 32`) `num_rows = B*H*S_q` is tiny, so only a handful of workgroups run serially over a huge `S_kv*D`, starving the GPU (the (2,1,1)@103ms dispatch). **Solution:** add a per-entry QK kernel (one thread per (b,h,s,c) attention entry, 2D-folded) and host-route to it when `num_rows` is below an occupancy floor (4096); standard attention keeps the per-row + vec4 path unchanged. **Before:** `et_vk_sdpa_qk` (per-row, vec4) — the only QK kernel; channel-attn shapes are occupancy-starved. **After:** router picks `et_vk_sdpa_qk_entry` (per-entry, scalar, 2D-folded) for small `num_rows`, else the unchanged per-row kernel. **Implementation:** - New `et_vk_sdpa_qk_entry.wgsl` (+ generated header) — same bindings and `Params` as the per-row kernel, so it is a drop-in under `layout:"auto"`; writes a layout-identical `attn[B,H,S_q,S_kv]` (`attn[idx]`), so softmax/AV are unchanged and either branch is numerically correct — the floor is a pure perf knob. - `EtVkSdpa.cpp` selects the shader and a 2D dispatch (`compute_2d_workgroup_count`, mirroring the softmax grid) when routed, else the existing 1D per-row dispatch; the grid + dispatch-limit check is computed up front (throw before any buffer alloc -> no leak). - Mirrors the codebase's host shape-router precedents (`LinearFp32.cpp` `K%4` vec4 selection, `Sdpa.cpp` variant selection). **Constraints:** per-entry drops vec4, so it only wins when the per-row path is occupancy-starved (small `num_rows`); the 4096 floor is Canary-tuned. Co-authored-with: Claude Code. ghstack-source-id: 405949750 @exported-using-ghexport Differential Revision: [D110994975](https://our.internmc.facebook.com/intern/diff/D110994975/) --- .../runtime/ops/et_vk_sdpa/EtVkSdpa.cpp | 46 ++++++++--- .../ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl | 54 +++++++++++++ .../ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h | 78 +++++++++++++++++++ 3 files changed, 168 insertions(+), 10 deletions(-) create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp index dbe6b03f30c..92a53cb7574 100644 --- a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -149,13 +150,33 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { const size_t aw_bytes = static_cast(aw_numel) * sizeof(float); // Up-front dispatch-limit checks (throw BEFORE any buffer alloc → no leak). - // QK = one thread per (b,h,s) row; AV = one per (b,h,s,d4) vec4 (4 output - // elements/thread); softmax = one workgroup per row. compute_1d_workgroup_ - // count throws past the limit. - const uint32_t qk_wg_size = - utils::clamp_workgroup_size(device, kEtVkSdpaQkWorkgroupSizeX); - const uint32_t qk_wg_count = utils::compute_1d_workgroup_count( - device, static_cast(num_rows), qk_wg_size, "et_vk_sdpa_qk"); + // QK: per-row (one thread per (b,h,s) row, vec4 loads) is fastest for + // standard attention, but starves the GPU on channel attention (S_q = + // head_dim, so num_rows is tiny → few workgroups serial over a huge S_kv*D). + // Route to the per-entry kernel (one thread per (b,h,s,c) attn entry, + // 2D-folded) below an occupancy floor; both write a layout-identical + // attn[B,H,S_q,S_kv] so softmax/AV are unchanged, and either branch is + // numerically correct, so the floor is a perf knob only (Canary M4 Pro: + // per-entry ~15-30x faster at num_rows <= 256, per-row wins at num_rows >= + // 8192). AV = one per (b,h,s,d4) vec4; softmax = one workgroup per row. + constexpr uint32_t kQkEntryOccupancyFloor = 4096u; + const bool qk_per_entry = num_rows < kQkEntryOccupancyFloor; + const uint32_t qk_wg_size = utils::clamp_workgroup_size( + device, + qk_per_entry ? kEtVkSdpaQkEntryWorkgroupSizeX + : kEtVkSdpaQkWorkgroupSizeX); + uint32_t qk_wg_count = 0; + utils::WgCount qk_entry_grid = {}; + if (qk_per_entry) { + qk_entry_grid = utils::compute_2d_workgroup_count( + device, + static_cast(aw_numel), + qk_wg_size, + "et_vk_sdpa_qk_entry"); + } else { + qk_wg_count = utils::compute_1d_workgroup_count( + device, static_cast(num_rows), qk_wg_size, "et_vk_sdpa_qk"); + } const uint32_t av_wg_size = utils::clamp_workgroup_size(device, kEtVkSdpaAvWorkgroupSizeX); const uint64_t out_numel4 = @@ -179,7 +200,7 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer attn_buf = graph.create_scratch_buffer(aw_bytes); WGPUBuffer softmax_buf = graph.create_scratch_buffer(aw_bytes); - // ---- Dispatch 1: QK (one thread per (b,h,s) row) ---- + // ---- Dispatch 1: QK (per-row for standard attn, per-entry for channel) ---- { QkParams p = {}; p.B = B; @@ -205,7 +226,7 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, - kEtVkSdpaQkWGSL, + qk_per_entry ? kEtVkSdpaQkEntryWGSL : kEtVkSdpaQkWGSL, { {0, WGPUBufferBindingType_Storage, attn_buf, aw_bytes}, {1, WGPUBufferBindingType_ReadOnlyStorage, q.buffer, q.nbytes}, @@ -222,7 +243,12 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { &wg_const, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count}); + if (qk_per_entry) { + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, qk_entry_grid.x, qk_entry_grid.y); + } else { + graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count}); + } wgpuBufferRelease(uniform_buffer); if (mask.owned_dummy != nullptr) { diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl new file mode 100644 index 00000000000..605ac52c89e --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl @@ -0,0 +1,54 @@ +@group(0) @binding(0) var attn: array; +@group(0) @binding(1) var q: array; +@group(0) @binding(2) var k: array; +@group(0) @binding(3) var mask: array; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + has_mask: u32, + _pad0: u32, + scale: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D], +// k [B, H, S_kv, D]. ONE thread per ENTRY (b,h,s,c) of attn_weights +// [B, H, S_q, S_kv] = one D-length dot. Previously this was one thread per ROW +// (looping c) — fine for window/self attention but catastrophic for DaViT +// CHANNEL attention where S_q = head_dim (~32), so B*H*S_q was 128/256/512/1024 +// → only 2/4/8/16 workgroups serial over the huge spatial D (the (2,1,1)@103ms +// dispatch). Parallelizing over all B*H*S_q*S_kv entries (2D-folded past the +// 65535 ceiling, mirroring the softmax phase) gives S_kv× more threads. +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) nwg: vec3) { + let aw_numel = params.B * params.H * params.S_q * params.S_kv; + let idx = gid.x + gid.y * (nwg.x * wg_size); // 2D-folded linear entry id + if (idx >= aw_numel) { + return; + } + let c = idx % params.S_kv; + let row = idx / params.S_kv; // (b,h,s) flattened + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let qbase = ((b * params.H + h) * params.S_q + s) * params.D; + let kbase = ((b * params.H + h) * params.S_kv + c) * params.D; + var acc: f32 = 0.0; + for (var d: u32 = 0u; d < params.D; d = d + 1u) { + acc = acc + q[qbase + d] * k[kbase + d]; + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[idx]; + } + attn[idx] = acc; // attn is [B,H,S_q,S_kv] row-major -> index == idx +} diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h new file mode 100644 index 00000000000..6918c6e2621 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h @@ -0,0 +1,78 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from et_vk_sdpa_qk_entry.wgsl - DO NOT EDIT. +// wgsl-sha256: 95ff9f19757d23c41d82054b779c02fb8d83531eef3e8cda5f3a78dde42a4e58 +inline constexpr const char* kEtVkSdpaQkEntryWGSL = R"( +@group(0) @binding(0) var attn: array; +@group(0) @binding(1) var q: array; +@group(0) @binding(2) var k: array; +@group(0) @binding(3) var mask: array; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + has_mask: u32, + _pad0: u32, + scale: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D], +// k [B, H, S_kv, D]. ONE thread per ENTRY (b,h,s,c) of attn_weights +// [B, H, S_q, S_kv] = one D-length dot. Previously this was one thread per ROW +// (looping c) — fine for window/self attention but catastrophic for DaViT +// CHANNEL attention where S_q = head_dim (~32), so B*H*S_q was 128/256/512/1024 +// → only 2/4/8/16 workgroups serial over the huge spatial D (the (2,1,1)@103ms +// dispatch). Parallelizing over all B*H*S_q*S_kv entries (2D-folded past the +// 65535 ceiling, mirroring the softmax phase) gives S_kv× more threads. +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) nwg: vec3) { + let aw_numel = params.B * params.H * params.S_q * params.S_kv; + let idx = gid.x + gid.y * (nwg.x * wg_size); // 2D-folded linear entry id + if (idx >= aw_numel) { + return; + } + let c = idx % params.S_kv; + let row = idx / params.S_kv; // (b,h,s) flattened + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let qbase = ((b * params.H + h) * params.S_q + s) * params.D; + let kbase = ((b * params.H + h) * params.S_kv + c) * params.D; + var acc: f32 = 0.0; + for (var d: u32 = 0u; d < params.D; d = d + 1u) { + acc = acc + q[qbase + d] * k[kbase + d]; + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[idx]; + } + attn[idx] = acc; // attn is [B,H,S_q,S_kv] row-major -> index == idx +} +)"; + +inline constexpr uint32_t kEtVkSdpaQkEntryWorkgroupSizeX = 64; +inline constexpr uint32_t kEtVkSdpaQkEntryWorkgroupSizeY = 1; +inline constexpr uint32_t kEtVkSdpaQkEntryWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 70c66d1d81f8e164b8d3e4f84624a1cb5bad0d9f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:23 -0700 Subject: [PATCH 06/22] [ExecuTorch][WebGPU] et_vk.sdpa channel-attention tests Pull Request resolved: https://github.com/pytorch/executorch/pull/20872 Splits the channel-attention routing case out of the `et_vk.sdpa` per-entry-QK op diff into its own test diff, stacked directly above it (op below, tests above). Adds the `chattn_davit` case to the `et_vk_sdpa` suite in `test/op_tests/cases.py`, exercising the per-entry QK kernel path (num_rows below the per-row floor). Co-authored-with: Claude Code. ghstack-source-id: 405949753 @exported-using-ghexport Differential Revision: [D111072706](https://our.internmc.facebook.com/intern/diff/D111072706/) --- backends/webgpu/test/op_tests/cases.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index f37e157866d..d803bd453a7 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -501,7 +501,10 @@ def _sdpa_mask(b, h, sq, skv): def _et_vk_sdpa_suite() -> WebGPUTestSuite: # Non-causal fused attention (Florence-2 vision + BART, via the et_vk source # transform). Covers self-attn, an asymmetric S_q != S_kv (cross-attn) case, - # an additive mask (BART), and D=128 (Voxtral/DaViT) through the vec4 kernels. + # an additive mask (BART), and D=128 (Voxtral/DaViT). The QK dispatch is + # occupancy-routed: chattn_davit (num_rows = B*H*S_q = 256 < the 4096 floor) + # exercises the per-entry QK kernel; selfattn_siglip (num_rows = 6912) is the + # per-row guard. Both branches must match the same fp32 golden. def qkv(b, h, sq, skv, d): return ( InputSpec(shape=(b, h, sq, d), gen=_sdpa_randn), @@ -521,6 +524,7 @@ def qkv(b, h, sq, skv, d): inputs=qkv(1, 4, 8, 8, 16), ), Case(name="d128_voxtral", inputs=qkv(1, 4, 6, 6, 128)), + Case(name="chattn_davit", inputs=qkv(1, 8, 32, 256, 64)), ], golden_dtype="float32", atol=1e-4, From 16dc8c28d7df63493050776762af904b5314c8e2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:24 -0700 Subject: [PATCH 07/22] [ExecuTorch][WebGPU] et_vk conv2d: route standard (groups==1) conv through an im2col tiled GEMM (1.1-2.4x) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20873 **Problem:** the direct conv2d kernel runs one thread per output element and re-reads the input receptive field from global memory for every output — zero cross-thread reuse. For the patch-embed stem (3-channel RGB) the vec4-over-IC path is inert (icpg=3 fails the `%4` gate), so it runs the scalar direct path with no reuse at all. **Solution:** route groups==1 non-transposed convs through an implicit-im2col tiled GEMM that reuses the linear tiled-GEMM skeleton — M=OC, N=B*OH*OW, K=IC*KH*KW; shared-memory 32x32 tiles + 4x4 register blocking; the input is im2col-sampled on the fly (out-of-range -> 0.0 implements padding). Grouped/depthwise/transpose stay on the direct/gather kernels. **Before:** every conv -> direct kernel (scalar, or vec4-over-IC when icpg%4==0), no input reuse. **After:** groups==1 -> `conv2d_gemm` (shared-mem tiling + register blocking, input-tile reuse across output positions); grouped/transpose -> unchanged. **Implementation:** - New `conv2d_gemm.wgsl` (+ generated header): forks `linear_fp32_tiled.wgsl` — `read_a` loads the weight `[OC, K]`, `read_b` im2col-samples the input (decodes n->(b,oh,ow), kk->(ic,kh,kw); ih=oh*sH-pH+kh*dH; bounds-check->0), bias per-row (OC), output written NCHW. Reuses the existing `ConvParams` uniform. - `Conv2d.cpp` branches on `groups==1`: GEMM via `compute_tile_grid_2d` + `add_dispatch_2d` (mirrors `LinearFp32.cpp`); else the existing direct dispatch. The grouped path is byte-identical; both grids are computed before any buffer alloc (throw-before-leak). Mirrors Vulkan's own `should_use_conv2d_im2col` groups==1 routing. **Constraints:** scalar GEMM (no vec4) — NCHW's channel stride isn't contiguous, so vec4-over-K would be a strided gather (no compute win on Apple's scalar ALU); ORT skips vec4 for NCHW too. Co-authored-with: Claude Code. ghstack-source-id: 405949751 @exported-using-ghexport Differential Revision: [D110995347](https://our.internmc.facebook.com/intern/diff/D110995347/) --- .../runtime/ops/et_vk_conv2d/Conv2d.cpp | 120 +++++++++++---- .../runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl | 116 +++++++++++++++ .../ops/et_vk_conv2d/conv2d_gemm_wgsl.h | 140 ++++++++++++++++++ 3 files changed, 345 insertions(+), 31 deletions(-) create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp index 432fdb9de5b..5f127f6d5f2 100644 --- a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -325,15 +326,39 @@ void conv2d_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t icpg = IC / groups; const bool use_vec4 = (icpg % 4u == 0u); - // Up-front (throw before any buffer alloc → no leak-on-throw). - // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 - // ceiling (the documented SAM FpnNeck @1008^2 blocker). stride_x lets the - // shader decode i = gid.y*stride_x + gid.x. - utils::DispatchGrid grid = utils::compute_dispatch_grid( - device, - utils::checked_u32(out_numel, "conv2d"), - kConv2dWorkgroupSizeX, - "et_vk_conv2d"); + // groups==1 non-transposed -> im2col tiled GEMM (reuses the linear tiled-GEMM + // skeleton: M=OC, N=B*OH*OW, K=IC*KH*KW; input im2col-sampled on the fly, + // out-of-range -> 0.0 for padding). Canary M4 Pro: 1.1-2.4x over the direct + // kernel across stem/FPN shapes (biggest on the RGB stem, where the + // vec4-over-IC path is inert). Grouped/depthwise stay on the direct kernel + // (grouped GEMM is block-diagonal; mirrors ORT). + const bool use_gemm = (groups == 1u); + + // Compute the dispatch grid UP-FRONT, before any buffer alloc, so a throw + // (grid exceeds the device dispatch/tile limit) can't leak the uniform/bias. + constexpr uint32_t kConv2dGemmTile = + 32u; // fixed @workgroup_size(8,8), TILE=32 + utils::WgCount gemm_grid = {}; + utils::DispatchGrid direct_grid = {}; + if (use_gemm) { + // 2D tile grid over (N cols, M rows); folds past the 65535 per-dim ceiling. + gemm_grid = utils::compute_tile_grid_2d( + device, + utils::checked_u32( + uint64_t(B) * uint64_t(OH) * uint64_t(OW), "et_vk_conv2d_gemm"), + OC, + kConv2dGemmTile, + "et_vk_conv2d_gemm"); + } else { + // Adaptive 1D->2D dispatch: wg=clamp(256) + 2D-spill past the 65535 ceiling + // (the SAM FpnNeck @1008^2 blocker); stride_x decodes + // i=gid.y*stride_x+gid.x. + direct_grid = utils::compute_dispatch_grid( + device, + utils::checked_u32(out_numel, "conv2d"), + kConv2dWorkgroupSizeX, + "et_vk_conv2d"); + } ConvParams params = {}; params.B = B; @@ -364,29 +389,62 @@ void conv2d_impl(WebGPUGraph& graph, const std::vector& args) { has_bias, has_bias ? graph.get_tensor(bias_id).buffer : nullptr, has_bias ? graph.get_tensor(bias_id).nbytes : 0); - auto constants = utils::make_grid_constants(grid); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - device, - use_vec4 ? kConv2dVec4WGSL : kConv2dWGSL, - { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - weight.buffer, - weight.nbytes}, - {3, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes}, - {4, - WGPUBufferBindingType_Uniform, - uniform_buffer, - sizeof(ConvParams)}, - }, - constants.data(), - constants.size()); - graph.add_dispatch_2d( - bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + if (use_gemm) { + // vec4-over-IC is inert for NCHW (strided channel gather), so the GEMM is + // scalar — ORT skips vec4 for NCHW too. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv2dGemmWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + bias.buffer, + bias.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvParams)}, + }); + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, gemm_grid.x, gemm_grid.y); + } else { + // Direct conv (grouped/depthwise); vec4-over-IC when icpg%4==0 (register- + // packed, not coalesced — NCHW's channel stride isn't contiguous). + auto constants = utils::make_grid_constants(direct_grid); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + use_vec4 ? kConv2dVec4WGSL : kConv2dWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + bias.buffer, + bias.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvParams)}, + }, + constants.data(), + constants.size()); + graph.add_dispatch_2d( + bundle.pipeline, + bundle.bind_group, + direct_grid.count_x, + direct_grid.count_y); + } wgpuBufferRelease(uniform_buffer); if (bias.owned_dummy != nullptr) { diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl new file mode 100644 index 00000000000..ecae242dbad --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl @@ -0,0 +1,116 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +// Standard (groups==1) conv2d as an implicit-im2col tiled GEMM, reusing the +// linear_fp32_tiled skeleton (TILE=32, RPT=4, shared-mem tiles). M=OC, N=B*OH*OW, +// K=IC*KH*KW. A=weight [OC,K] (contiguous OIHW); B=input sampled on the fly (im2col), +// out-of-range -> 0.0 implements padding. bias is per-ROW (OC). Output written NCHW. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn gK() -> u32 { return params.IC * params.KH * params.KW; } +fn gN() -> u32 { return params.B * params.OH * params.OW; } + +fn read_a(m: u32, kk: u32) -> f32 { // weight[OC, IC*KH*KW] row-major + if (m < params.OC && kk < gK()) { + return weight[m * gK() + kk]; + } + return 0.0; +} + +fn read_b(kk: u32, n: u32) -> f32 { // im2col sample of input[B,IC,IH,IW] + if (kk >= gK() || n >= gN()) { + return 0.0; + } + let ohw = params.OH * params.OW; + let b = n / ohw; + let sp = n % ohw; + let oh = sp / params.OW; + let ow = sp % params.OW; + let khw = params.KH * params.KW; + let ic = kk / khw; + let r = kk % khw; + let kh = r / params.KW; + let kw = r % params.KW; + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (ih < 0 || ih >= i32(params.IH) || iw < 0 || iw >= i32(params.IW)) { + return 0.0; // padding + } + return input[((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw)]; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let M = params.OC; + let N = gN(); + let K = gK(); + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0; + } + } + + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + a_sub[arow][col] = read_a(tile_row0 + arow, k_start + col); + b_sub[arow][col] = read_b(k_start + arow, tile_col0 + col); + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = acc[ir][ic] + aval * b_sub[k][tile_col + ic]; + } + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let m = tile_row0 + tile_row + ir; // oc + let n = tile_col0 + tile_col + ic; // spatial index + if (m < M && n < N) { + var val = acc[ir][ic]; + if (params.has_bias != 0u) { + val = val + bias[m]; + } + let ohw = params.OH * params.OW; + let b = n / ohw; + let sp = n % ohw; + let oh = sp / params.OW; + let ow = sp % params.OW; + out[((b * params.OC + m) * params.OH + oh) * params.OW + ow] = val; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h new file mode 100644 index 00000000000..b86007f4d6a --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h @@ -0,0 +1,140 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from conv2d_gemm.wgsl - DO NOT EDIT. +// wgsl-sha256: 0cd4ee9c946a7d7029062fcd6cb137fe32db886ce7f07c3514ec58ede0fa695e +inline constexpr const char* kConv2dGemmWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +// Standard (groups==1) conv2d as an implicit-im2col tiled GEMM, reusing the +// linear_fp32_tiled skeleton (TILE=32, RPT=4, shared-mem tiles). M=OC, N=B*OH*OW, +// K=IC*KH*KW. A=weight [OC,K] (contiguous OIHW); B=input sampled on the fly (im2col), +// out-of-range -> 0.0 implements padding. bias is per-ROW (OC). Output written NCHW. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn gK() -> u32 { return params.IC * params.KH * params.KW; } +fn gN() -> u32 { return params.B * params.OH * params.OW; } + +fn read_a(m: u32, kk: u32) -> f32 { // weight[OC, IC*KH*KW] row-major + if (m < params.OC && kk < gK()) { + return weight[m * gK() + kk]; + } + return 0.0; +} + +fn read_b(kk: u32, n: u32) -> f32 { // im2col sample of input[B,IC,IH,IW] + if (kk >= gK() || n >= gN()) { + return 0.0; + } + let ohw = params.OH * params.OW; + let b = n / ohw; + let sp = n % ohw; + let oh = sp / params.OW; + let ow = sp % params.OW; + let khw = params.KH * params.KW; + let ic = kk / khw; + let r = kk % khw; + let kh = r / params.KW; + let kw = r % params.KW; + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (ih < 0 || ih >= i32(params.IH) || iw < 0 || iw >= i32(params.IW)) { + return 0.0; // padding + } + return input[((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw)]; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let M = params.OC; + let N = gN(); + let K = gK(); + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0; + } + } + + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + a_sub[arow][col] = read_a(tile_row0 + arow, k_start + col); + b_sub[arow][col] = read_b(k_start + arow, tile_col0 + col); + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = acc[ir][ic] + aval * b_sub[k][tile_col + ic]; + } + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let m = tile_row0 + tile_row + ir; // oc + let n = tile_col0 + tile_col + ic; // spatial index + if (m < M && n < N) { + var val = acc[ir][ic]; + if (params.has_bias != 0u) { + val = val + bias[m]; + } + let ohw = params.OH * params.OW; + let b = n / ohw; + let sp = n % ohw; + let oh = sp / params.OW; + let ow = sp % params.OW; + out[((b * params.OC + m) * params.OH + oh) * params.OW + ow] = val; + } + } + } +} +)"; + +inline constexpr uint32_t kConv2dGemmWorkgroupSizeX = 8; +inline constexpr uint32_t kConv2dGemmWorkgroupSizeY = 8; +inline constexpr uint32_t kConv2dGemmWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 823c8391aaee842471da071abbf05d8ebea9814d Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:24 -0700 Subject: [PATCH 08/22] [ExecuTorch][WebGPU] et_vk conv2d im2col-GEMM tests Pull Request resolved: https://github.com/pytorch/executorch/pull/20874 Splits the im2col-GEMM routing cases out of the `conv2d` im2col-GEMM op diff into their own test diff, stacked directly above it (op below, tests above). Adds the `grouped_vec4` and `gemm_batched` cases to the `conv2d` suite in `test/op_tests/cases.py`, covering `groups==1` im2col-GEMM routing versus the direct vec4 / scalar kernels. Co-authored-with: Claude Code. ghstack-source-id: 405949758 @exported-using-ghexport Differential Revision: [D111072713](https://our.internmc.facebook.com/intern/diff/D111072713/) --- backends/webgpu/test/op_tests/cases.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index d803bd453a7..9d9c2404dad 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -431,7 +431,11 @@ def _linear_fp32_suite() -> WebGPUTestSuite: @register_op_test("conv2d") def _conv2d_suite() -> WebGPUTestSuite: # DaViT patch-embed / downsample convs + conv_transpose2d (same registration, - # folded by the `transposed` arg). NCHW fp32. + # folded by the `transposed` arg). NCHW fp32. Routing coverage (all vs the + # same fp64 golden): patch_embed/conv3x3_pad1/strided/gemm_batched are + # groups==1 → im2col tiled GEMM (gemm_batched pins the B>1 output write); + # grouped_vec4 (groups=2, icpg=4) → direct vec4 kernel; depthwise (groups=8, + # icpg=1) → direct scalar; transpose2x → conv_transpose2d. return WebGPUTestSuite( module_factory=make_conv, cases=[ @@ -467,6 +471,22 @@ def _conv2d_suite() -> WebGPUTestSuite: }, inputs=(InputSpec(shape=(1, 8, 8, 8), gen=_chw_ramp),), ), + Case( + name="grouped_vec4", + construct={ + "in_ch": 8, + "out_ch": 8, + "kernel": 3, + "padding": 1, + "groups": 2, + }, + inputs=(InputSpec(shape=(1, 8, 8, 8), gen=_chw_ramp),), + ), + Case( + name="gemm_batched", + construct={"in_ch": 8, "out_ch": 16, "kernel": 3, "padding": 1}, + inputs=(InputSpec(shape=(2, 8, 16, 16), gen=_chw_ramp),), + ), Case( name="transpose2x", construct={ From 9df2414ef91a51bd42dee179d0b58bdf588d51e0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:24 -0700 Subject: [PATCH 09/22] [ExecuTorch][WebGPU] et_vk.apply_rotary_emb_hf: HuggingFace rotate-half RoPE runtime op (unblocks Qwen3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20875 **Problem:** the WebGPU runtime registers only `et_vk.apply_rotary_emb` (the interleaved/Meta RoPE convention). HuggingFace-derived models (Qwen3, etc.) export the rotate-half convention, which fuses under VulkanPartitioner into `et_vk.apply_rotary_emb_hf` — an op the runtime graph builder has no handler for, so `WebGPUGraph::build()` throws and the delegate is rejected at load with `DelegateInvalidCompatibility` (et_load error 48). The whole model then fails to load on WebGPU. **Solution:** add the `et_vk.apply_rotary_emb_hf` runtime kernel + handler as a rotate-half sibling of the interleaved op. **Before:** only `apply_rotary_emb` (interleaved) is registered; HF-RoPE models throw at load. **After:** both conventions are handled; HF-RoPE models (Qwen3) load and run. **Implementation:** - New `rotary_embedding_hf.wgsl`: one thread per (i, i+half_dim) pair (rotate-half pairing vs the interleaved even/odd), reading a full `[max_seq, rotary_dim]` freqs table indexed at row `start_pos + s`. Scalar, `wg_size` 64 — structural + optimization parity with the interleaved kernel (RoPE is ~1% of runtime; vec4 is neutral for this elementwise-class op on Apple's scalar ALU). - `RotaryEmbedding.cpp`: `apply_rotary_emb_hf_impl` mirrors the interleaved handler; it parses the extra `start_pos` arg as a build-time Int (baked) or a runtime SymInt (dynamic KV-cache decode) exactly as `Sdpa.cpp` handles `input_pos`, and registers a seq resize hook (xq/xk) plus a start_pos resize hook (dynamic decode). Full rotary only (`rotary_dim == head_dim`); partial-rotary passthrough throws (documented follow-up; Qwen3 uses full RoPE). Mirrors Vulkan `et_vk.apply_rotary_emb_hf` (`backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp`). - Registers `et_vk.apply_rotary_emb_hf.default`. **Constraints:** full rotary only for now; scalar one-thread-per-pair, kept at parity with the interleaved sibling rather than vec4 (neutral for RoPE per the closed vec4 sweep). Co-authored-with: Claude Code. ghstack-source-id: 405949770 @exported-using-ghexport Differential Revision: [D111009173](https://our.internmc.facebook.com/intern/diff/D111009173/) --- .../runtime/ops/rope/RotaryEmbedding.cpp | 375 ++++++++++++++++++ .../runtime/ops/rope/rotary_embedding_hf.wgsl | 49 +++ .../ops/rope/rotary_embedding_hf_wgsl.h | 73 ++++ 3 files changed, 497 insertions(+) create mode 100644 backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl create mode 100644 backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 76eafb0e738..af1783496d2 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -378,10 +379,384 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { // pipeline_q/pipeline_k owned by their dispatches; graph dtor frees. } +// HuggingFace rotate-half RoPE (Qwen3 etc.). Structural sibling of the +// interleaved path above (same one-thread-per-pair scalar dispatch, wg_size, +// and resize hook); differs only in element pairing (i with i+half_dim vs +// even/odd), a full [max_seq, rotary_dim] freqs table, and a start_pos offset. +// Mirrors Vulkan's et_vk.apply_rotary_emb_hf +// (backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp:211). + +// Uniform layout matching the HF WGSL Params struct (32 bytes). +struct RotaryHfParams { + uint32_t n_heads; + uint32_t seq; + uint32_t head_dim; + uint32_t half_dim; + uint32_t num_pairs; + uint32_t rotary_dim; + uint32_t start_pos; + uint32_t _pad0; +}; +static_assert(sizeof(RotaryHfParams) == 32, "RotaryHfParams must be 32 bytes"); + +RopeDispatch add_rope_hf_dispatch( + WebGPUGraph& graph, + WGPUDevice device, + uint32_t wg_size, + const WebGPUTensor& x, + const WebGPUTensor& out, + const WebGPUTensor& freqs_cos, + const WebGPUTensor& freqs_sin, + uint32_t n_heads, + uint32_t seq, + uint32_t head_dim, + uint32_t half_dim, + uint32_t rotary_dim, + uint32_t start_pos, + uint32_t workgroup_count) { + const uint32_t num_pairs = + static_cast(utils::numel_of(out.dims) / 2u); + + RotaryHfParams params = {}; + params.n_heads = n_heads; + params.seq = seq; + params.head_dim = head_dim; + params.half_dim = half_dim; + params.num_pairs = num_pairs; + params.rotary_dim = rotary_dim; + params.start_pos = start_pos; + + WGPUBufferDescriptor uniform_desc = {}; + uniform_desc.size = sizeof(RotaryHfParams); + uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + uniform_desc.mappedAtCreation = true; + WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); + void* mapped = + wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(RotaryHfParams)); + std::memcpy(mapped, ¶ms, sizeof(RotaryHfParams)); + wgpuBufferUnmap(uniform_buffer); + graph.add_uniform_buffer_bytes(sizeof(RotaryHfParams)); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kRotaryEmbeddingHfWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_cos.buffer, + freqs_cos.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_sin.buffer, + freqs_sin.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(RotaryHfParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_index = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count, + "apply_rotary_emb_hf"}); + + graph.own_uniform_buffer(uniform_buffer); + return {uniform_buffer, dispatch_index}; +} + +// Resize hook body: recompute S/num_pairs + (dynamic) start_pos for both +// dispatches; out follows xq/xk. Fires on xq/xk seq resize and, when start_pos +// is a runtime SymInt (KV-cache decode), on each start_pos change; idempotent. +void resize_rope_hf( + WebGPUGraph& g, + int xq_id, + int xk_id, + int xq_out_id, + int xk_out_id, + int start_pos_id, + bool dynamic_pos, + uint32_t baked_start_pos, + uint32_t n_heads_q, + uint32_t n_heads_k, + uint32_t head_dim, + uint32_t half_dim, + uint32_t rotary_dim, + uint32_t wg_size, + size_t q_idx, + size_t k_idx, + WGPUBuffer q_ubuf, + WGPUBuffer k_ubuf) { + const auto& qd = g.cur_dims(xq_id); + const auto& kd = g.cur_dims(xk_id); + if (qd.size() < 3 || kd.size() < 3) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); + } + const uint32_t s = static_cast(qd[qd.size() - 3]); + const uint64_t qn = utils::numel_of(qd); + const uint64_t kn = utils::numel_of(kd); + if (static_cast(kd[kd.size() - 3]) != s) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q and k seq lengths differ"); + } + uint32_t start_pos = baked_start_pos; + if (dynamic_pos) { + const int32_t pos = g.read_symint(start_pos_id); + if (pos < 0) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos must be non-negative"); + } + start_pos = static_cast(pos); + } + RotaryHfParams pq = {}; + pq.n_heads = n_heads_q; + pq.seq = s; + pq.head_dim = head_dim; + pq.half_dim = half_dim; + pq.num_pairs = static_cast(qn / 2u); + pq.rotary_dim = rotary_dim; + pq.start_pos = start_pos; + RotaryHfParams pk = pq; + pk.n_heads = n_heads_k; + pk.num_pairs = static_cast(kn / 2u); + wgpuQueueWriteBuffer(g.queue(), q_ubuf, 0, &pq, sizeof(pq)); + wgpuQueueWriteBuffer(g.queue(), k_ubuf, 0, &pk, sizeof(pk)); + g.dispatch_at(q_idx).workgroup_count_x = utils::compute_1d_workgroup_count( + g.device(), + static_cast(qn / 2u), + wg_size, + "apply_rotary_emb_hf(resize)"); + g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( + g.device(), + static_cast(kn / 2u), + wg_size, + "apply_rotary_emb_hf(resize)"); + g.set_cur_dims(xq_out_id, qd); + g.set_cur_dims(xk_out_id, kd); +} + +// args: [xq, xk, freqs_cos, freqs_sin, start_pos, out_list(ValueList[xq_out, +// xk_out])]. freqs is the FULL [max_seq, rotary_dim] table (start_pos offsets +// into it), unlike the pre-sliced interleaved freqs. +void apply_rotary_emb_hf_impl( + WebGPUGraph& graph, + const std::vector& args) { + const int xq_id = args.at(0); + const int xk_id = args.at(1); + const int freqs_cos_id = args.at(2); + const int freqs_sin_id = args.at(3); + const int start_pos_id = args.at(4); + + const std::vector& out_list = graph.get_value_list(args.at(5)); + if (out_list.size() != 2) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: expected an output ValueList of size 2"); + } + + WGPUDevice device = graph.device(); + + const auto& xq = graph.get_tensor(xq_id); + const auto& xk = graph.get_tensor(xk_id); + const auto& freqs_cos = graph.get_tensor(freqs_cos_id); + const auto& freqs_sin = graph.get_tensor(freqs_sin_id); + const auto& xq_out = graph.get_tensor(out_list[0]); + const auto& xk_out = graph.get_tensor(out_list[1]); + + // Shape contract: xq/xk (B,S,n_heads,head_dim), freqs (max_seq, rotary_dim). + if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { + throw std::runtime_error("WebGPU apply_rotary_emb_hf: malformed dims"); + } + const uint32_t head_dim = static_cast(xq.dims.back()); + const uint32_t seq = static_cast(xq.dims[xq.dims.size() - 3]); + const uint32_t n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); + const uint32_t n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); + const uint32_t seq_k = static_cast(xk.dims[xk.dims.size() - 3]); + const uint32_t max_seq = + static_cast(freqs_cos.dims[freqs_cos.dims.size() - 2]); + const uint32_t rotary_dim = static_cast(freqs_cos.dims.back()); + + if (head_dim == 0 || head_dim % 2 != 0) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: head_dim must be a nonzero multiple of 2"); + } + if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: xq/xk head_dim and seq must match"); + } + // Full rotary only (rotary_dim == head_dim); partial-rotary passthrough is a + // documented follow-up (Qwen3 uses full RoPE). Throw rather than mis-rotate. + if (rotary_dim != head_dim) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: partial rotary (rotary_dim != head_dim) " + "not supported"); + } + if (freqs_cos.dims != freqs_sin.dims) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: freqs_cos and freqs_sin shapes differ"); + } + if (max_seq < seq) { + throw std::runtime_error("WebGPU apply_rotary_emb_hf: freqs max_seq < seq"); + } + + if (xq.buffer == nullptr || xk.buffer == nullptr || + freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || + xq_out.buffer == nullptr || xk_out.buffer == nullptr) { + throw std::runtime_error("WebGPU apply_rotary_emb_hf: null buffer binding"); + } + + const uint32_t half_dim = rotary_dim / 2u; + + // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); + // mirrors sdpa's input_pos handling. + int64_t start_pos = 0; + const auto start_pos_type = graph.get_value_type(start_pos_id); + const bool dynamic_pos = start_pos_type == WebGPUGraph::ValueType::SymInt; + if (dynamic_pos) { + start_pos = graph.read_symint(start_pos_id); // build placeholder (e.g. 0) + } else if (start_pos_type == WebGPUGraph::ValueType::Int) { + start_pos = graph.get_int(start_pos_id); + } else { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos must be Int or SymInt"); + } + if (start_pos < 0) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); + } + + // All tensors are fp32; output shapes equal their inputs. + const uint64_t xq_numel = utils::numel_of(xq.dims); + const uint64_t xk_numel = utils::numel_of(xk.dims); + const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); + if (freqs_numel != static_cast(max_seq) * rotary_dim || + xq.nbytes != xq_numel * sizeof(float) || + xk.nbytes != xk_numel * sizeof(float) || + freqs_cos.nbytes != freqs_numel * sizeof(float) || + freqs_sin.nbytes != freqs_numel * sizeof(float) || + xq_out.nbytes != xq_numel * sizeof(float) || + xk_out.nbytes != xk_numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: dtype/byte-size mismatch (all fp32) or " + "freqs shape != [max_seq, rotary_dim]"); + } + + if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: pair count exceeds uint32 dispatch range"); + } + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kRotaryEmbeddingHfWorkgroupSizeX); + // Validate both dispatches before any GPU-object alloc (no leak on throw). + const uint32_t xq_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(xq_numel / 2u), + wg_size, + "apply_rotary_emb_hf"); + const uint32_t xk_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(xk_numel / 2u), + wg_size, + "apply_rotary_emb_hf"); + + RopeDispatch q_disp = add_rope_hf_dispatch( + graph, + device, + wg_size, + xq, + xq_out, + freqs_cos, + freqs_sin, + n_heads_q, + seq, + head_dim, + half_dim, + rotary_dim, + static_cast(start_pos), + xq_wgc); + RopeDispatch k_disp = add_rope_hf_dispatch( + graph, + device, + wg_size, + xk, + xk_out, + freqs_cos, + freqs_sin, + n_heads_k, + seq, + head_dim, + half_dim, + rotary_dim, + static_cast(start_pos), + xk_wgc); + WGPUBuffer q_ubuf = q_disp.uniform; + WGPUBuffer k_ubuf = k_disp.uniform; + const size_t q_idx = q_disp.dispatch_index; + const size_t k_idx = k_disp.dispatch_index; + + const int xq_out_id = out_list[0]; + const int xk_out_id = out_list[1]; + const uint32_t baked_start_pos = static_cast(start_pos); + auto rope_hook = [xq_id, + xk_id, + xq_out_id, + xk_out_id, + start_pos_id, + dynamic_pos, + baked_start_pos, + n_heads_q, + n_heads_k, + head_dim, + half_dim, + rotary_dim, + wg_size, + q_idx, + k_idx, + q_ubuf, + k_ubuf](WebGPUGraph& g) { + resize_rope_hf( + g, + xq_id, + xk_id, + xq_out_id, + xk_out_id, + start_pos_id, + dynamic_pos, + baked_start_pos, + n_heads_q, + n_heads_k, + head_dim, + half_dim, + rotary_dim, + wg_size, + q_idx, + k_idx, + q_ubuf, + k_ubuf); + }; + graph.add_tensor_resize_hook(xq_id, rope_hook); + graph.add_tensor_resize_hook(xk_id, rope_hook); + // Dynamic decode: re-fire when the runtime start_pos SymInt changes. + if (dynamic_pos) { + graph.add_resize_hook(start_pos_id, rope_hook); + } +} + } // namespace WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(et_vk.apply_rotary_emb.default, apply_rotary_emb_impl); + WEBGPU_REGISTER_OP( + et_vk.apply_rotary_emb_hf.default, apply_rotary_emb_hf_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl new file mode 100644 index 00000000000..14a6853afa3 --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl @@ -0,0 +1,49 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_freqs_cos: array; +@group(0) @binding(3) var t_freqs_sin: array; + +struct Params { + n_heads: u32, + seq: u32, + head_dim: u32, + half_dim: u32, + num_pairs: u32, + rotary_dim: u32, + start_pos: u32, + _pad0: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +// One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated +// halves) indexed at row (start_pos + s); only the first-half column is read. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let pair = gid.x; + if (pair >= params.num_pairs) { + return; + } + let half_dim = params.half_dim; + let pair_i = pair % half_dim; + let t1 = pair / half_dim; + let head = t1 % params.n_heads; + let t2 = t1 / params.n_heads; + let s = t2 % params.seq; + let b = t2 / params.seq; + + let head_base = + ((b * params.seq + s) * params.n_heads + head) * params.head_dim; + let a_idx = head_base + pair_i; + let b_idx = head_base + pair_i + half_dim; + let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + + let c = t_freqs_cos[freqs_idx]; + let si = t_freqs_sin[freqs_idx]; + let x_a = t_in[a_idx]; + let x_b = t_in[b_idx]; + t_out[a_idx] = x_a * c - x_b * si; + t_out[b_idx] = x_b * c + x_a * si; +} diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h new file mode 100644 index 00000000000..191ec710e66 --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from rotary_embedding_hf.wgsl - DO NOT EDIT. +// wgsl-sha256: 5ba8d45925f00f12af17bf3092a1af9513a9e501c5c35e6b0d48cfb3dac7b5d6 +inline constexpr const char* kRotaryEmbeddingHfWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_freqs_cos: array; +@group(0) @binding(3) var t_freqs_sin: array; + +struct Params { + n_heads: u32, + seq: u32, + head_dim: u32, + half_dim: u32, + num_pairs: u32, + rotary_dim: u32, + start_pos: u32, + _pad0: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +// One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated +// halves) indexed at row (start_pos + s); only the first-half column is read. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let pair = gid.x; + if (pair >= params.num_pairs) { + return; + } + let half_dim = params.half_dim; + let pair_i = pair % half_dim; + let t1 = pair / half_dim; + let head = t1 % params.n_heads; + let t2 = t1 / params.n_heads; + let s = t2 % params.seq; + let b = t2 / params.seq; + + let head_base = + ((b * params.seq + s) * params.n_heads + head) * params.head_dim; + let a_idx = head_base + pair_i; + let b_idx = head_base + pair_i + half_dim; + let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + + let c = t_freqs_cos[freqs_idx]; + let si = t_freqs_sin[freqs_idx]; + let x_a = t_in[a_idx]; + let x_b = t_in[b_idx]; + t_out[a_idx] = x_a * c - x_b * si; + t_out[b_idx] = x_b * c + x_a * si; +} +)"; + +inline constexpr uint32_t kRotaryEmbeddingHfWorkgroupSizeX = 64; +inline constexpr uint32_t kRotaryEmbeddingHfWorkgroupSizeY = 1; +inline constexpr uint32_t kRotaryEmbeddingHfWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 4752939ff5e36ae82b7603d33e2fa85389b7104d Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:25 -0700 Subject: [PATCH 10/22] [ExecuTorch][WebGPU] apply_rotary_emb_hf op tests Pull Request resolved: https://github.com/pytorch/executorch/pull/20876 Splits the `apply_rotary_emb_hf` op tests into their own diff, stacked directly above the op (op below, tests above). Adds `test/ops/test_rope_hf.py`, the per-op export test for the HuggingFace rotate-half RoPE runtime op. Co-authored-with: Claude Code. ghstack-source-id: 405949767 @exported-using-ghexport Differential Revision: [D111072714](https://our.internmc.facebook.com/intern/diff/D111072714/) --- backends/webgpu/test/ops/test_rope_hf.py | 136 +++++++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 backends/webgpu/test/ops/test_rope_hf.py diff --git a/backends/webgpu/test/ops/test_rope_hf.py b/backends/webgpu/test/ops/test_rope_hf.py new file mode 100644 index 00000000000..6e26fe53b17 --- /dev/null +++ b/backends/webgpu/test/ops/test_rope_hf.py @@ -0,0 +1,136 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""HuggingFace rotate-half rotary positional embedding +(`et_vk.apply_rotary_emb_hf`) export + goldens for the WebGPU backend. + +Qwen3 (and other HF-derived models) use the rotate-half RoPE convention, which +fuses under VulkanPartitioner into `et_vk.apply_rotary_emb_hf.default` (a full +[max_seq, rotary_dim] freqs table + a start_pos offset, two outputs xq_out, +xk_out as a ValueList). This is the counterpart of test_rope.py for the +interleaved (Llama) convention. Full rotary only (rotary_dim == head_dim), which +is what Qwen3 uses; the runtime rejects partial rotary. + +Inputs are deterministic /16 ramps so the native binary reconstructs them +bit-for-bit; the two torch-computed goldens are written for the native binary to +compare (it has no ATen). +""" + +import unittest +from collections import namedtuple + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + +import torch +from executorch.backends.vulkan import VulkanPartitioner +from executorch.examples.models.llama.rope import hf_apply_rotary_emb +from executorch.exir import to_edge_transform_and_lower + +# B batch, S tokens, NH query heads, NKV kv heads (NH != NKV so the two outputs +# are distinguishable by numel), HD head dim (even; full rotary, rotary_dim==HD). +Shape = namedtuple("Shape", ["name", "b", "s", "nh", "nkv", "hd"]) +SHAPES = [ + Shape("multi", 1, 5, 8, 2, 64), + # Single-token decode at a Qwen3-0.6B-like head config (GQA 16:8, head_dim + # 128) so the seq=1 / batch decompositions are covered at decode too. + Shape("decode", 1, 1, 16, 8, 128), +] + + +class HfRope(torch.nn.Module): + # unsqueeze_dim=1: freqs [S, HD] -> [S, 1, HD] broadcasts over (B, NH) of the + # [B, S, NH, HD] q/k, matching the WebGPU kernel + HfRotaryEmbeddingPattern. + def forward(self, xq, xk, freqs_cos, freqs_sin): + return hf_apply_rotary_emb(xq, xk, freqs_cos, freqs_sin, unsqueeze_dim=1) + + +def _ramp(numel: int, mod: int, off: int) -> torch.Tensor: + # ((i % mod) - off) / 16: exact in fp32, matches test_webgpu_native.cpp. + idx = torch.arange(numel, dtype=torch.int64) + return ((idx % mod) - off).to(torch.float32) / 16.0 + + +def _inputs( + shape: Shape, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + xq = _ramp(shape.b * shape.s * shape.nh * shape.hd, 17, 8).reshape( + shape.b, shape.s, shape.nh, shape.hd + ) + xk = _ramp(shape.b * shape.s * shape.nkv * shape.hd, 13, 6).reshape( + shape.b, shape.s, shape.nkv, shape.hd + ) + # HF freqs are the FULL rotary_dim (== head_dim) table, not head_dim/2. + freqs_cos = _ramp(shape.s * shape.hd, 11, 5).reshape(shape.s, shape.hd) + freqs_sin = _ramp(shape.s * shape.hd, 7, 3).reshape(shape.s, shape.hd) + return xq, xk, freqs_cos, freqs_sin + + +def _golden( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + # Reference = the registered et_vk op the kernel implements (start_pos=0). + return torch.ops.et_vk.apply_rotary_emb_hf.default(xq, xk, freqs_cos, freqs_sin, 0) + + +def _export(inputs): + ep = torch.export.export(HfRope().eval(), inputs) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class TestRopeHf(unittest.TestCase): + def test_export_delegates(self) -> None: + for shape in SHAPES: + with self.subTest(shape=shape.name): + et = _export(_inputs(shape)) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue( + found, + "Expected a VulkanBackend delegate (apply_rotary_emb_hf " "fusion)", + ) + + def test_golden_matches_eager(self) -> None: + # The et_vk golden must equal the real HF rotate-half apply_rotary_emb, + # so a buggy golden can't fake-pass the native kernel. Run at both shapes + # so the S=1 decode position indexing is covered. + for shape in SHAPES: + with self.subTest(shape=shape.name): + xq, xk, fc, fs = _inputs(shape) + gq, gk = _golden(xq, xk, fc, fs) + eq, ek = hf_apply_rotary_emb(xq, xk, fc, fs, unsqueeze_dim=1) + torch.testing.assert_close(gq, eq, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(gk, ek, atol=1e-5, rtol=1e-5) + + +def export_rope_hf_model( + pte_path: str, xq_golden_path: str, xk_golden_path: str, shape_name: str = "multi" +) -> None: + """Write the apply_rotary_emb_hf .pte + the xq_out and xk_out torch goldens + (raw LE fp32). Inputs are /16 ramps reconstructed in the native test.""" + shape = next(s for s in SHAPES if s.name == shape_name) + xq, xk, fc, fs = _inputs(shape) + gq, gk = _golden(xq, xk, fc, fs) + et = _export((xq, xk, fc, fs)) + with open(pte_path, "wb") as f: + f.write(et.buffer) + gq.detach().numpy().astype(" Date: Wed, 22 Jul 2026 22:05:25 -0700 Subject: [PATCH 11/22] [ExecuTorch][WebGPU] Tests for aten.sub.Tensor broadcast Pull Request resolved: https://github.com/pytorch/executorch/pull/20986 **Tests for `aten.sub.Tensor` broadcast** Adds op-test coverage for `aten.sub.Tensor`, stacked directly on the sub op diff (op below, tests above). `test/ops/test_sub.py` provides `SubModule` + `CONFIGS` (same-shape, the middle/spatial broadcast `[N,C,H,W] - [N,C,1,1]`, and an alpha != 1 case) plus the export-delegation smoke test; `test/op_tests/cases.py` registers the matching numeric suite (fp64 torch golden on Dawn, mirroring `_mul_suite`), with `alpha` baked into the `.pte` as a construct constant. Co-authored-with: Claude Code. ghstack-source-id: 405949771 @exported-using-ghexport Differential Revision: [D112378930](https://our.internmc.facebook.com/intern/diff/D112378930/) --- backends/webgpu/test/op_tests/cases.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 9d9c2404dad..9d453832ad8 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -771,3 +771,23 @@ def _relu_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_sub import ( + CONFIGS as _SUB_CONFIGS, + SubModule, +) + + +@register_op_test("sub") +def _sub_suite() -> WebGPUTestSuite: + # Full numeric coverage incl. the spatial broadcast + alpha (binary_sub.wgsl + # over a TensorMeta UBO); fp64 golden. Mirrors _mul_suite. alpha is a + # construct kwarg baked into the .pte, never a serialized input. + return WebGPUTestSuite( + module_factory=lambda alpha=1.0: SubModule(alpha), + cases=[ + Case(name=name, construct={"alpha": alpha}, inputs=(sa, sb)) + for name, (sa, sb, alpha) in _SUB_CONFIGS.items() + ], + ) From 1bc407dcd8e9e20642b1d7214555d79763a5cbbd Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:26 -0700 Subject: [PATCH 12/22] [ExecuTorch][WebGPU] Add aten._to_copy.default with int<->float numeric convert MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20987 **Add `aten._to_copy.default` with int↔float numeric convert** **Problem:** The copy-family ops byte-copied across dtypes, so an int32 -> fp32 cast reinterpreted the raw bits — int32 `2` = `0x2` decodes as the fp32 denormal `2.8e-45` — producing wrong values (and div-by-~0 `inf` downstream). Separately, `aten._to_copy.default` was unregistered, so any delegate containing it failed to load. **Solution:** Add `add_to_copy_node`: same-dtype copies stay a flat byte copy, while int<->float copies run a numeric-convert compute shader (`f32(i32)` / `i32(f32)`). Register `aten._to_copy.default`, and route the dim-order copy ops (`dim_order_ops._clone_dim_order.default` / `._to_dim_order_copy.default`) through the same convert-aware path so an int<->float dim-order copy numeric-converts instead of byte-reinterpreting; `view_copy` / `clone` / `alias_copy` stay on the flat copy. Mirrors Vulkan `ToCopy.cpp` (BlitNode vs the view_convert path). **Implementation:** `runtime/ops/to_copy/{ToCopy.cpp,to_copy.h,to_copy_int_to_float.wgsl,to_copy_float_to_int.wgsl}` provide `add_to_copy_node` and register `aten._to_copy.default`; `runtime/ops/view_copy/ViewCopy.cpp` re-points the two dim-order copy ops at `add_to_copy_node`. One `WEBGPU_SRCS` entry. **Constraints:** 32-bit only (int64 constants are downcast to int32 by the Vulkan serializer); fails loud on any other element width. Co-authored-with: Claude Code. ghstack-source-id: 405949773 @exported-using-ghexport Differential Revision: [D112378932](https://our.internmc.facebook.com/intern/diff/D112378932/) --- .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 163 ++++++++++++++++++ backends/webgpu/runtime/ops/to_copy/to_copy.h | 18 ++ .../ops/to_copy/to_copy_float_to_int.wgsl | 18 ++ .../ops/to_copy/to_copy_float_to_int_wgsl.h | 42 +++++ .../ops/to_copy/to_copy_int_to_float.wgsl | 18 ++ .../ops/to_copy/to_copy_int_to_float_wgsl.h | 42 +++++ .../webgpu/runtime/ops/view_copy/ViewCopy.cpp | 14 +- 7 files changed, 313 insertions(+), 2 deletions(-) create mode 100644 backends/webgpu/runtime/ops/to_copy/ToCopy.cpp create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp new file mode 100644 index 00000000000..05f7ff6fe24 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -0,0 +1,163 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform buffer layout matching the WGSL Params struct; 16-byte aligned. +struct ConvertParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; + +// Elementwise int32<->fp32 convert; mirrors Vulkan add_view_copy_convert_node. +void add_convert_op( + WebGPUGraph& graph, + int in_id, + int out_id, + const char* wgsl_source, + uint32_t wg_size_x, + const char* op_name) { + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error(std::string(op_name) + ": null buffer binding"); + } + + // 32-bit only: int64 consts are downcast to int32 by the Vulkan serializer. + if (in_tensor.elem_size != 4 || out_tensor.elem_size != 4) { + throw std::runtime_error( + std::string(op_name) + ": only 32-bit int<->float convert supported"); + } + if (in_tensor.nbytes != out_tensor.nbytes) { + throw std::runtime_error(std::string(op_name) + ": numel mismatch"); + } + + uint32_t num_elements = static_cast(out_tensor.nbytes / 4); + + uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); + uint32_t workgroup_count = + utils::compute_1d_workgroup_count(device, num_elements, wg_size, op_name); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl_source, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + // Dynamic shapes: recompute num_elements/dispatch for the live shape. + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(d); + g.set_cur_dims(out_id, d); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy(resize)"); + }); + + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +void to_copy_impl(WebGPUGraph& graph, const std::vector& args) { + // aten._to_copy.default args: [self, ...kwargs, out]; out = last value id. + add_to_copy_node(graph, args.at(0), args.at(args.size() - 1)); +} + +} // namespace + +void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + + // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain. + if (in_tensor.is_int == out_tensor.is_int && + in_tensor.elem_size == out_tensor.elem_size) { + add_flat_copy(graph, in_id, out_id); + return; + } + + // int<->float = numeric convert (mirrors Vulkan add_view_copy_convert_node). + if (in_tensor.is_int && !out_tensor.is_int) { + add_convert_op( + graph, + in_id, + out_id, + kToCopyIntToFloatWGSL, + kToCopyIntToFloatWorkgroupSizeX, + "to_copy_int_to_float"); + } else { + add_convert_op( + graph, + in_id, + out_id, + kToCopyFloatToIntWGSL, + kToCopyFloatToIntWorkgroupSizeX, + "to_copy_float_to_int"); + } +} + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten._to_copy.default, to_copy_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy.h b/backends/webgpu/runtime/ops/to_copy/to_copy.h new file mode 100644 index 00000000000..714e9a5224e --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy.h @@ -0,0 +1,18 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// Copy: flat byte for same dtype, numeric convert for int<->float (Vulkan). +void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl new file mode 100644 index 00000000000..3eb6cb44595 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl @@ -0,0 +1,18 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + output[idx] = i32(input[idx]); +} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h new file mode 100644 index 00000000000..e7e0391dd13 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h @@ -0,0 +1,42 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from to_copy_float_to_int.wgsl - DO NOT EDIT. +// wgsl-sha256: c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f +inline constexpr const char* kToCopyFloatToIntWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + output[idx] = i32(input[idx]); +} +)"; + +inline constexpr uint32_t kToCopyFloatToIntWorkgroupSizeX = 64; +inline constexpr uint32_t kToCopyFloatToIntWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyFloatToIntWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl new file mode 100644 index 00000000000..87affe78290 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl @@ -0,0 +1,18 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + output[idx] = f32(input[idx]); +} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h new file mode 100644 index 00000000000..fd18700f17c --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h @@ -0,0 +1,42 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from to_copy_int_to_float.wgsl - DO NOT EDIT. +// wgsl-sha256: e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195 +inline constexpr const char* kToCopyIntToFloatWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + output[idx] = f32(input[idx]); +} +)"; + +inline constexpr uint32_t kToCopyIntToFloatWorkgroupSizeX = 64; +inline constexpr uint32_t kToCopyIntToFloatWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyIntToFloatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp index 7b3b44eb86f..e718cae4478 100644 --- a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp +++ b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -98,14 +99,23 @@ void clone_impl(WebGPUGraph& graph, const std::vector& args) { add_flat_copy(graph, args.at(0), args.at(args.size() - 1)); } +// dim_order copies may change dtype -> convert-aware add_to_copy_node. +void clone_dim_order_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [self, non_blocking?, dim_order?, out]; out = last value id. + add_to_copy_node(graph, args.at(0), args.at(args.size() - 1)); +} + } // namespace WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.view_copy.default, view_copy_impl); WEBGPU_REGISTER_OP(aten.clone.default, clone_impl); WEBGPU_REGISTER_OP(aten.alias_copy.default, clone_impl); - // _clone_dim_order ignores dim_order (AOT pass elides via shape+dtype). - WEBGPU_REGISTER_OP(dim_order_ops._clone_dim_order.default, clone_impl); + // dim_order_ops._clone_dim_order.default is owned by DimOrder.cpp + // (numel-preserving flat copy); do NOT re-register it here (first-wins + // OperatorRegistry would nondeterministically shadow that impl). + WEBGPU_REGISTER_OP( + dim_order_ops._to_dim_order_copy.default, clone_dim_order_impl); } } // namespace executorch::backends::webgpu From a3ba9d02b3b2668925598c852aeda077554c4153 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:26 -0700 Subject: [PATCH 13/22] [ExecuTorch][WebGPU] Tests for aten._to_copy.default int<->float convert MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20988 **Tests for `aten._to_copy.default` int↔float convert** Adds `test/ops/test_to_copy.py`, stacked directly on the to_copy op diff (op below, tests above). Two export-delegation smoke tests (mirroring `test_view_copy.py`): int32 -> fp32 (input int `[1, 2, 3]`, the numeric-convert path) and fp32 -> fp32 (same-dtype flat copy, `copy=True` so the op is not elided). The int -> float value correctness — `[1, 2, 3]` -> `[1.0, 2.0, 3.0]`, NOT the bit-reinterpretation `0x1 -> 1.4e-45` — is checked by the lvp golden. Co-authored-with: Claude Code. ghstack-source-id: 405949769 @exported-using-ghexport Differential Revision: [D112378931](https://our.internmc.facebook.com/intern/diff/D112378931/) --- backends/webgpu/test/ops/test_to_copy.py | 65 ++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 backends/webgpu/test/ops/test_to_copy.py diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py new file mode 100644 index 00000000000..baf068a3438 --- /dev/null +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -0,0 +1,65 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten._to_copy.default` modules for the WebGPU op-test framework. + +`ToCopyIntToFloatModule` / `ToCopyFloatModule` drive the export-delegation smoke +test (mirroring `test_view_copy.py`). The int32 -> fp32 numeric convert — input +int `[1, 2, 3]` -> float `[1.0, 2.0, 3.0]`, NOT the byte-reinterpretation +`0x1 -> 1.4e-45` — and the fp32 same-dtype passthrough are value-checked by the +lvp golden (yolo11n / Depth-Anything-V2). +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class ToCopyIntToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + # int32 -> fp32 dtype cast (the numeric-convert path). + return x.to(torch.float32) + + +class ToCopyFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Same-dtype copy (flat byte-copy path); copy=True keeps the op from + # being elided as a no-op. + return x.to(torch.float32, copy=True) + + +def _export(model: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(model.eval(), (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +class ToCopyTest(unittest.TestCase): + def test_int_to_float_delegates(self) -> None: + x = torch.tensor([1, 2, 3], dtype=torch.int32) + et = _export(ToCopyIntToFloatModule(), x) + self.assertTrue( + _delegated(et), "Expected a VulkanBackend delegate (to_copy int->float)" + ) + + def test_float_passthrough_delegates(self) -> None: + x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + et = _export(ToCopyFloatModule(), x) + self.assertTrue( + _delegated(et), "Expected a VulkanBackend delegate (to_copy float->float)" + ) From 9181a9926f1a9c022bc90b637719fc5082be8011 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:26 -0700 Subject: [PATCH 14/22] [ExecuTorch][WebGPU] Add leaky_relu op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20989 **Add `aten.leaky_relu.default` to the WebGPU backend** — the SRVGGNetCompact body activation in Real-ESRGAN x4plus super-resolution, so that model can fully delegate to the GPU. **Problem**: The WebGPU delegate had no `leaky_relu` handler, so a model using it could not produce a fully-delegated `.pte`. **Solution**: A scalar-parameter elementwise fp32 kernel computing `x >= 0 ? x : negative_slope * x`, with `negative_slope` carried in the uniform and a 2D-spill dispatch for tensors exceeding the 1D workgroup-count limit. **Implementation**: - `runtime/ops/leaky_relu/{LeakyRelu.cpp,leaky_relu.wgsl,leaky_relu_wgsl.h}` registering `aten.leaky_relu.default`; uses `utils::make_compute_pipeline` + `utils::compute_dispatch_grid`. - Mirrors the Vulkan `leaky_relu.default` delegate (scalar-in-uniform, like `pow.Tensor_Scalar`). - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only — throws on non-fp32 or input/output size mismatch (fail-loud, never a silent zero output); no change to existing ops. ghstack-source-id: 405949766 @exported-using-ghexport Differential Revision: [D112417289](https://our.internmc.facebook.com/intern/diff/D112417289/) --- .../runtime/ops/leaky_relu/LeakyRelu.cpp | 100 ++++++++++++++++++ .../runtime/ops/leaky_relu/leaky_relu.wgsl | 24 +++++ .../runtime/ops/leaky_relu/leaky_relu_wgsl.h | 48 +++++++++ 3 files changed, 172 insertions(+) create mode 100644 backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp create mode 100644 backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl create mode 100644 backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h diff --git a/backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp b/backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp new file mode 100644 index 00000000000..4cb23bdfa72 --- /dev/null +++ b/backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp @@ -0,0 +1,100 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct LeakyReluParams { + uint32_t num_elements; + float neg_slope; + uint32_t _pad[2]; +}; + +// aten.leaky_relu.default: fp32 elementwise; mirrors Vulkan leaky_relu. +void leaky_relu_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 2) { + throw std::runtime_error("WebGPU leaky_relu: expected >=2 args"); + } + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + const float neg_slope = + (args.size() >= 3) ? utils::scalar_or(graph, args.at(1), 0.01f) : 0.01f; + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("leaky_relu: null buffer binding"); + } + if (in_tensor.nbytes != out_tensor.nbytes || + out_tensor.nbytes % sizeof(float) != 0) { + throw std::runtime_error("leaky_relu: fp32-only / size mismatch"); + } + const uint32_t num_elements = + utils::checked_u32(out_tensor.nbytes / sizeof(float), "leaky_relu"); + + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, num_elements, kLeakyReluWorkgroupSizeX, "leaky_relu"); + + auto constants = utils::make_grid_constants(grid); + + LeakyReluParams params = {}; + params.num_elements = num_elements; + params.neg_slope = neg_slope; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(LeakyReluParams)); + graph.add_uniform_buffer_bytes(sizeof(LeakyReluParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLeakyReluWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LeakyReluParams)}, + }, + constants.data(), + constants.size()); + + WebGPUDispatch dispatch{}; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; + dispatch.workgroup_count_x = grid.count_x; + dispatch.workgroup_count_y = grid.count_y; + graph.add_dispatch(dispatch); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.leaky_relu.default, leaky_relu_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl new file mode 100644 index 00000000000..317333a7e5b --- /dev/null +++ b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl @@ -0,0 +1,24 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + neg_slope: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// leaky_relu(x) = max(x,0) + neg_slope*min(x,0); fp32 elementwise, 2D-spill. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let x = input[i]; + output[i] = max(x, 0.0) + params.neg_slope * min(x, 0.0); +} diff --git a/backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h new file mode 100644 index 00000000000..48e5b925fb1 --- /dev/null +++ b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h @@ -0,0 +1,48 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from leaky_relu.wgsl - DO NOT EDIT. +// wgsl-sha256: 98add607747a0f47521645f6cb046c500d59709f6641846a0099c5acd3b6a66b +inline constexpr const char* kLeakyReluWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + neg_slope: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// leaky_relu(x) = max(x,0) + neg_slope*min(x,0); fp32 elementwise, 2D-spill. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let x = input[i]; + output[i] = max(x, 0.0) + params.neg_slope * min(x, 0.0); +} +)"; + +inline constexpr uint32_t kLeakyReluWorkgroupSizeX = 256; +inline constexpr uint32_t kLeakyReluWorkgroupSizeY = 1; +inline constexpr uint32_t kLeakyReluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 4c98ebfe4e6417f89f2ea07d2fc5aeefb95fbc88 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:27 -0700 Subject: [PATCH 15/22] [ExecuTorch][WebGPU] leaky_relu op tests Pull Request resolved: https://github.com/pytorch/executorch/pull/20990 **Op-test suite for `aten.leaky_relu.default`** (stacked on the leaky_relu op diff). Adds the declarative op-test entry: `test/ops/test_leaky_relu.py` (`LeakyReluModule`) + a `@register_op_test("leaky_relu")` suite in `test/op_tests/cases.py`. The framework exports each case via `VulkanPartitioner`, computes the fp64 torch golden, and compares the on-GPU output at `atol=rtol=1e-3`. Cases: `default_slope` (4D `[1,16,8,8]`, slope 0.01) + `slope_0_2` (2D `[3,32]`, slope 0.2). The deterministic input spans negatives so the `negative_slope` branch is exercised. ghstack-source-id: 405949781 @exported-using-ghexport Differential Revision: [D112417280](https://our.internmc.facebook.com/intern/diff/D112417280/) --- backends/webgpu/test/op_tests/cases.py | 26 +++++++++++++++++++++ backends/webgpu/test/ops/test_leaky_relu.py | 21 +++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100644 backends/webgpu/test/ops/test_leaky_relu.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 9d453832ad8..6bab48176bb 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -703,6 +703,32 @@ def _upsample_nearest2d_suite() -> WebGPUTestSuite: ) +from executorch.backends.webgpu.test.ops.test_leaky_relu import LeakyReluModule + + +@register_op_test("leaky_relu") +def _leaky_relu_suite() -> WebGPUTestSuite: + # Real-ESRGAN SRVGGNetCompact body activation. The det input spans negatives, + # exercising the negative_slope branch; a 4D and a 2D case. + return WebGPUTestSuite( + module_factory=lambda negative_slope: LeakyReluModule(negative_slope), + cases=[ + Case( + name="default_slope", + construct={"negative_slope": 0.01}, + inputs=(InputSpec(shape=(1, 16, 8, 8), gen=_upsample_det_input),), + ), + Case( + name="slope_0_2", + construct={"negative_slope": 0.2}, + inputs=(InputSpec(shape=(3, 32), gen=_upsample_det_input),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) + + from executorch.backends.webgpu.test.ops.test_max_pool2d import ( _det_input as _maxpool_det_input, MaxPool2dModule, diff --git a/backends/webgpu/test/ops/test_leaky_relu.py b/backends/webgpu/test/ops/test_leaky_relu.py new file mode 100644 index 00000000000..57020749e43 --- /dev/null +++ b/backends/webgpu/test/ops/test_leaky_relu.py @@ -0,0 +1,21 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.leaky_relu.default` module for the WebGPU op-test framework. + +LeakyReLU is the activation in Real-ESRGAN's SRVGGNetCompact body. +""" + +import torch + + +class LeakyReluModule(torch.nn.Module): + def __init__(self, negative_slope: float) -> None: + super().__init__() + self.negative_slope = negative_slope + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.leaky_relu(x, self.negative_slope) From 98f5313499ae283c751cca4eff2103e5df1d42b8 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:27 -0700 Subject: [PATCH 16/22] [ExecuTorch][WebGPU] Add upsample_bilinear2d op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20991 **Add `aten.upsample_bilinear2d.vec` to the WebGPU backend** — the bilinear resize on the Depth-Anything / DPT reassemble+fusion head, which upsamples the ViT patch grid back to image resolution. **Problem**: The WebGPU delegate had only nearest-neighbor upsample, so depth-estimation models using bilinear resize could not fully delegate to the GPU. **Solution**: A 4D NCHW fp32 kernel where each output pixel bilinearly interpolates its four source neighbors. `align_corners` selects the source-index formula (matching ATen `area_pixel_compute_source_index`); output H/W come from the output tensor's own dims. **Implementation**: - `runtime/ops/upsample_bilinear2d/{UpsampleBilinear2d.cpp,upsample_bilinear2d.wgsl,upsample_bilinear2d_wgsl.h}` registering `aten.upsample_bilinear2d.vec`; uses `utils::make_compute_pipeline` + `utils::compute_dispatch_grid` + `utils::make_grid_constants`. - Mirrors the Vulkan `upsample_bilinear2d.vec` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only, 4D in/out with N/C preserved — throws on rank/shape/dtype mismatch (fail-loud, never a silent zero output); no change to existing ops. ghstack-source-id: 405949780 @exported-using-ghexport Differential Revision: [D112417281](https://our.internmc.facebook.com/intern/diff/D112417281/) --- .../UpsampleBilinear2d.cpp | 132 ++++++++++++++++++ .../upsample_bilinear2d.wgsl | 68 +++++++++ .../upsample_bilinear2d_wgsl.h | 92 ++++++++++++ 3 files changed, 292 insertions(+) create mode 100644 backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp create mode 100644 backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl create mode 100644 backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp b/backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp new file mode 100644 index 00000000000..bb4b3d88e12 --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp @@ -0,0 +1,132 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct UpsampleParams { + uint32_t N; + uint32_t C; + uint32_t IH; + uint32_t IW; + uint32_t OH; + uint32_t OW; + uint32_t align_corners; + uint32_t _p0; +}; +static_assert(sizeof(UpsampleParams) == 32, "UpsampleParams must be 32 bytes"); + +// aten.upsample_bilinear2d.vec: bilinear NCHW resize; align_corners src-index. +void upsample_bilinear2d_impl( + WebGPUGraph& graph, + const std::vector& args) { + if (args.size() < 3) { + throw std::runtime_error("WebGPU upsample_bilinear2d: expected >=3 args"); + } + const int in_id = args.at(0); + const int align_corners_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& in = graph.get_tensor(in_id); + const auto& out = graph.get_tensor(out_id); + + if (in.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU upsample_bilinear2d: expected 4D in/out"); + } + + const uint32_t N = static_cast(in.dims[0]); + const uint32_t C = static_cast(in.dims[1]); + const uint32_t IH = static_cast(in.dims[2]); + const uint32_t IW = static_cast(in.dims[3]); + const uint32_t OH = static_cast(out.dims[2]); + const uint32_t OW = static_cast(out.dims[3]); + + if (static_cast(out.dims[0]) != N || + static_cast(out.dims[1]) != C) { + throw std::runtime_error("WebGPU upsample_bilinear2d: N/C mismatch"); + } + if (IH == 0 || IW == 0 || OH == 0 || OW == 0) { + throw std::runtime_error("WebGPU upsample_bilinear2d: zero spatial dim"); + } + + uint64_t out_numel = uint64_t(N) * C * OH * OW; + if (out.nbytes != out_numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU upsample_bilinear2d: fp32-only (byte mismatch)"); + } + if (out_numel > 0xFFFFFFFFull) { + throw std::runtime_error("WebGPU upsample_bilinear2d: output too large"); + } + + const uint32_t align_corners = graph.get_bool(align_corners_id) ? 1u : 0u; + + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + static_cast(out_numel), + kUpsampleBilinear2dWorkgroupSizeX, + "upsample_bilinear2d"); + + UpsampleParams params = {}; + params.N = N; + params.C = C; + params.IH = IH; + params.IW = IW; + params.OH = OH; + params.OW = OW; + params.align_corners = align_corners; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(UpsampleParams)); + graph.add_uniform_buffer_bytes(sizeof(UpsampleParams)); + + auto constants = utils::make_grid_constants(grid); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kUpsampleBilinear2dWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(UpsampleParams)}, + }, + constants.data(), + constants.size()); + + WebGPUDispatch dispatch{}; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; + dispatch.workgroup_count_x = grid.count_x; + dispatch.workgroup_count_y = grid.count_y; + graph.add_dispatch(dispatch); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.upsample_bilinear2d.vec, upsample_bilinear2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl new file mode 100644 index 00000000000..7656f3478a9 --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl @@ -0,0 +1,68 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + align_corners: u32, + _p0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Bilinear NCHW fp32 upsample; src-index matches ATen upsample_bilinear2d. +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = sh - f32(h0); + let lw = sw - f32(w0); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = (n * params.C + c) * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + out[i] = top + (bot - top) * lh; +} diff --git a/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h new file mode 100644 index 00000000000..bbdc7f77c11 --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h @@ -0,0 +1,92 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from upsample_bilinear2d.wgsl - DO NOT EDIT. +// wgsl-sha256: fa8d89f2ec0e316650315ce39a426efec9ed34daef1fb5b7bc19ae9cea64565a +inline constexpr const char* kUpsampleBilinear2dWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + align_corners: u32, + _p0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Bilinear NCHW fp32 upsample; src-index matches ATen upsample_bilinear2d. +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = sh - f32(h0); + let lw = sw - f32(w0); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = (n * params.C + c) * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + out[i] = top + (bot - top) * lh; +} +)"; + +inline constexpr uint32_t kUpsampleBilinear2dWorkgroupSizeX = 256; +inline constexpr uint32_t kUpsampleBilinear2dWorkgroupSizeY = 1; +inline constexpr uint32_t kUpsampleBilinear2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From f47325d744273efaee77e0a1ad82cc118fe6f6bc Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:28 -0700 Subject: [PATCH 17/22] [ExecuTorch][WebGPU] upsample_bilinear2d op tests Pull Request resolved: https://github.com/pytorch/executorch/pull/20992 **Op-test suite for `aten.upsample_bilinear2d.vec`** (stacked on the upsample_bilinear2d op diff). Adds `test/ops/test_upsample_bilinear2d.py` (`UpsampleBilinear2dModule`) + a `@register_op_test("upsample_bilinear2d")` suite in `test/op_tests/cases.py` (5 cases). Covers both `align_corners` branches and a non-integer ratio (5->8) that discriminates the two source-index formulas. ghstack-source-id: 405949782 @exported-using-ghexport Differential Revision: [D112417283](https://our.internmc.facebook.com/intern/diff/D112417283/) --- backends/webgpu/test/op_tests/cases.py | 46 +++++++++++++++++++ .../test/ops/test_upsample_bilinear2d.py | 29 ++++++++++++ 2 files changed, 75 insertions(+) create mode 100644 backends/webgpu/test/ops/test_upsample_bilinear2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 6bab48176bb..d256a154fba 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -729,6 +729,52 @@ def _leaky_relu_suite() -> WebGPUTestSuite: ) +from executorch.backends.webgpu.test.ops.test_upsample_bilinear2d import ( + UpsampleBilinear2dModule, +) + + +@register_op_test("upsample_bilinear2d") +def _upsample_bilinear2d_suite() -> WebGPUTestSuite: + # DPT/Depth-Anything bilinear resize head. Both align_corners branches + + # a non-integer ratio (5->8) that discriminates the two source-index + # formulas (see upsample_bilinear2d.wgsl). + return WebGPUTestSuite( + module_factory=lambda scale_factor, align_corners: UpsampleBilinear2dModule( + scale_factor, align_corners + ), + cases=[ + Case( + name="af_false_2x", + construct={"scale_factor": 2.0, "align_corners": False}, + inputs=(InputSpec(shape=(1, 8, 36, 36), gen=_upsample_det_input),), + ), + Case( + name="af_false_non_2x", + construct={"scale_factor": 1.6, "align_corners": False}, + inputs=(InputSpec(shape=(1, 2, 5, 5), gen=_upsample_det_input),), + ), + Case( + name="af_false_tiny", + construct={"scale_factor": 2.0, "align_corners": False}, + inputs=(InputSpec(shape=(1, 4, 5, 7), gen=_upsample_det_input),), + ), + Case( + name="af_true_2x", + construct={"scale_factor": 2.0, "align_corners": True}, + inputs=(InputSpec(shape=(1, 4, 7, 7), gen=_upsample_det_input),), + ), + Case( + name="af_true_non_2x", + construct={"scale_factor": 1.6, "align_corners": True}, + inputs=(InputSpec(shape=(1, 2, 5, 5), gen=_upsample_det_input),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) + + from executorch.backends.webgpu.test.ops.test_max_pool2d import ( _det_input as _maxpool_det_input, MaxPool2dModule, diff --git a/backends/webgpu/test/ops/test_upsample_bilinear2d.py b/backends/webgpu/test/ops/test_upsample_bilinear2d.py new file mode 100644 index 00000000000..baa34f4b3e9 --- /dev/null +++ b/backends/webgpu/test/ops/test_upsample_bilinear2d.py @@ -0,0 +1,29 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.upsample_bilinear2d.vec` module for the WebGPU op-test framework. + +Bilinear upsample is on the Depth-Anything / DPT reassemble+fusion head +(`F.interpolate(..., mode="bilinear", align_corners=...)`), which resizes the +ViT patch grid back to image resolution. +""" + +import torch + + +class UpsampleBilinear2dModule(torch.nn.Module): + def __init__(self, scale_factor: float, align_corners: bool) -> None: + super().__init__() + self.scale_factor = scale_factor + self.align_corners = align_corners + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.interpolate( + x, + scale_factor=self.scale_factor, + mode="bilinear", + align_corners=self.align_corners, + ) From 0e2fec61c8784197764c2c599e15214157439377 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:28 -0700 Subject: [PATCH 18/22] [ExecuTorch][WebGPU] Add batch_norm op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20993 **Add `aten._native_batch_norm_legit_no_training.default` to the WebGPU backend** — inference batch norm on MODNet's decoder and CNN backbones. **Problem**: The WebGPU delegate had no batch-norm handler, so MODNet (background removal) and other CNN models could not fully delegate to the GPU. **Solution**: A 4D NCHW fp32 kernel applying the per-channel inference affine `y = (x - running_mean) / sqrt(running_var + eps) * weight + bias`. `weight`/`bias` are optional (affine=False → unit scale / zero shift), bound via `utils::make_optional_binding` with a dummy buffer when absent. **Implementation**: - `runtime/ops/batch_norm/{BatchNorm.cpp,batch_norm.wgsl,batch_norm_wgsl.h}` registering `aten._native_batch_norm_legit_no_training.default`; uses `utils::make_compute_pipeline` + `utils::make_optional_binding`. - Multi-output op: reads the `out` entry of the output ValueList (`save_mean`/`save_invstd` unused in inference). - Mirrors the Vulkan `_native_batch_norm_legit_no_training` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only, 4D in/out — throws on rank/shape/dtype mismatch (fail-loud); no change to existing ops. ghstack-source-id: 405949779 @exported-using-ghexport Differential Revision: [D112417288](https://our.internmc.facebook.com/intern/diff/D112417288/) --- .../runtime/ops/batch_norm/BatchNorm.cpp | 172 ++++++++++++++++++ .../runtime/ops/batch_norm/batch_norm.wgsl | 41 +++++ .../runtime/ops/batch_norm/batch_norm_wgsl.h | 65 +++++++ 3 files changed, 278 insertions(+) create mode 100644 backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp create mode 100644 backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl create mode 100644 backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h diff --git a/backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp b/backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp new file mode 100644 index 00000000000..dafde31f1f7 --- /dev/null +++ b/backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp @@ -0,0 +1,172 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct BNParams { + uint32_t num_elements; + uint32_t C; + uint32_t HW; + uint32_t has_weight; + uint32_t has_bias; + float eps; + uint32_t _p0; + uint32_t _p1; +}; +static_assert(sizeof(BNParams) == 32, "BNParams must be 32 bytes"); + +// _native_batch_norm_legit_no_training: per-channel affine from running stats. +void batch_norm_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 8) { + throw std::runtime_error("WebGPU batch_norm: expected >=8 args"); + } + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int mean_id = args.at(3); + const int var_id = args.at(4); + const double eps = graph.get_double(args.at(6)); + const std::vector& outs = graph.get_value_list(args.at(args.size() - 1)); + if (outs.empty()) { + throw std::runtime_error("WebGPU batch_norm: empty out ValueList"); + } + const int out_id = outs.at(0); + + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& out = graph.get_tensor(out_id); + const auto& mean = graph.get_tensor(mean_id); + const auto& var = graph.get_tensor(var_id); + + if (in.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU batch_norm: expected 4D in/out"); + } + const uint32_t N = static_cast(in.dims[0]); + const uint32_t C = static_cast(in.dims[1]); + const uint32_t HW = + static_cast(in.dims[2]) * static_cast(in.dims[3]); + + uint64_t out_numel = uint64_t(N) * C * HW; + if (out.nbytes != out_numel * sizeof(float)) { + throw std::runtime_error("WebGPU batch_norm: fp32-only (byte mismatch)"); + } + if (out_numel > 0xFFFFFFFFull) { + throw std::runtime_error("WebGPU batch_norm: output too large"); + } + + const bool has_weight = + graph.get_value_type(weight_id) == WebGPUGraph::ValueType::Tensor; + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + // Per-channel params are indexed [0, C); guard their byte sizes so an + // undersized buffer can't drive an out-of-bounds shader read. mean/var are + // required; weight/bias only when present. + const uint64_t c_bytes = uint64_t(C) * sizeof(float); + if (mean.nbytes != c_bytes || var.nbytes != c_bytes) { + throw std::runtime_error( + "WebGPU batch_norm: mean/var must be fp32 of length C"); + } + if (has_weight && graph.get_tensor(weight_id).nbytes != c_bytes) { + throw std::runtime_error( + "WebGPU batch_norm: weight must be fp32 of length C"); + } + if (has_bias && graph.get_tensor(bias_id).nbytes != c_bytes) { + throw std::runtime_error( + "WebGPU batch_norm: bias must be fp32 of length C"); + } + + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + static_cast(out_numel), + kBatchNormWorkgroupSizeX, + "batch_norm"); + + BNParams params = {}; + params.num_elements = static_cast(out_numel); + params.C = C; + params.HW = HW; + params.has_weight = has_weight ? 1u : 0u; + params.has_bias = has_bias ? 1u : 0u; + params.eps = static_cast(eps); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(BNParams)); + graph.add_uniform_buffer_bytes(sizeof(BNParams)); + + utils::OptionalBinding wbind = utils::make_optional_binding( + device, + has_weight, + has_weight ? graph.get_tensor(weight_id).buffer : nullptr, + has_weight ? graph.get_tensor(weight_id).nbytes : 0); + utils::OptionalBinding bbind = utils::make_optional_binding( + device, + has_bias, + has_bias ? graph.get_tensor(bias_id).buffer : nullptr, + has_bias ? graph.get_tensor(bias_id).nbytes : 0); + + auto constants = utils::make_grid_constants(grid); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBatchNormWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + wbind.buffer, + wbind.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + bbind.buffer, + bbind.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, mean.buffer, mean.nbytes}, + {5, WGPUBufferBindingType_ReadOnlyStorage, var.buffer, var.nbytes}, + {6, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(BNParams)}, + }, + constants.data(), + constants.size()); + + WebGPUDispatch dispatch{}; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; + dispatch.workgroup_count_x = grid.count_x; + dispatch.workgroup_count_y = grid.count_y; + graph.add_dispatch(dispatch); + + wgpuBufferRelease(uniform_buffer); + if (wbind.owned_dummy != nullptr) { + wgpuBufferRelease(wbind.owned_dummy); + } + if (bbind.owned_dummy != nullptr) { + wgpuBufferRelease(bbind.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + aten._native_batch_norm_legit_no_training.default, batch_norm_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl b/backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl new file mode 100644 index 00000000000..015f6c7f8a2 --- /dev/null +++ b/backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl @@ -0,0 +1,41 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var mean: array; +@group(0) @binding(5) var var_: array; + +struct Params { + num_elements: u32, + C: u32, + HW: u32, + has_weight: u32, + has_bias: u32, + eps: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(6) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +// Inference batch_norm NCHW fp32: per-channel affine from running stats. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let c = (i / params.HW) % params.C; + let inv = inverseSqrt(var_[c] + params.eps); + var g: f32 = 1.0; + if (params.has_weight == 1u) { + g = weight[c]; + } + var b: f32 = 0.0; + if (params.has_bias == 1u) { + b = bias[c]; + } + out[i] = (inp[i] - mean[c]) * inv * g + b; +} diff --git a/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h b/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h new file mode 100644 index 00000000000..f78d8cdc765 --- /dev/null +++ b/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from batch_norm.wgsl - DO NOT EDIT. +// wgsl-sha256: ac7208cec0fe1454698b9b16985172c1ef34c994880b568238eaeaf1d62c0342 +inline constexpr const char* kBatchNormWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var mean: array; +@group(0) @binding(5) var var_: array; + +struct Params { + num_elements: u32, + C: u32, + HW: u32, + has_weight: u32, + has_bias: u32, + eps: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(6) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +// Inference batch_norm NCHW fp32: per-channel affine from running stats. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let c = (i / params.HW) % params.C; + let inv = inverseSqrt(var_[c] + params.eps); + var g: f32 = 1.0; + if (params.has_weight == 1u) { + g = weight[c]; + } + var b: f32 = 0.0; + if (params.has_bias == 1u) { + b = bias[c]; + } + out[i] = (inp[i] - mean[c]) * inv * g + b; +} +)"; + +inline constexpr uint32_t kBatchNormWorkgroupSizeX = 256; +inline constexpr uint32_t kBatchNormWorkgroupSizeY = 1; +inline constexpr uint32_t kBatchNormWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 2a6e7187df2034b28e25a878036c630517f6793b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:28 -0700 Subject: [PATCH 19/22] [ExecuTorch][WebGPU] batch_norm op tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20994 **Op-test suite for `aten._native_batch_norm_legit_no_training.default`** (stacked on the batch_norm op diff). Adds `test/ops/test_batch_norm.py` (`BatchNorm2dModule` — `nn.BatchNorm2d.eval()` with deterministic running stats + affine) + a `@register_op_test("batch_norm")` suite in `test/op_tests/cases.py` (3 cases). Covers affine + non-affine (optional weight/bias) and an odd H*W; only the populated `out` ValueList entry is compared (out_index 0). ghstack-source-id: 405949784 @exported-using-ghexport Differential Revision: [D112417282](https://our.internmc.facebook.com/intern/diff/D112417282/) --- backends/webgpu/test/op_tests/cases.py | 37 ++++++++++++++++++ backends/webgpu/test/ops/test_batch_norm.py | 42 +++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 backends/webgpu/test/ops/test_batch_norm.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index d256a154fba..9ccb8ad1376 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -775,6 +775,43 @@ def _upsample_bilinear2d_suite() -> WebGPUTestSuite: ) +from executorch.backends.webgpu.test.ops.test_batch_norm import ( + _det_input as _bn_det_input, + BatchNorm2dModule, +) + + +@register_op_test("batch_norm") +def _batch_norm_suite() -> WebGPUTestSuite: + # MODNet decoder / CNN-backbone inference batch norm (eval -> the no-training + # variant). Covers affine + non-affine (optional weight/bias) and an odd H*W. + # Only the `out` ValueList entry is compared (out_index 0). + return WebGPUTestSuite( + module_factory=lambda num_features, affine: BatchNorm2dModule( + num_features, affine + ).eval(), + cases=[ + Case( + name="affine_c8", + construct={"num_features": 8, "affine": True}, + inputs=(InputSpec(shape=(1, 8, 12, 12), gen=_bn_det_input),), + ), + Case( + name="no_affine_c8", + construct={"num_features": 8, "affine": False}, + inputs=(InputSpec(shape=(1, 8, 12, 12), gen=_bn_det_input),), + ), + Case( + name="c16_odd", + construct={"num_features": 16, "affine": True}, + inputs=(InputSpec(shape=(1, 16, 5, 7), gen=_bn_det_input),), + ), + ], + atol=1e-3, + rtol=1e-3, + ) + + from executorch.backends.webgpu.test.ops.test_max_pool2d import ( _det_input as _maxpool_det_input, MaxPool2dModule, diff --git a/backends/webgpu/test/ops/test_batch_norm.py b/backends/webgpu/test/ops/test_batch_norm.py new file mode 100644 index 00000000000..9b604f1893a --- /dev/null +++ b/backends/webgpu/test/ops/test_batch_norm.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten._native_batch_norm_legit_no_training.default` module for the WebGPU op-test framework. + +Inference batch norm is on MODNet's decoder and many CNN backbones; `nn.BatchNorm2d` +in eval mode lowers to `_native_batch_norm_legit_no_training`, applying the per-channel +running-stat affine. Deterministic running stats + affine make the golden non-trivial. +""" + +import torch + + +class BatchNorm2dModule(torch.nn.Module): + def __init__( + self, num_features: int, affine: bool = True, eps: float = 1e-5 + ) -> None: + super().__init__() + bn = torch.nn.BatchNorm2d(num_features, eps=eps, affine=affine) + with torch.no_grad(): + bn.running_mean.copy_(torch.linspace(-1.0, 1.0, num_features)) + bn.running_var.copy_(torch.linspace(0.5, 2.0, num_features)) + if affine: + bn.weight.copy_(torch.linspace(0.5, 1.5, num_features)) + bn.bias.copy_(torch.linspace(-0.5, 0.5, num_features)) + bn.eval() + self.bn = bn + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.bn(x) + + +def _det_input(shape) -> torch.Tensor: + # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives. + n = 1 + for s in shape: + n *= s + idx = torch.arange(n, dtype=torch.int64) + return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape) From c5caff1c128ccbb4e76141a4db9882d9c9d5d324 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:29 -0700 Subject: [PATCH 20/22] [ExecuTorch][WebGPU] Add split_with_sizes_copy op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20995 **Add `aten.split_with_sizes_copy.default` to the WebGPU backend** — the split on YOLO's Detect head (separating the concatenated box / objectness / class predictions). **Problem**: The WebGPU delegate had no `split_with_sizes_copy`, so YOLO object-detection could not fully delegate to the GPU. **Solution**: Split `self` along `dim` into N contiguous chunks. Each chunk is a step-1 slice from the running offset, reusing the `slice.wgsl` gather kernel — one dispatch per output. Outputs arrive as a serialized ValueList. Each chunk writes its own distinct output buffer and reads only the shared input, so there is no cross-dispatch read-after-write hazard. **Implementation**: - `runtime/ops/split_with_sizes/SplitWithSizes.cpp` registering `aten.split_with_sizes_copy.default`; reuses `slice_wgsl.h` (no new shader). Uses `utils::make_compute_pipeline` (auto-derived bind-group layout) rather than hand-rolling the layout / pipeline / bind group. - Mirrors the Vulkan `split_with_sizes_copy` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only; `dim` normalized + range-checked; `outputs == sizes` count enforced (fail-loud). Reuses the `slice` op's `slice_wgsl.h`, so this diff stacks above `slice` and must land after it. No change to existing ops. ghstack-source-id: 405949786 @exported-using-ghexport Differential Revision: [D112417284](https://our.internmc.facebook.com/intern/diff/D112417284/) --- .../ops/split_with_sizes/SplitWithSizes.cpp | 147 ++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp diff --git a/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp b/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp new file mode 100644 index 00000000000..75542377f69 --- /dev/null +++ b/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp @@ -0,0 +1,147 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct SliceParams { + uint32_t dim; + uint32_t start; + uint32_t step; + uint32_t _pad; +}; + +// aten.split_with_sizes_copy: N contiguous chunks via per-output slice gather. +void split_with_sizes_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 4) { + throw std::runtime_error("WebGPU split_with_sizes: expected >=4 args"); + } + const int in_id = args.at(0); + const std::vector& sizes = graph.get_int_list(args.at(1)); + int64_t dim = graph.get_int(args.at(2)); + const std::vector& outs = graph.get_value_list(args.at(args.size() - 1)); + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const int in_ndim = static_cast(in_tensor.dims.size()); + if (dim < 0) { + dim += in_ndim; + } + if (dim < 0 || dim >= in_ndim) { + throw std::runtime_error("split_with_sizes: dim out of range"); + } + if (outs.size() != sizes.size()) { + throw std::runtime_error("split_with_sizes: outputs != sizes count"); + } + + // Validate the split contract up front (before any dispatch): each size is + // non-negative and they sum to the split dim's extent -- a negative or + // oversize value would otherwise wrap when cast to u32 into an out-of-bounds + // gather offset. + int64_t sizes_sum = 0; + for (int64_t s : sizes) { + if (s < 0) { + throw std::runtime_error("split_with_sizes: negative split size"); + } + sizes_sum += s; + } + if (sizes_sum != in_tensor.dims[dim]) { + throw std::runtime_error( + "split_with_sizes: sizes must sum to the split dim extent"); + } + + TensorMeta in_meta; + fill_tensor_meta(in_tensor, &in_meta); + if (in_tensor.nbytes != static_cast(in_meta.numel) * sizeof(float)) { + throw std::runtime_error("split_with_sizes: non-fp32 input"); + } + + uint32_t start = 0; + for (size_t i = 0; i < outs.size(); i++) { + const auto& out_tensor = graph.get_tensor(outs[i]); + TensorMeta out_meta; + fill_tensor_meta(out_tensor, &out_meta); + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float)) { + throw std::runtime_error("split_with_sizes: non-fp32 output"); + } + + SliceParams params = {}; + params.dim = static_cast(dim); + params.start = start; + params.step = 1u; + start += static_cast(sizes[i]); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, out_meta.numel, wg_size, "split_with_sizes"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in_meta_buf = + utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(SliceParams)); + graph.add_uniform_buffer_bytes( + 2 * sizeof(TensorMeta) + sizeof(SliceParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSliceWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + out_meta_buf, + sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(SliceParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + graph.own_uniform_buffer(out_meta_buf); + graph.own_uniform_buffer(in_meta_buf); + graph.own_uniform_buffer(params_buf); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.split_with_sizes_copy.default, split_with_sizes_impl); +} + +} // namespace executorch::backends::webgpu From 019cf753898ac63e93876683b385fe88f3c63b25 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:29 -0700 Subject: [PATCH 21/22] [ExecuTorch][WebGPU] split_with_sizes_copy op tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20996 **Op-test suite for `aten.split_with_sizes_copy.default`** (stacked on the split_with_sizes_copy op diff). Adds `test/ops/test_split_with_sizes_copy.py` (`SplitWithSizesModule` — `torch.split` by a size list) + a `@register_op_test("split_with_sizes_copy")` suite in `test/op_tests/cases.py` (3 cases: a 3-way channel split, a dim-0 split, and a last-dim split). Multi-output: the framework compares chunk 0 (out_index 0) while each case exercises all N per-chunk dispatches. `copy` is bit-exact, so the golden is float32. ghstack-source-id: 405949787 @exported-using-ghexport Differential Revision: [D112417279](https://our.internmc.facebook.com/intern/diff/D112417279/) --- backends/webgpu/test/op_tests/cases.py | 46 +++++++++++++++++++ .../test/ops/test_split_with_sizes_copy.py | 39 ++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 backends/webgpu/test/ops/test_split_with_sizes_copy.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 9ccb8ad1376..a73e8e8d04a 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -812,6 +812,52 @@ def _batch_norm_suite() -> WebGPUTestSuite: ) +from executorch.backends.webgpu.test.ops.test_split_with_sizes_copy import ( + _det_input as _split_det_input, + SplitWithSizesModule, +) + + +@register_op_test("split_with_sizes_copy") +def _split_with_sizes_copy_suite() -> WebGPUTestSuite: + # YOLO Detect-head split of concatenated predictions. Multi-output: the + # framework compares chunk 0 (out_index 0) while each case runs all N + # per-chunk dispatches. Covers a 3-way channel split, a dim-0 split, and + # a last-dim split. copy is bit-exact -> float32 golden. + return WebGPUTestSuite( + module_factory=lambda sizes, dim, out_order=None: SplitWithSizesModule( + sizes, dim, out_order + ), + cases=[ + Case( + name="three_dim1", + construct={"sizes": [2, 3, 3], "dim": 1}, + inputs=(InputSpec(shape=(1, 8, 4, 4), gen=_split_det_input),), + ), + Case( + name="two_dim0", + construct={"sizes": [3, 2], "dim": 0}, + inputs=(InputSpec(shape=(5, 4), gen=_split_det_input),), + ), + Case( + name="dim_last", + construct={"sizes": [4, 4], "dim": -1}, + inputs=(InputSpec(shape=(2, 8), gen=_split_det_input),), + ), + # Reorder so chunk 1 (running offset > 0) is output 0 -> verifies the + # per-chunk start accumulation (out_index 0 is all the framework checks). + Case( + name="offset_chunk1_first", + construct={"sizes": [2, 3, 3], "dim": 1, "out_order": [1, 0, 2]}, + inputs=(InputSpec(shape=(1, 8, 4, 4), gen=_split_det_input),), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) + + from executorch.backends.webgpu.test.ops.test_max_pool2d import ( _det_input as _maxpool_det_input, MaxPool2dModule, diff --git a/backends/webgpu/test/ops/test_split_with_sizes_copy.py b/backends/webgpu/test/ops/test_split_with_sizes_copy.py new file mode 100644 index 00000000000..c1771039e41 --- /dev/null +++ b/backends/webgpu/test/ops/test_split_with_sizes_copy.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.split_with_sizes_copy.default` module for the WebGPU op-test framework. + +split_with_sizes is on YOLO's Detect head (splitting the concatenated +box/objectness/class predictions). `torch.split` by a size list lowers to +`split_with_sizes_copy` in the edge dialect. +""" + +import torch + + +class SplitWithSizesModule(torch.nn.Module): + def __init__(self, sizes, dim: int, out_order=None) -> None: + super().__init__() + self.sizes = sizes + self.dim = dim + # Reorder the returned chunks so a non-first chunk can be output 0 (the + # op-test framework compares out_index 0) -- exercises the running offset. + self.out_order = out_order + + def forward(self, x: torch.Tensor): + chunks = torch.split(x, self.sizes, self.dim) + if self.out_order is not None: + return tuple(chunks[i] for i in self.out_order) + return chunks + + +def _det_input(shape) -> torch.Tensor: + # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives. + n = 1 + for s in shape: + n *= s + idx = torch.arange(n, dtype=torch.int64) + return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape) From 2f11ef01637ab67ec41ebdd7c0f906604f5c7cbf Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:30 -0700 Subject: [PATCH 22/22] [ExecuTorch][WebGPU] Migrate 14 single-pipeline ops to make_compute_pipeline Pull Request resolved: https://github.com/pytorch/executorch/pull/20868 Route `add`, `mul`, `index`, `permute`, `select`, `slice`, `update_cache`, `rms_norm`, and `embedding_q4gsw` through the shared `utils::make_compute_pipeline` helper (which uses `layout:"auto"`), replacing each op's hand-written `WGPUBindGroupLayoutEntry[]` + pipeline-layout + bind-group boilerplate (~50-110 lines each) with a single helper call. The driver now derives the bind-group layout from the shader's statically-used bindings. Byte-behavior is preserved: identical binding indices/types/buffers/sizes, dispatch workgroup counts, resize hooks, and validations; override constants (`wg_size`) passed via the helper's `constants` param. Extends the Diff 1 layout:"auto" adoption to the trivial single-dispatch ops. ghstack-source-id: 405949795 @exported-using-ghexport Differential Revision: [D110836665](https://our.internmc.facebook.com/intern/diff/D110836665/) --- .../runtime/ops/boolean_op/BooleanOp.cpp | 70 +++-------- backends/webgpu/runtime/ops/div/BinaryOp.cpp | 116 ++++-------------- .../runtime/ops/expand_copy/ExpandCopy.cpp | 86 +++---------- backends/webgpu/runtime/ops/fill/Fill.cpp | 72 +++-------- backends/webgpu/runtime/ops/gather/Gather.cpp | 98 ++++----------- backends/webgpu/runtime/ops/index/Index.cpp | 92 ++++---------- .../runtime/ops/log_softmax/LogSoftmax.cpp | 72 +++-------- .../webgpu/runtime/ops/permute/Permute.cpp | 93 ++++---------- backends/webgpu/runtime/ops/reduce/Reduce.cpp | 76 +++--------- backends/webgpu/runtime/ops/select/Select.cpp | 89 +++----------- backends/webgpu/runtime/ops/slice/Slice.cpp | 92 +++----------- .../webgpu/runtime/ops/softmax/Softmax.cpp | 76 +++--------- backends/webgpu/runtime/ops/sub/BinaryOp.cpp | 116 ++++-------------- .../runtime/ops/update_cache/UpdateCache.cpp | 84 ++++--------- backends/webgpu/runtime/ops/where/Where.cpp | 104 +++++----------- 15 files changed, 318 insertions(+), 1018 deletions(-) diff --git a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp index 1da013f77a3..e62f387d258 100644 --- a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp +++ b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp @@ -84,61 +84,22 @@ void dispatch_bool_op( utils::make_uniform(device, ¶ms, sizeof(BoolOpParams)); graph.add_uniform_buffer_bytes(sizeof(BoolOpParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - for (uint32_t i = 0; i < 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - } - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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); - - 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[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = self_tensor.buffer; - bg_entries[0].size = self_bind_size; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_bind_size; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(BoolOpParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_bind_size}, + {1, WGPUBufferBindingType_Storage, out_tensor.buffer, out_bind_size}, + {2, WGPUBufferBindingType_Uniform, params_buf, sizeof(BoolOpParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); WGPUBuffer p_buf = params_buf; auto resize = @@ -158,9 +119,6 @@ void dispatch_bool_op( }; graph.add_tensor_resize_hook(self_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/div/BinaryOp.cpp b/backends/webgpu/runtime/ops/div/BinaryOp.cpp index ec359b6684d..279b50e5ec5 100644 --- a/backends/webgpu/runtime/ops/div/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/div/BinaryOp.cpp @@ -78,94 +78,35 @@ void div_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in2_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(3 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBinaryDivWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - entries[5].binding = 5; - entries[5].visibility = WGPUShaderStage_Compute; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - 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); - - 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[6] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - - bg_entries[4].binding = 4; - bg_entries[4].buffer = in1_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - - bg_entries[5].binding = 5; - bg_entries[5].buffer = in2_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBinaryDivWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in1_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, in2_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "div", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "div", + workgroup_count.y}); // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; @@ -206,9 +147,6 @@ void div_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, div_resize); graph.add_tensor_resize_hook(in2_id, div_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so a resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in1_meta_buf); diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp index 507677c7bca..5bfc0fa3bf1 100644 --- a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -60,72 +60,26 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kExpandCopyWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - 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); - - 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[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kExpandCopyWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); } diff --git a/backends/webgpu/runtime/ops/fill/Fill.cpp b/backends/webgpu/runtime/ops/fill/Fill.cpp index 2ccd48cc8fa..8da3b31577d 100644 --- a/backends/webgpu/runtime/ops/fill/Fill.cpp +++ b/backends/webgpu/runtime/ops/fill/Fill.cpp @@ -67,57 +67,28 @@ void add_fill( utils::make_uniform(device, ¶ms, sizeof(FillParams)); graph.add_uniform_buffer_bytes(sizeof(FillParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kFillWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[2] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 2; - 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); - - 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[2] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out_tensor.buffer; - bg_entries[0].size = out_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = uniform_buffer; - bg_entries[1].size = sizeof(FillParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 2; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kFillWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(FillParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "", + workgroup_count.y}); // Dynamic shapes: recompute num_elements/dispatch from the live output dims. WGPUBuffer params_buf = uniform_buffer; @@ -136,9 +107,6 @@ void add_fill( g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/gather/Gather.cpp b/backends/webgpu/runtime/ops/gather/Gather.cpp index d7e170129d1..9236a6714a2 100644 --- a/backends/webgpu/runtime/ops/gather/Gather.cpp +++ b/backends/webgpu/runtime/ops/gather/Gather.cpp @@ -93,81 +93,31 @@ void gather_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(GatherParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(GatherParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kGatherWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - entries[0].binding = 0; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - entries[5].binding = 5; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - for (auto& e : entries) { - e.visibility = WGPUShaderStage_Compute; - } + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kGatherWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + index_tensor.buffer, + index_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, self_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, params_buf, sizeof(GatherParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - 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); - - 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[6] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = self_tensor.buffer; - bg_entries[0].size = self_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = index_tensor.buffer; - bg_entries[1].size = index_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = self_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - bg_entries[5].binding = 5; - bg_entries[5].buffer = params_buf; - bg_entries[5].size = sizeof(GatherParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(self_meta_buf); wgpuBufferRelease(params_buf); diff --git a/backends/webgpu/runtime/ops/index/Index.cpp b/backends/webgpu/runtime/ops/index/Index.cpp index 0c33d616c03..b122cc3af50 100644 --- a/backends/webgpu/runtime/ops/index/Index.cpp +++ b/backends/webgpu/runtime/ops/index/Index.cpp @@ -109,73 +109,33 @@ void index_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(IndexParams)); graph.add_uniform_buffer_bytes(sizeof(IndexParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kIndexWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // self (read), out (read_write), index (read i32), params (uniform). - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - 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); - - 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[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = self_tensor.buffer; - bg_entries[0].size = self_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = index_tensor.buffer; - bg_entries[2].size = index_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(IndexParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kIndexWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + index_tensor.buffer, + index_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(IndexParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + // The bind group keeps the uniform buffer alive until release. wgpuBufferRelease(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp index f4fb9952706..e88114e6516 100644 --- a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp +++ b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp @@ -120,67 +120,26 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector& args) { wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(LogSoftmaxParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kLogSoftmaxWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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(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[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out.buffer; - bg_entries[1].size = out.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(LogSoftmaxParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLogSoftmaxWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LogSoftmaxParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "log_softmax"}); + {bundle.pipeline, bundle.bind_group, workgroup_count, "log_softmax"}); // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. WGPUBuffer params_buf = uniform_buffer; @@ -215,9 +174,6 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, std::vector(d.begin(), d.end())); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/permute/Permute.cpp b/backends/webgpu/runtime/ops/permute/Permute.cpp index c3b9b5c70c6..fb23fce7019 100644 --- a/backends/webgpu/runtime/ops/permute/Permute.cpp +++ b/backends/webgpu/runtime/ops/permute/Permute.cpp @@ -108,80 +108,33 @@ void permute_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_uniform_buffer_bytes( 2 * sizeof(TensorMeta) + sizeof(PermuteParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kPermuteWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - 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); - - 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 = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(PermuteParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kPermuteWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(PermuteParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "permute", + workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Drop our refs; the bind group keeps the uniforms alive until release. wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/reduce/Reduce.cpp b/backends/webgpu/runtime/ops/reduce/Reduce.cpp index b276354d2a8..7e185d38f82 100644 --- a/backends/webgpu/runtime/ops/reduce/Reduce.cpp +++ b/backends/webgpu/runtime/ops/reduce/Reduce.cpp @@ -132,67 +132,30 @@ void reduce_impl( wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(ReduceParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kReduceWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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(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[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out.buffer; - bg_entries[1].size = out.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(ReduceParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kReduceWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ReduceParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, op_name, workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + op_name, + workgroup_count.y}); // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. WGPUBuffer params_buf = uniform_buffer; @@ -248,9 +211,6 @@ void reduce_impl( g.set_cur_dims(out_id, od); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/select/Select.cpp b/backends/webgpu/runtime/ops/select/Select.cpp index 5d3e8103b7d..f2e69823974 100644 --- a/backends/webgpu/runtime/ops/select/Select.cpp +++ b/backends/webgpu/runtime/ops/select/Select.cpp @@ -108,76 +108,28 @@ void select_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(SelectParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(SelectParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSelectWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - 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); - - 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 = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(SelectParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSelectWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(SelectParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); // Dynamic shapes: out = in minus `dim`; re-resolve index, meta, dispatch. graph.add_tensor_resize_hook( @@ -219,9 +171,6 @@ void select_impl(WebGPUGraph& graph, const std::vector& args) { g.device(), out_numel, wg_size, "select(resize)"); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so the resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/slice/Slice.cpp b/backends/webgpu/runtime/ops/slice/Slice.cpp index 4fbbfccc5fe..f4fcb6088a8 100644 --- a/backends/webgpu/runtime/ops/slice/Slice.cpp +++ b/backends/webgpu/runtime/ops/slice/Slice.cpp @@ -167,75 +167,26 @@ void slice_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(SliceParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(SliceParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSliceWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - 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); - - 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 = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(SliceParams); - - 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}); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSliceWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(SliceParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); const size_t dispatch_idx = graph.num_dispatches() - 1; // Dynamic shapes: live start/end -> out[dim] len + meta/params/dispatch. @@ -289,9 +240,6 @@ void slice_impl(WebGPUGraph& graph, const std::vector& args) { } graph.add_tensor_resize_hook(in_id, recompute); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns the uniforms so the resize hook can rewrite them; freed in dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/softmax/Softmax.cpp b/backends/webgpu/runtime/ops/softmax/Softmax.cpp index 72b48cd8bdd..e5f60f1d4c2 100644 --- a/backends/webgpu/runtime/ops/softmax/Softmax.cpp +++ b/backends/webgpu/runtime/ops/softmax/Softmax.cpp @@ -109,67 +109,26 @@ void softmax_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(SoftmaxParams)); graph.add_uniform_buffer_bytes(sizeof(SoftmaxParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSoftmaxWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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(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[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out.buffer; - bg_entries[1].size = out.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(SoftmaxParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count, "softmax"}); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSoftmaxWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(SoftmaxParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, workgroup_count, "softmax"}); // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. WGPUBuffer params_buf = uniform_buffer; @@ -204,9 +163,6 @@ void softmax_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp index 6945feb8da2..13b609cc1e7 100644 --- a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp @@ -87,94 +87,35 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in2_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(3 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBinarySubWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - entries[5].binding = 5; - entries[5].visibility = WGPUShaderStage_Compute; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - 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); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 2; - pipeline_desc.compute.constants = constants; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[6] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - - bg_entries[4].binding = 4; - bg_entries[4].buffer = in1_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - - bg_entries[5].binding = 5; - bg_entries[5].buffer = in2_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBinarySubWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in1_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, in2_meta_buf, sizeof(TensorMeta)}, + }, + constants, + 2); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "sub", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "sub", + workgroup_count.y}); // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; @@ -215,9 +156,6 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, sub_resize); graph.add_tensor_resize_hook(in2_id, sub_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so a resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in1_meta_buf); diff --git a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp index dc23a45eb91..5bdba562633 100644 --- a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp +++ b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp @@ -119,72 +119,32 @@ void update_cache_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_uniform_buffer_bytes(sizeof(UpdateCacheParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kUpdateCacheWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: cache (rw storage) + value (ro storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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(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[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = cache_tensor.buffer; - bg_entries[0].size = cache_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = value_tensor.buffer; - bg_entries[1].size = value_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(UpdateCacheParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count_x}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kUpdateCacheWGSL, + { + {0, + WGPUBufferBindingType_Storage, + cache_tensor.buffer, + cache_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + value_tensor.buffer, + value_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(UpdateCacheParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count_x}); + // Drop our ref; the bind group keeps the uniform buffer alive until release. wgpuBufferRelease(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/where/Where.cpp b/backends/webgpu/runtime/ops/where/Where.cpp index bde779498dd..f8aba4a5d01 100644 --- a/backends/webgpu/runtime/ops/where/Where.cpp +++ b/backends/webgpu/runtime/ops/where/Where.cpp @@ -92,81 +92,36 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &b_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(4 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kWhereWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[8] = {}; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].buffer.type = WGPUBufferBindingType_Storage; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - entries[6].buffer.type = WGPUBufferBindingType_Uniform; - entries[7].buffer.type = WGPUBufferBindingType_Uniform; - for (uint32_t i = 0; i < 8; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - } - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 8; - 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); - - 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[8] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = cond_tensor.buffer; - bg_entries[0].size = cond_bind_size; - bg_entries[1].binding = 1; - bg_entries[1].buffer = a_tensor.buffer; - bg_entries[1].size = a_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = b_tensor.buffer; - bg_entries[2].size = b_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_tensor.buffer; - bg_entries[3].size = out_tensor.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = out_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - bg_entries[5].binding = 5; - bg_entries[5].buffer = cond_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - bg_entries[6].binding = 6; - bg_entries[6].buffer = a_meta_buf; - bg_entries[6].size = sizeof(TensorMeta); - bg_entries[7].binding = 7; - bg_entries[7].buffer = b_meta_buf; - bg_entries[7].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 8; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kWhereWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + cond_tensor.buffer, + cond_bind_size}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + a_tensor.buffer, + a_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + b_tensor.buffer, + b_tensor.nbytes}, + {3, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {4, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, cond_meta_buf, sizeof(TensorMeta)}, + {6, WGPUBufferBindingType_Uniform, a_meta_buf, sizeof(TensorMeta)}, + {7, WGPUBufferBindingType_Uniform, b_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); // Dynamic shapes: rebuild the 4 broadcast TensorMeta UBOs + dispatch count. WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf, @@ -223,9 +178,6 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(a_id, where_resize); graph.add_tensor_resize_hook(b_id, where_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(cond_meta_buf); graph.own_uniform_buffer(a_meta_buf);