From 1f620f956d264cf05f435b53a3241ec4dccb9b4c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:13 -0700 Subject: [PATCH 01/41] [ExecuTorch][WebGPU] Shared pipeline + dispatch-grid + vec4-align helpers (op-suite foundation) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20842 **Foundation for the vision/voice/seg WebGPU op suite: the shared helpers every op builds on.** Adds to `WebGPUUtils.h` the `make_compute_pipeline` bundle (`BindingSpec`, `ComputePipelineBundle`, `make_storage_zeros`), `check_vec4_aligned`, and the validation/dispatch-config helpers every handler reuses: `check_fp32` and `check_elementwise_fp32_io` (fp32 byte-size + elementwise input/output guards), `checked_u32` (the `uint64_t`->`uint32_t` dispatch-count guard), `make_wg_size_constant` (the 1-element override-constant sibling of `make_grid_constants`), and `compute_tile_grid_2d` (the tiled-GEMM 2D grid, device-limit-queried); a new `WebGPUDispatchMath.h` with the near-square 2D dispatch-grid fold (`div_up`, `numel`); a `WebGPUGraph::add_dispatch_2d` convenience for the 2D dispatch fill; shared `WebGPUGraph` hooks; and op-test-driver infra (+ `test_webgpu_utils`). No op is registered here — each op lands on top using these helpers. Co-authored-with: Claude Code. ghstack-source-id: 405949691 @exported-using-ghexport Differential Revision: [D110836666](https://our.internmc.facebook.com/intern/diff/D110836666/) --- backends/webgpu/CMakeLists.txt | 8 +- backends/webgpu/runtime/WebGPUDispatchMath.h | 108 +++++ backends/webgpu/runtime/WebGPUGraph.cpp | 9 + backends/webgpu/runtime/WebGPUGraph.h | 19 + backends/webgpu/runtime/WebGPUUtils.h | 381 ++++++++++++++++-- .../webgpu/test/native/test_webgpu_utils.cpp | 64 +++ backends/webgpu/test/op_tests/driver_util.cpp | 15 + backends/webgpu/test/op_tests/driver_util.h | 2 + .../webgpu/test/op_tests/generate_op_tests.py | 18 +- .../webgpu/test/op_tests/op_test_driver.cpp | 12 +- 10 files changed, 597 insertions(+), 39 deletions(-) create mode 100644 backends/webgpu/runtime/WebGPUDispatchMath.h create mode 100644 backends/webgpu/test/native/test_webgpu_utils.cpp diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 78d9c4f30dd..3060d52bf03 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -109,6 +109,7 @@ function(add_webgpu_native_test test_name test_src) target_link_libraries( ${test_name} PRIVATE webgpu_backend + vulkan_schema ${WEBGPU_GPU_LIB} executorch_core extension_module_static @@ -153,10 +154,11 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) ) # Device-free util unit test: no backend/Dawn link (pure manifest/tolerance - # helpers), so it does NOT use the native-test helper. + # + dispatch-grid-math helpers), so it does NOT use the native-test helper. add_executable( - webgpu_op_test_util_test test/op_tests/test_driver_util.cpp - test/op_tests/driver_util.cpp + webgpu_op_test_util_test + test/op_tests/test_driver_util.cpp test/op_tests/driver_util.cpp + test/native/test_webgpu_utils.cpp ) target_include_directories( webgpu_op_test_util_test diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h new file mode 100644 index 00000000000..56544c6d496 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -0,0 +1,108 @@ +/* + * 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 + +// Pure dispatch-grid math with zero WebGPU/Dawn dependency, so it is +// unit-testable without a WGPUDevice (split out of WebGPUUtils.h, which +// requires for its device-facing functions). + +#include +#include +#include +#include +#include + +namespace executorch::backends::webgpu::utils { + +// Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). +template +inline T div_up(T a, T b) { + return (a + b - 1) / b; +} + +// Product of a tensor's dims; the same accumulation was duplicated per-op. +inline uint64_t numel(const std::vector& dims) { + uint64_t n = 1; + for (int64_t d : dims) { + if (d < 0) { + throw std::runtime_error("numel: negative dimension"); + } + uint64_t ud = static_cast(d); + if (ud != 0 && n > UINT64_MAX / ud) { + throw std::runtime_error("numel: element count overflow"); + } + n *= ud; + } + return n; +} + +// Broadcasts a 1- or 2-element int list to (h, w); PyTorch's convention for +// kernel_size/stride/padding/dilation args. Was duplicated as a local `hw` +// lambda in conv2d/conv_transpose2d/max_pool2d. +inline void parse_hw( + const std::vector& v, + uint32_t& h, + uint32_t& w, + const char* op_name, + const char* arg_name) { + if (v.size() == 1) { + h = w = static_cast(v[0]); + } else if (v.size() == 2) { + h = static_cast(v[0]); + w = static_cast(v[1]); + } else { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + arg_name + + " must be 1 or 2 elements"); + } +} + +// Adaptive 1D->2D dispatch grid. `count_x`/`count_y` are the dispatch dims; +// `stride_x` (= count_x * wg_size) lets the shader decode a flat index as +// `gid.y * stride_x + gid.x`. Used by ops whose element count can exceed the +// 65535 per-dimension ceiling that compute_1d_workgroup_count throws on. +struct DispatchGrid { + uint32_t wg_size; + uint32_t count_x; + uint32_t count_y; + uint32_t stride_x; +}; + +// Given the workgroup count needed (1D) and the device's per-dimension +// dispatch-count ceiling, compute a near-square 2D grid rather than +// {max_dim, div_up(total, max_dim)} — maxing one dim pads the other with +// mostly-idle workgroups (up to ~2x the needed launch) when total isn't a +// clean multiple of max_dim. +inline DispatchGrid compute_dispatch_grid_from_limits( + uint32_t total, // workgroups needed (1D) + uint32_t wg_size, + uint32_t max_dim, + const char* op_name) { + DispatchGrid g; + g.wg_size = wg_size; + if (total <= max_dim) { + g.count_x = total; + g.count_y = 1; + } else { + uint32_t sq = + static_cast(std::ceil(std::sqrt(static_cast(total)))); + g.count_x = sq < max_dim ? sq : max_dim; + g.count_y = div_up(total, g.count_x); + if (g.count_y > + max_dim) { // > max_dim^2 * wg threads — astronomically large + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": dispatch exceeds 2D grid capacity"); + } + } + g.stride_x = g.count_x * g.wg_size; + return g; +} + +} // namespace executorch::backends::webgpu::utils diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 2bb091d391e..abffa25e969 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -418,6 +418,7 @@ void WebGPUGraph::build( value_lists_.resize(num_vals); doubles_.resize(num_vals, 0.0); bools_.resize(num_vals, false); + strings_.resize(num_vals); // Pre-scan the op chain: a constant may be DEFERRED (no eager GPU buffer; the // prepack node materializes it once) only if it is a prepack source AND never @@ -625,6 +626,14 @@ void WebGPUGraph::build( bools_[i] = val->value_as_Bool()->bool_val(); break; } + case vkgraph::GraphTypes::String: { + value_types_[i] = ValueType::String; + const auto* sv = val->value_as_String()->string_val(); + if (sv) { + strings_[i] = sv->str(); + } + break; + } case vkgraph::GraphTypes::SymInt: { // Live scalar: small Uniform buffer the CPU rewrites per execute. value_types_[i] = ValueType::SymInt; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 6f99c6162a4..c193de9fe35 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -152,6 +152,10 @@ class WebGPUGraph { bool get_bool(int id) const { return bools_[id]; } + // String value (e.g. gelu's `approximate` kwarg). + const std::string& get_string(int id) const { + return strings_[id]; + } // Live-scalar (SymInt) API; mirrors the Vulkan SymInt/ParamsBuffer UBO. // set_symint writes the buffer + marks dirty only if the value changed. @@ -240,6 +244,20 @@ class WebGPUGraph { return dispatches_.size() - 1; } + // 2D sibling of add_dispatch (sets workgroup_count_y); returns the index. + size_t add_dispatch_2d( + WGPUComputePipeline pipeline, + WGPUBindGroup bind_group, + uint32_t count_x, + uint32_t count_y) { + WebGPUDispatch d; + d.pipeline = pipeline; + d.bind_group = bind_group; + d.workgroup_count_x = count_x; + d.workgroup_count_y = count_y; + return add_dispatch(d); + } + // In-graph buffer-to-buffer DMA (e.g. flat copy); returns the dispatch index. size_t add_buffer_copy(WGPUBuffer src, WGPUBuffer dst, size_t nbytes) { WebGPUDispatch d; @@ -376,6 +394,7 @@ class WebGPUGraph { std::vector> value_lists_; std::vector doubles_; std::vector bools_; + std::vector strings_; // SymInt (live scalar): id -> {live Uniform buffer, current value}, sparse. struct SymIntSlot { diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index b41b02bf6e4..dfcb0c1b03a 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -8,9 +8,13 @@ #pragma once +#include +#include + #include #include +#include #include #include #include @@ -20,22 +24,10 @@ namespace executorch::backends::webgpu::utils { -// Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). -template -inline T div_up(T a, T b) { - return (a + b - 1) / b; -} - -// Product of dims (live element count); used by dynamic-resize hooks. +// Product of dims (live element count); used by dynamic-resize hooks. Delegates +// to numel (single impl; keeps the negative-dim guard, no caller churn). inline uint64_t numel_of(const std::vector& dims) { - uint64_t n = 1; - for (int64_t v : dims) { - if (v < 0) { - throw std::runtime_error("numel_of: negative dimension"); - } - n *= static_cast(v); - } - return n; + return numel(dims); } // Clamp workgroup size to device limit (SwiftShader caps at 128). @@ -81,24 +73,15 @@ inline uint32_t queried_max_workgroups(WGPUDevice device) { // split keeps the waste to O(sqrt(count)). Throws if even a max*max grid is too // small (a 3rd dispatch dimension, out of scope). The shader reconstructs the // linear index from @builtin(num_workgroups), so any x/y factoring works. +// Now delegates to the DispatchMath fold (single fold impl); wg_size=1 makes +// the returned grid's stride_x collapse harmlessly. inline WgCount fold_workgroup_count_2d( uint32_t count, uint32_t max_count, const char* op_name) { - if (count <= max_count) { - return {count, 1u}; - } - uint32_t x = - static_cast(std::ceil(std::sqrt(static_cast(count)))); - x = std::min(x, max_count); - // ceil-div written overflow-safe (count >= 1 here) as count nears UINT32_MAX. - uint32_t y = 1u + (count - 1u) / x; - if (y > max_count) { - throw std::runtime_error( - std::string("WebGPU ") + op_name + - ": workgroup count needs a 3rd dispatch dimension (unsupported)"); - } - return {x, y}; + DispatchGrid g = compute_dispatch_grid_from_limits( + count, /*wg_size=*/1, max_count, op_name); + return {g.count_x, g.count_y}; } // 1D dispatch count (mirrors Vulkan div_up); throws if > device limit. @@ -128,6 +111,60 @@ inline WgCount compute_2d_workgroup_count( count, queried_max_workgroups(device), op_name); } +inline DispatchGrid compute_dispatch_grid( + WGPUDevice device, + uint32_t num_threads, + uint32_t desired_wg, + const char* op_name) { + // Single limits query shared by wg-size clamp + max-dim (avoid 2 queries). + WGPULimits limits = {}; + bool got_limits = wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success; + uint32_t wg = (got_limits && limits.maxComputeInvocationsPerWorkgroup > 0) + ? std::min(desired_wg, limits.maxComputeInvocationsPerWorkgroup) + : desired_wg; + if (wg == 0) { + wg = 1; + } + uint32_t max_dim = (got_limits && limits.maxComputeWorkgroupsPerDimension > 0) + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; // WebGPU spec-default floor + uint32_t total = div_up(num_threads, wg); // workgroups needed (1D) + return compute_dispatch_grid_from_limits(total, wg, max_dim, op_name); +} + +// 2D tile grid for tiled kernels: {div_up(n, tile), div_up(m, tile)} workgroups +// (one per tile); throws if either dim exceeds the device's per-dim limit. +inline WgCount compute_tile_grid_2d( + WGPUDevice device, + uint32_t n, + uint32_t m, + uint32_t tile, + const char* op_name) { + uint32_t max_wgs = queried_max_workgroups(device); + uint32_t x = div_up(n, tile); + uint32_t y = div_up(m, tile); + if (x > max_wgs || y > max_wgs) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": tile grid exceeds dispatch limit"); + } + return {x, y}; +} + +// For "one workgroup per row" kernels (native_layer_norm, sdpa_softmax) where +// num_rows IS the workgroup count already, unlike compute_dispatch_grid's +// element-count-needing-division use case. wg_size=1 makes the returned +// grid's stride_x collapse to exactly count_x, giving a correct flat +// *workgroup*-index decode (row_idx = workgroup_id.y * stride_x + +// workgroup_id.x). +inline DispatchGrid compute_row_dispatch_grid( + WGPUDevice device, + uint32_t num_rows, + const char* op_name) { + return compute_dispatch_grid_from_limits( + num_rows, /*wg_size=*/1, queried_max_workgroups(device), op_name); +} + // Create a uniform buffer mapped-at-creation, copy `size` bytes in, and unmap. inline WGPUBuffer make_uniform(WGPUDevice device, const void* data, size_t size) { @@ -155,4 +192,290 @@ inline uint32_t clamp_workgroup_count(WGPUDevice device, uint32_t desired) { return std::min(desired, queried_max_workgroups(device)); } +// Zero-filled storage buffer; used as a dummy binding for an optional tensor +// arg (bias/mask/affine) the shader gates off and never reads. +inline WGPUBuffer make_storage_zeros(WGPUDevice device, size_t size) { + WGPUBufferDescriptor desc = {}; + desc.size = size; + desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + desc.mappedAtCreation = true; + WGPUBuffer buf = wgpuDeviceCreateBuffer(device, &desc); + if (!buf) { + throw std::runtime_error("make_storage_zeros: buffer creation failed"); + } + void* ptr = wgpuBufferGetMappedRange(buf, 0, size); + if (!ptr) { + wgpuBufferRelease(buf); + throw std::runtime_error("make_storage_zeros: mapped range is null"); + } + std::memset(ptr, 0, size); + wgpuBufferUnmap(buf); + return buf; +} + +// Validates a dim is a multiple of 4 (vec4-alignment precondition for a +// vec4-typed buffer view); throws loud, mirroring this backend's no-silent- +// return convention. Was independently duplicated in sdpa/Sdpa.cpp, +// sdpa_fd_decode/SdpaFdDecode.cpp, and et_vk_sdpa/EtVkSdpa.cpp. +inline void +check_vec4_aligned(uint32_t dim, const char* op_name, const char* dim_name) { + if (dim % 4 != 0) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + dim_name + + " must be a multiple of 4"); + } +} + +// fp32 byte-size guard (no runtime dtype): the serialized bytes must equal the +// element count times sizeof(float), else the tensor is not fp32. Returns the +// element count; replaces the `numel(dims)*sizeof(float) != nbytes` check +// duplicated across the fp32-only op handlers. +inline uint64_t +check_fp32(const WebGPUTensor& t, const char* op_name, const char* label) { + uint64_t n = numel(t.dims); + if (t.nbytes != n * sizeof(float)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + label + " must be fp32"); + } + return n; +} + +// Elementwise unary I/O guard: both buffers bound, byte counts fp32-aligned, +// and input/output sizes equal. Mirrors the guard inlined in every elementwise +// op handler (gelu/sigmoid/...). +inline void check_elementwise_fp32_io( + const WebGPUTensor& in, + const WebGPUTensor& out, + const char* op_name) { + if (in.buffer == nullptr || out.buffer == nullptr) { + throw std::runtime_error(std::string(op_name) + ": null buffer binding"); + } + if (in.nbytes % sizeof(float) != 0 || out.nbytes % sizeof(float) != 0) { + throw std::runtime_error( + std::string(op_name) + ": operand not 4-byte aligned"); + } + if (in.nbytes != out.nbytes) { + throw std::runtime_error( + std::string(op_name) + ": input/output size mismatch"); + } +} + +// Narrow a u64 element/byte count to u32 for a dispatch/param field; throws if +// it would truncate (the count exceeds the u32 addressing range). +inline uint32_t checked_u32(uint64_t v, const char* op_name) { + if (v > 0xFFFFFFFFull) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": output too large"); + } + return static_cast(v); +} + +// Extract a Scalar arg as float: Double -> cast, Int -> cast, Null -> fallback. +// Any other type throws (this backend's no-silent-return convention: an +// unexpected scalar type is a validation failure, not a "use the default"). +// Scope: the 3 already-identical Double/Int/Null-permissive copies only +// (addmm's beta/alpha, constant_pad_nd's value, native_layer_norm's epsilon) +// -- NOT et_vk.sdpa/sdpa's scale extraction, which is a stricter +// Double-or-Null-only policy (rejects Int) that would silently change +// behavior if forced into this signature. +inline float scalar_or(WebGPUGraph& graph, int id, float fallback) { + const auto vt = graph.get_value_type(id); + if (vt == WebGPUGraph::ValueType::Double) { + return static_cast(graph.get_double(id)); + } + if (vt == WebGPUGraph::ValueType::Int) { + return static_cast(graph.get_int(id)); + } + if (vt == WebGPUGraph::ValueType::Null) { + return fallback; + } + throw std::runtime_error("scalar_or: expected Double, Int, or None"); +} + +// Resolves an optional tensor arg (bias/mask/affine) to the binding to use: +// the real tensor's buffer+size if present, else a zero-filled dummy. Caller +// releases `.owned_dummy` (if non-null) after the dispatch is recorded. +struct OptionalBinding { + WGPUBuffer buffer; + uint64_t nbytes; + WGPUBuffer owned_dummy; // non-null only when a dummy was allocated +}; + +inline OptionalBinding make_optional_binding( + WGPUDevice device, + bool present, + WGPUBuffer real_buffer, + uint64_t real_nbytes) { + if (present) { + return {real_buffer, real_nbytes, nullptr}; + } + // 16 bytes (not 4): WebGPU validates a binding's size against the shader's + // DECLARED type at dispatch time regardless of which branch runs, and a + // vec4-typed binding requires >=16 bytes even when never read. + WGPUBuffer dummy = make_storage_zeros(device, 16); + return {dummy, 16, dummy}; +} + +// One compute-shader binding: a bind-group-layout entry + its bind-group +// entry share every field except `type` (layout) vs `buffer`/`size` (bind). +struct BindingSpec { + uint32_t binding; + WGPUBufferBindingType type; + WGPUBuffer buffer; + uint64_t size; +}; + +// Owns the shader module, bind-group layout, and pipeline layout, releasing +// them on destruction. `pipeline` and `bind_group` are NOT released here — +// every op hands them to WebGPUGraph::add_dispatch, which keeps them alive +// for the life of the graph (mirrors the "kept by dispatch" comment repeated +// across every op handler). +struct ComputePipelineBundle { + WGPUShaderModule shader = nullptr; + WGPUBindGroupLayout bind_group_layout = nullptr; + WGPUPipelineLayout pipeline_layout = nullptr; + WGPUComputePipeline pipeline = nullptr; + WGPUBindGroup bind_group = nullptr; + + ComputePipelineBundle() = default; + ComputePipelineBundle(const ComputePipelineBundle&) = delete; + ComputePipelineBundle& operator=(const ComputePipelineBundle&) = delete; + ComputePipelineBundle& operator=(ComputePipelineBundle&&) = delete; + + ComputePipelineBundle(ComputePipelineBundle&& other) noexcept + : shader(other.shader), + bind_group_layout(other.bind_group_layout), + pipeline_layout(other.pipeline_layout), + pipeline(other.pipeline), + bind_group(other.bind_group) { + other.shader = nullptr; + other.bind_group_layout = nullptr; + other.pipeline_layout = nullptr; + other.pipeline = nullptr; + other.bind_group = nullptr; + } + + ~ComputePipelineBundle() { + if (shader != nullptr) { + wgpuShaderModuleRelease(shader); + } + if (bind_group_layout != nullptr) { + wgpuBindGroupLayoutRelease(bind_group_layout); + } + if (pipeline_layout != nullptr) { + wgpuPipelineLayoutRelease(pipeline_layout); + } + } +}; + +// Builds shader module -> bind-group layout -> pipeline layout -> compute +// pipeline -> bind group from one binding list, replacing the ~50-line +// sequence duplicated (with varying binding counts) across every op handler. +inline ComputePipelineBundle make_compute_pipeline( + WGPUDevice device, + const char* wgsl_source, + const std::vector& bindings, + const WGPUConstantEntry* constants = nullptr, + size_t constant_count = 0, + const char* entry_point = "main") { + // Build into the bundle incrementally so an allocation failure below is + // leak-free: ~ComputePipelineBundle releases shader/bind_group_layout/ + // pipeline_layout, and pipeline/bind_group are handled explicitly. + ComputePipelineBundle bundle; + + 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; + bundle.shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + if (bundle.shader == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: shader module creation failed"); + } + + std::vector layout_entries(bindings.size()); + std::vector bind_entries(bindings.size()); + for (size_t i = 0; i < bindings.size(); i++) { + layout_entries[i] = {}; + layout_entries[i].binding = bindings[i].binding; + layout_entries[i].visibility = WGPUShaderStage_Compute; + layout_entries[i].buffer.type = bindings[i].type; + + bind_entries[i] = {}; + bind_entries[i].binding = bindings[i].binding; + bind_entries[i].buffer = bindings[i].buffer; + bind_entries[i].size = bindings[i].size; + } + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = layout_entries.size(); + bgl_desc.entries = layout_entries.data(); + bundle.bind_group_layout = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + if (bundle.bind_group_layout == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: bind group layout creation failed"); + } + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bundle.bind_group_layout; + bundle.pipeline_layout = wgpuDeviceCreatePipelineLayout(device, &pl_desc); + if (bundle.pipeline_layout == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: pipeline layout creation failed"); + } + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = bundle.pipeline_layout; + pipeline_desc.compute.module = bundle.shader; + pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN}; + pipeline_desc.compute.constantCount = constant_count; + pipeline_desc.compute.constants = constants; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + if (pipeline == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: compute pipeline creation failed"); + } + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bundle.bind_group_layout; + bg_desc.entryCount = bind_entries.size(); + bg_desc.entries = bind_entries.data(); + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + if (bind_group == nullptr) { + // pipeline is not yet owned by the bundle (whose dtor intentionally does + // not release pipeline/bind_group); release it before unwinding. + wgpuComputePipelineRelease(pipeline); + throw std::runtime_error( + "make_compute_pipeline: bind group creation failed"); + } + + bundle.pipeline = pipeline; + bundle.bind_group = bind_group; + return bundle; +} + +// The {wg_size, stride_x} override-constant pair every 2D-spill dispatch +// builds from its DispatchGrid; was hand-rolled identically at 7 call sites. +inline std::array make_grid_constants( + const DispatchGrid& grid) { + std::array constants = {}; + constants[0].key = {"wg_size", WGPU_STRLEN}; + constants[0].value = static_cast(grid.wg_size); + constants[1].key = {"stride_x", WGPU_STRLEN}; + constants[1].value = static_cast(grid.stride_x); + return constants; +} + +// The 1-element sibling of make_grid_constants: the single wg_size override- +// constant, hand-rolled at 5 call sites. +inline WGPUConstantEntry make_wg_size_constant(uint32_t wg_size) { + WGPUConstantEntry constant = {}; + constant.key = {"wg_size", WGPU_STRLEN}; + constant.value = static_cast(wg_size); + return constant; +} + } // namespace executorch::backends::webgpu::utils diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp new file mode 100644 index 00000000000..edc0f315294 --- /dev/null +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -0,0 +1,64 @@ +/* + * 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. + */ + +// Device-free unit tests for the dispatch-grid math (WebGPUDispatchMath.h has +// zero WebGPU/Dawn dependency, unlike WebGPUUtils.h which needs a WGPUDevice +// for its other helpers). + +#include + +#include + +using namespace executorch::backends::webgpu; + +TEST(WebGPUUtils, DispatchGridStaysOneDimUnderCeiling) { + utils::DispatchGrid g = + utils::compute_dispatch_grid_from_limits(1000u, 256u, 65535u, "test"); + EXPECT_EQ(g.count_x, 1000u); + EXPECT_EQ(g.count_y, 1u); + EXPECT_EQ(g.stride_x, g.count_x * g.wg_size); +} + +TEST(WebGPUUtils, DispatchGridPastCeilingIsNearSquareNotMaxedOut) { + // total=65536, max_dim=65535: one workgroup past the 1D ceiling. + // The old {max_dim, div_up(total,max_dim)} fold gave (65535, 2) = 131070 + // launched workgroups for 65536 needed (~100% overhead). A near-square + // grid should launch close to the needed count instead. + const uint32_t total = 65536u; + const uint32_t max_dim = 65535u; + utils::DispatchGrid g = + utils::compute_dispatch_grid_from_limits(total, 256u, max_dim, "test"); + + EXPECT_LE(static_cast(g.count_x) * g.count_y, total + total / 10) + << "near-square grid should launch within ~10% of the needed count, " + "not ~2x it"; + // Grid must still cover every needed workgroup. + EXPECT_GE(static_cast(g.count_x) * g.count_y, total); + // Not the old maxed-out-count_x behavior. + EXPECT_NE(g.count_x, max_dim); + EXPECT_EQ(g.stride_x, g.count_x * g.wg_size); +} + +TEST(WebGPUUtils, DispatchGridExactSquareCase) { + // total=65536 factors exactly as 256*256 — the near-square grid should + // find this with zero waste. + utils::DispatchGrid g = + utils::compute_dispatch_grid_from_limits(65536u, 1u, 65535u, "test"); + EXPECT_EQ(g.count_x, 256u); + EXPECT_EQ(g.count_y, 256u); + EXPECT_EQ(static_cast(g.count_x) * g.count_y, 65536u); +} + +TEST(WebGPUUtils, DispatchGridThrowsPastCapacity) { + // total > max_dim^2: even a near-square grid can't fit in the 2D ceiling. + const uint32_t max_dim = 4u; + EXPECT_THROW( + utils::compute_dispatch_grid_from_limits( + static_cast(max_dim) * max_dim + 1u, 1u, max_dim, "test"), + std::runtime_error); +} diff --git a/backends/webgpu/test/op_tests/driver_util.cpp b/backends/webgpu/test/op_tests/driver_util.cpp index aa6c337e0c8..24003a2c3fb 100644 --- a/backends/webgpu/test/op_tests/driver_util.cpp +++ b/backends/webgpu/test/op_tests/driver_util.cpp @@ -49,6 +49,7 @@ std::vector parse_manifest(const std::string& manifest_path) { InputRef ir; ir.path = join(base, ie.at("path").get()); ir.shape = ie.at("shape").get>(); + ir.dtype = ie.value("dtype", std::string("float32")); m.inputs.push_back(std::move(ir)); } const auto& g = e.at("golden"); @@ -64,6 +65,20 @@ std::vector parse_manifest(const std::string& manifest_path) { return out; } +std::vector load_int32_bin(const std::string& path, size_t numel) { + FILE* f = std::fopen(path.c_str(), "rb"); + if (!f) { + return {}; + } + std::vector g(numel); + const size_t n = std::fread(g.data(), sizeof(int32_t), numel, f); + std::fclose(f); + if (n != numel) { + return {}; + } + return g; +} + std::vector load_fp32_bin(const std::string& path, size_t numel) { FILE* f = std::fopen(path.c_str(), "rb"); if (!f) { diff --git a/backends/webgpu/test/op_tests/driver_util.h b/backends/webgpu/test/op_tests/driver_util.h index 3f5e7c47058..4c638c4f57d 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -17,6 +17,7 @@ namespace executorch::backends::webgpu { struct InputRef { std::string path; std::vector shape; + std::string dtype = "float32"; }; struct GoldenRef { @@ -42,6 +43,7 @@ std::vector parse_manifest(const std::string& manifest_path); /// Load raw little-endian fp32; empty on size/IO mismatch. std::vector load_fp32_bin(const std::string& path, size_t numel); +std::vector load_int32_bin(const std::string& path, size_t numel); /// Element OK if abs_err <= atol OR rel_err <= rtol (rel floored at /// |golden|=1e-6). Sets the reported maxima; true iff all elements pass. diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 66e7e55bafc..d0b972ec74e 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -38,7 +38,10 @@ def _materialize(spec) -> torch.Tensor: else: shape, gen = spec, "randn" if callable(gen): - return gen(shape).to(torch.float32) + _t = gen(shape) + return ( + _t.to(torch.int32) if not _t.is_floating_point() else _t.to(torch.float32) + ) if gen == "randn": return torch.randn(*shape) if gen == "ramp": @@ -97,7 +100,9 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: elif golden_dtype == "float64": # fp64 oracle (~1e-15); deepcopy keeps the original module fp32. double_module = copy.deepcopy(module).double() - golden = double_module(*[x.double() for x in inputs]) + golden = double_module( + *[x.double() if x.is_floating_point() else x for x in inputs] + ) else: golden = eager # gather/copy: fp64 is bit-identical, skip it golden_t = golden[out_index] if isinstance(golden, (tuple, list)) else golden @@ -119,8 +124,13 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: input_entries: list[dict] = [] for i, t in enumerate(inputs): rel = f"{case_id}.in{i}.bin" - _write_fp32(t, os.path.join(out_dir, rel)) - input_entries.append({"path": rel, "shape": list(t.shape), "dtype": "float32"}) + if t.dtype == torch.int32: + t.detach().cpu().numpy().astype(" tensors; for (const auto& in : e_.inputs) { const size_t n = numel(in.shape); - auto data = load_fp32_bin(in.path, n); - ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; std::vector sizes( in.shape.begin(), in.shape.end()); - tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); + if (in.dtype == "int32") { + auto data = load_int32_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); + } else { + auto data = load_fp32_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); + } } std::vector inputs; for (auto& t : tensors) { From a95b422bd64a242dc2036d94167de834c188c2be Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:14 -0700 Subject: [PATCH 02/41] [ExecuTorch][WebGPU] Add optimized gelu 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/20843 **Adds the `gelu` activation to the WebGPU backend.** GELU is the feed-forward (MLP) activation in the vision and VLM transformer blocks this backend targets — the Florence-2 DaViT encoder, BART, and the Whisper/Voxtral stacks — so it is on the critical path for delegating those models without graph breaks. **Problem** — The WebGPU backend had no `aten.gelu.default` kernel, so any model whose FFN uses GELU could not be fully delegated. PyTorch's default is the exact (erf) formulation, and the target encoders use it; the tanh approximation must also be supported. **Solution** - Before: `aten.gelu.default` was unsupported — no kernel, forcing a partition break around every GELU. - After: a single `gelu.wgsl` computes GELU on the GPU. The exact (erf) vs tanh choice is resolved by compiling the matching entry point (`main_erf` / `main_tanh`) into the pipeline, so each dispatch runs only its own formula — no per-invocation `select()` and no double-eval of both branches. **Implementation** - 1D dispatch with one thread per 4 elements: a vec4 body + scalar-tail idiom loads up to 4 elements via `select()`-guarded scalar reads, computes GELU as one `vec4` op, and scatters back only the in-bounds lanes — so arbitrary FFN widths that are not a multiple of 4 are handled without a separate remainder pass. - The erf path uses the Abramowitz & Stegun 7.1.26 rational approximation (max abs err ~1.5e-7); the tanh path uses the clamped `0.5*x*(1+tanh(...))` form. - The `approximate` arg (args[1]: "none" selects erf, anything else selects tanh) picks the entry point at build time through the shared `utils::make_compute_pipeline`; workgroup size comes from `utils::clamp_workgroup_size`, the dispatch count from `utils::compute_1d_workgroup_count`, and the uniform from `utils::make_uniform`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp` (its `gelu` handler; args[1] is the `approximate` string). Note the WebGPU kernel additionally implements the exact/erf path, which the Vulkan handler does not currently support. **Constraints** — fp32 only (both operands must be 4-byte aligned, enforced in the handler); input and output must have identical byte size; 1D dispatch only (throws if the workgroup count would exceed the 65535 cap). Element count need not be a multiple of 4 (scalar tail); static shape (no resize hook). Co-authored-with: Claude Code. ghstack-source-id: 405949686 @exported-using-ghexport Differential Revision: [D110836674](https://our.internmc.facebook.com/intern/diff/D110836674/) --- backends/webgpu/runtime/ops/gelu/Gelu.cpp | 100 +++++++++++++++++++ backends/webgpu/runtime/ops/gelu/gelu.wgsl | 65 ++++++++++++ backends/webgpu/runtime/ops/gelu/gelu_wgsl.h | 89 +++++++++++++++++ 3 files changed, 254 insertions(+) create mode 100644 backends/webgpu/runtime/ops/gelu/Gelu.cpp create mode 100644 backends/webgpu/runtime/ops/gelu/gelu.wgsl create mode 100644 backends/webgpu/runtime/ops/gelu/gelu_wgsl.h diff --git a/backends/webgpu/runtime/ops/gelu/Gelu.cpp b/backends/webgpu/runtime/ops/gelu/Gelu.cpp new file mode 100644 index 00000000000..943012515ff --- /dev/null +++ b/backends/webgpu/runtime/ops/gelu/Gelu.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 +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform buffer layout matching the WGSL Params struct; 16-byte aligned. +struct GeluParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; + +// aten.gelu.default args: [in, approximate, out] (mirrors Vulkan UnaryOp.cpp +// gelu — args[1] is the `approximate` string). approximate="none" selects the +// exact (erf) entry point; anything else (e.g. "tanh") selects the tanh +// approximation entry point. +void gelu_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int out_id = args.at(2); + const bool exact = graph.get_string(args.at(1)) == "none"; + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu"); + + uint32_t num_elements = + static_cast(out_tensor.nbytes / sizeof(float)); + + // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom). + uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); + uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_vec4_threads, wg_size, "gelu"); + + WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size); + + GeluParams params = {}; + params.num_elements = num_elements; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(GeluParams)); + graph.add_uniform_buffer_bytes(sizeof(GeluParams)); + + // input (read storage) + output (storage) + params. The exact/approximate + // choice is baked into the compiled pipeline via the entry point (mirrors + // onnxruntime's WebGPU EP), not a per-invocation select() — `main_tanh`/ + // `main_erf` each contain only their own formula, no double-eval. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kGeluWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(GeluParams)}, + }, + &wg_constant, + 1, + exact ? "main_erf" : "main_tanh"); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + // Drop our ref; the bind group keeps the uniform buffer alive until release. + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.gelu.default, gelu_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/gelu/gelu.wgsl b/backends/webgpu/runtime/ops/gelu/gelu.wgsl new file mode 100644 index 00000000000..4f7eb68bc96 --- /dev/null +++ b/backends/webgpu/runtime/ops/gelu/gelu.wgsl @@ -0,0 +1,65 @@ +@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; + +// Abramowitz & Stegun 7.1.26 erf approximation (max abs err ~1.5e-7). +fn erf_approx4(x: vec4) -> vec4 { + let s = sign(x); + let a = abs(x); + let t = 1.0 / (1.0 + 0.3275911 * a); + let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-a * a); + return s * y; +} + +fn gelu_tanh4(x: vec4) -> vec4 { + let inner = clamp(0.7978845608028654 * (x + 0.044715 * x * x * x), vec4(-10.0), vec4(10.0)); + return 0.5 * x * (1.0 + tanh(inner)); +} + +fn gelu_erf4(x: vec4) -> vec4 { + return 0.5 * x * (1.0 + erf_approx4(x * 0.7071067811865476)); +} + +// Vec4 body + scalar-tail idiom (num_elements isn't guaranteed % 4 == 0 for +// arbitrary FFN activation shapes): each thread gathers up to 4 elements via +// select()-guarded scalar loads (the out-of-bounds lanes read a WebGPU-spec- +// clamped, NOT zero, value — but that value is discarded by the same select() +// before use), computes GELU as one vec4 op, then scatters back only the +// in-bounds lanes. +@compute @workgroup_size(wg_size, 1, 1) +fn main_tanh(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_tanh4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main_erf(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_erf4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} diff --git a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h new file mode 100644 index 00000000000..6da12e229af --- /dev/null +++ b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h @@ -0,0 +1,89 @@ +/* + * 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 gelu.wgsl - DO NOT EDIT. +// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f +inline constexpr const char* kGeluWGSL = 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; + +// Abramowitz & Stegun 7.1.26 erf approximation (max abs err ~1.5e-7). +fn erf_approx4(x: vec4) -> vec4 { + let s = sign(x); + let a = abs(x); + let t = 1.0 / (1.0 + 0.3275911 * a); + let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-a * a); + return s * y; +} + +fn gelu_tanh4(x: vec4) -> vec4 { + let inner = clamp(0.7978845608028654 * (x + 0.044715 * x * x * x), vec4(-10.0), vec4(10.0)); + return 0.5 * x * (1.0 + tanh(inner)); +} + +fn gelu_erf4(x: vec4) -> vec4 { + return 0.5 * x * (1.0 + erf_approx4(x * 0.7071067811865476)); +} + +// Vec4 body + scalar-tail idiom (num_elements isn't guaranteed % 4 == 0 for +// arbitrary FFN activation shapes): each thread gathers up to 4 elements via +// select()-guarded scalar loads (the out-of-bounds lanes read a WebGPU-spec- +// clamped, NOT zero, value — but that value is discarded by the same select() +// before use), computes GELU as one vec4 op, then scatters back only the +// in-bounds lanes. +@compute @workgroup_size(wg_size, 1, 1) +fn main_tanh(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_tanh4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main_erf(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_erf4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} +)"; + +inline constexpr uint32_t kGeluWorkgroupSizeX = 64; +inline constexpr uint32_t kGeluWorkgroupSizeY = 1; +inline constexpr uint32_t kGeluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From a90388ea0b0a3dd41b105c204a95e61f5a930502 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:14 -0700 Subject: [PATCH 03/41] [ExecuTorch][WebGPU] gelu 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/20844 Splits the `gelu` op tests into their own diff, stacked directly above the `gelu` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_gelu.py` and registers the `gelu` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949721 @exported-using-ghexport Differential Revision: [D111072717](https://our.internmc.facebook.com/intern/diff/D111072717/) --- backends/webgpu/test/op_tests/cases.py | 39 +++++++++++++++++ backends/webgpu/test/ops/test_gelu.py | 59 ++++++++++++++++++++++++++ 2 files changed, 98 insertions(+) create mode 100644 backends/webgpu/test/ops/test_gelu.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 0216350a443..d054a39ef4f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -286,3 +286,42 @@ def _cat_suite() -> WebGPUTestSuite: ], golden_dtype="float32", # concatenation copies values; fp64 bit-identical ) + + +from executorch.backends.webgpu.test.ops.test_gelu import ( + _det_input as _gelu_det_input, + GeluModule, + N as _GELU_N, +) + + +def _gelu_full_range(_shape) -> torch.Tensor: + # Reuse the deterministic linspace(-6, 6) spanning negatives/zero/positives. + return _gelu_det_input() + + +@register_op_test("gelu") +def _gelu_suite() -> WebGPUTestSuite: + # erf ("none") is the Florence-2/BART + PyTorch default; tanh is the approx. + return WebGPUTestSuite( + module_factory=lambda approximate: GeluModule(approximate), + cases=[ + Case(name="erf_vec", construct={"approximate": "none"}, inputs=((M1,),)), + Case(name="erf_mat", construct={"approximate": "none"}, inputs=((M1, M2),)), + Case( + name="erf_rank3", + construct={"approximate": "none"}, + inputs=((S1, M1, M2),), + ), + Case( + name="tanh_mat", construct={"approximate": "tanh"}, inputs=((M1, M2),) + ), + Case( + name="erf_range", + construct={"approximate": "none"}, + inputs=(InputSpec(shape=(_GELU_N,), gen=_gelu_full_range),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_gelu.py b/backends/webgpu/test/ops/test_gelu.py new file mode 100644 index 00000000000..5fb6d73fb5d --- /dev/null +++ b/backends/webgpu/test/ops/test_gelu.py @@ -0,0 +1,59 @@ +# 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.gelu.default` module + input for the WebGPU op-test framework. + +`GeluModule`, `N`, and `_det_input` are imported by `cases.py` to drive the +declarative op-test suite. `GeluTest` is the export-delegation smoke test. The +`approximate` kwarg selects the exact (erf) path ("none", PyTorch's default and +the Florence-2/BART path) or the tanh approximation; the deterministic input +spans negatives, zero, and the saturation region. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# Input length; the deterministic input spans negatives, zero, and positives. +N = 64 + + +class GeluModule(torch.nn.Module): + def __init__(self, approximate: str = "none") -> None: + super().__init__() + self.approximate = approximate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(x, approximate=self.approximate) + + +def _det_input() -> torch.Tensor: + """Deterministic fp32 input spanning negatives, zero, and positives.""" + return torch.linspace(-6.0, 6.0, N, dtype=torch.float32) + + +def _export(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class GeluTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for approximate in ("none", "tanh"): + et = _export(GeluModule(approximate).eval(), _det_input()) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue( + found, f"Expected a VulkanBackend delegate (gelu {approximate})" + ) From e1de13ca3599466370513cd1e4adc94b614744b0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:15 -0700 Subject: [PATCH 04/41] [ExecuTorch][WebGPU] Add optimized layer_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/20845 **Adds `native_layer_norm` to the WebGPU backend.** LayerNorm is pervasive in the BART decoder and the Florence-2 DaViT vision encoder (and the Whisper/Voxtral hidden dims 768/1024/1280/3072), so it is required to delegate those transformer stacks end-to-end. **Problem** — The backend had no `aten.native_layer_norm.default` kernel. LayerNorm also needs a numerically robust mean/variance (a naive `E[x^2]-E[x]^2` cancels badly on large-mean/small-variance activations), and some graphs (the group_norm LN-reframe) pass `None` for weight and bias, so the affine step must be optional. **Solution** - Before: `aten.native_layer_norm.default` was unsupported. - After: `native_layer_norm.wgsl` computes per-row mean, rstd, and the normalized (optionally affine) output in a single kernel launch — one workgroup per row — writing all three outputs (`out`, `mean`, `rstd`) of the op's ValueList. **Implementation** - Single-pass, numerically robust mean+variance via Chan et al.'s parallel Welford: each of the 64 threads folds its strided `vec4` slice of the row into a running `(n, mean, M2)`, then a shared-memory tree reduction pairwise-merges the per-thread triples — no `E[x^2]-E[x]^2` cancellation. - `t_in`/`t_out`/`t_weight`/`t_bias` are viewed as `array>` over the row width for wide loads/stores; the second pass re-streams the row applying `(x-mean)*rstd` and the affine only when `has_affine == 1`. - Rows past the 65535 per-dimension ceiling are handled by a near-square 2D workgroup grid: `utils::compute_row_dispatch_grid` computes `count_x`/`count_y` and a `stride_x` override constant, and the shader decodes the flat row index as `wid.y * stride_x + wid.x`. - Weight/bias are optional: when either is `None` the handler binds dummy storage via `utils::make_optional_binding` and sets `has_affine == 0`; the pipeline is built with the shared `utils::make_compute_pipeline`, `epsilon` read via `utils::scalar_or`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp` (same arg layout and `[out, mean, rstd]` ValueList; the WebGPU kernel additionally supports the no-affine path, which the Vulkan handler rejects). **Constraints** — fp32 only (byte-size must equal `numel * sizeof(float)`); normalizes over the last dim only; the last dim (`row_width`) must be a multiple of 4 for the vec4 view (all in-scope hidden dims are); non-empty input; 2D dispatch grid for the row count. Co-authored-with: Claude Code. ghstack-source-id: 405949712 @exported-using-ghexport Differential Revision: [D110836681](https://our.internmc.facebook.com/intern/diff/D110836681/) --- .../ops/layer_norm/NativeLayerNorm.cpp | 183 ++++++++++++++++++ .../ops/layer_norm/native_layer_norm.wgsl | 116 +++++++++++ .../ops/layer_norm/native_layer_norm_wgsl.h | 140 ++++++++++++++ 3 files changed, 439 insertions(+) create mode 100644 backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp create mode 100644 backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl create mode 100644 backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h diff --git a/backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp b/backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp new file mode 100644 index 00000000000..7551d9dcdcf --- /dev/null +++ b/backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp @@ -0,0 +1,183 @@ +/* + * 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 LayerNormParams { + uint32_t num_rows; + uint32_t row_width; + float epsilon; + uint32_t has_affine; +}; +static_assert( + sizeof(LayerNormParams) == 16, + "LayerNormParams must be 16 bytes"); + +// aten.native_layer_norm.default args: [in, normalized_shape, weight, bias, +// eps, out] where out is a ValueList [out, mean, rstd] (mirrors Vulkan +// NativeLayerNorm.cpp). normalized_shape (args[1]) only constrains the last +// dim. +void native_layer_norm_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(2); + const int bias_id = args.at(3); + const int eps_id = args.at(4); + const int out_list_id = args.at(5); + + if (graph.get_value_type(out_list_id) != WebGPUGraph::ValueType::ValueList) { + throw std::runtime_error( + "WebGPU native_layer_norm: out is not a ValueList"); + } + const std::vector& outs = graph.get_value_list(out_list_id); + if (outs.size() != 3) { + throw std::runtime_error( + "WebGPU native_layer_norm: expected 3 outputs (out, mean, rstd)"); + } + const int out_id = outs.at(0); + const int mean_id = outs.at(1); + const int rstd_id = outs.at(2); + + WGPUDevice device = graph.device(); + + float epsilon = + utils::scalar_or(graph, eps_id, std::numeric_limits::epsilon()); + + const auto& in_tensor = graph.get_tensor(in_id); + if (in_tensor.dims.empty() || in_tensor.nbytes == 0) { + throw std::runtime_error("WebGPU native_layer_norm: empty input"); + } + const uint32_t row_width = static_cast(in_tensor.dims.back()); + if (row_width == 0) { + throw std::runtime_error("WebGPU native_layer_norm: zero row width"); + } + // The shader views t_in/t_out/t_weight/t_bias as vec4 over row_width; + // every model in scope (DaViT/BART/Whisper/Voxtral hidden dims 768/1024/ + // 1280/3072) is always %4==0. + if (row_width % 4 != 0) { + throw std::runtime_error( + "WebGPU native_layer_norm: row_width must be a multiple of 4"); + } + const uint64_t in_numel = + utils::check_fp32(in_tensor, "native_layer_norm", "input"); + const uint32_t num_rows = static_cast(in_numel / row_width); + if (num_rows == 0) { + throw std::runtime_error("WebGPU native_layer_norm: zero rows"); + } + + // Near-square 2D grid of workgroups (1 workgroup = 1 row) past the 65535 + // per-dimension ceiling; stride_x lets the shader decode a flat row index + // as workgroup_id.y * stride_x + workgroup_id.x. + utils::DispatchGrid grid = + utils::compute_row_dispatch_grid(device, num_rows, "native_layer_norm"); + + // weight/bias are optional: aten.native_layer_norm passes None for both when + // there is no affine (e.g. the group_norm LN-reframe). When absent, bind + // dummy storage buffers and gate the affine in the shader (has_affine == 0). + const bool has_affine = + graph.get_value_type(weight_id) == WebGPUGraph::ValueType::Tensor && + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + LayerNormParams params = {}; + params.num_rows = num_rows; + params.row_width = row_width; + params.epsilon = epsilon; + params.has_affine = has_affine ? 1u : 0u; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(LayerNormParams)); + graph.add_uniform_buffer_bytes(sizeof(LayerNormParams)); + + const auto& out_tensor = graph.get_tensor(out_id); + const auto& mean_tensor = graph.get_tensor(mean_id); + const auto& rstd_tensor = graph.get_tensor(rstd_id); + + utils::OptionalBinding weight = utils::make_optional_binding( + device, + has_affine, + has_affine ? graph.get_tensor(weight_id).buffer : nullptr, + has_affine ? graph.get_tensor(weight_id).nbytes : 0); + utils::OptionalBinding bias = utils::make_optional_binding( + device, + has_affine, + has_affine ? graph.get_tensor(bias_id).buffer : nullptr, + has_affine ? graph.get_tensor(bias_id).nbytes : 0); + + WGPUConstantEntry stride_const = {}; + stride_const.key = {"stride_x", WGPU_STRLEN}; + stride_const.value = static_cast(grid.stride_x); + + // out(rw,0), in(ro,1), weight(ro,2), bias(ro,3), mean(rw,4), rstd(rw,5), + // params(uniform,6). + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kNativeLayerNormWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes}, + {4, + WGPUBufferBindingType_Storage, + mean_tensor.buffer, + mean_tensor.nbytes}, + {5, + WGPUBufferBindingType_Storage, + rstd_tensor.buffer, + rstd_tensor.nbytes}, + {6, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LayerNormParams)}, + }, + &stride_const, + 1); + + static_assert( + kNativeLayerNormWorkgroupSizeX == 64, + "must match @workgroup_size and WG_SIZE in native_layer_norm.wgsl"); + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + + wgpuBufferRelease(uniform_buffer); + if (weight.owned_dummy != nullptr) { + wgpuBufferRelease(weight.owned_dummy); + } + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.native_layer_norm.default, native_layer_norm_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl new file mode 100644 index 00000000000..8e2656d7a65 --- /dev/null +++ b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl @@ -0,0 +1,116 @@ +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_in: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array>; +@group(0) @binding(4) var t_mean: array; +@group(0) @binding(5) var t_rstd: array; + +struct Params { + num_rows: u32, + row_width: u32, + epsilon: f32, + has_affine: u32, +} +@group(0) @binding(6) var params: Params; + +const WG_SIZE: u32 = 64u; + +override stride_x: u32 = 4294967295u; // = count_x; set by host for 2D-spill + +// Single-pass, numerically-robust mean+variance via Chan et al.'s parallel +// Welford: each thread folds its strided vec4 elements into a running +// (n, mean, M2), then this tree-reduces the WG_SIZE per-thread triples via +// pairwise merges — no naive E[x^2]-E[x]^2 cancellation risk, de-risked on +// CPU against large-mean/small-variance activations. +var shared_n: array; +var shared_mean: array; +var shared_m2: array; + +fn reduce_shared_welford(worker_id: u32) { + workgroupBarrier(); + var stride: u32 = WG_SIZE / 2u; + loop { + if (stride == 0u) { + break; + } + if (worker_id < stride) { + let na = shared_n[worker_id]; + let nb = shared_n[worker_id + stride]; + let n = na + nb; + if (n > 0.0) { + let delta = shared_mean[worker_id + stride] - shared_mean[worker_id]; + shared_mean[worker_id] = shared_mean[worker_id] + delta * (nb / n); + shared_m2[worker_id] = shared_m2[worker_id] + shared_m2[worker_id + stride] + + delta * delta * (na * nb / n); + shared_n[worker_id] = n; + } + } + workgroupBarrier(); + stride = stride >> 1u; + } +} + +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let row_idx = wid.y * stride_x + wid.x; + let worker_id = lid.x; + + if (row_idx >= params.num_rows) { + return; + } + + let row_width4 = params.row_width / 4u; + let base4 = row_idx * row_width4; + + var local_n: f32 = 0.0; + var local_mean: f32 = 0.0; + var local_m2: f32 = 0.0; + var x4: u32 = worker_id; + loop { + if (x4 >= row_width4) { + break; + } + let v4 = t_in[base4 + x4]; + for (var c: u32 = 0u; c < 4u; c = c + 1u) { + local_n = local_n + 1.0; + let v = v4[c]; + let delta = v - local_mean; + local_mean = local_mean + delta / local_n; + let delta2 = v - local_mean; + local_m2 = local_m2 + delta * delta2; + } + x4 = x4 + WG_SIZE; + } + shared_n[worker_id] = local_n; + shared_mean[worker_id] = local_mean; + shared_m2[worker_id] = local_m2; + reduce_shared_welford(worker_id); + let mean = shared_mean[0]; + let variance = shared_m2[0] / f32(params.row_width); + let rstd = inverseSqrt(variance + params.epsilon); + + if (worker_id == 0u) { + t_mean[row_idx] = mean; + t_rstd[row_idx] = rstd; + } + workgroupBarrier(); + + x4 = worker_id; + loop { + if (x4 >= row_width4) { + break; + } + let v4 = t_in[base4 + x4]; + let normed = (v4 - vec4(mean)) * rstd; + // weight/bias are optional (None when this is the group_norm LN-reframe); the + // dummy buffers are not bound then, so gate the affine on has_affine. + if (params.has_affine == 1u) { + t_out[base4 + x4] = normed * t_weight[x4] + t_bias[x4]; + } else { + t_out[base4 + x4] = normed; + } + x4 = x4 + WG_SIZE; + } +} diff --git a/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h new file mode 100644 index 00000000000..14495b7b40e --- /dev/null +++ b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_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 native_layer_norm.wgsl - DO NOT EDIT. +// wgsl-sha256: 135a274c6ee14e47f5de1f1255d88425f2b56cfbfa95c82ddb17aa82d725d88b +inline constexpr const char* kNativeLayerNormWGSL = R"( +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_in: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array>; +@group(0) @binding(4) var t_mean: array; +@group(0) @binding(5) var t_rstd: array; + +struct Params { + num_rows: u32, + row_width: u32, + epsilon: f32, + has_affine: u32, +} +@group(0) @binding(6) var params: Params; + +const WG_SIZE: u32 = 64u; + +override stride_x: u32 = 4294967295u; // = count_x; set by host for 2D-spill + +// Single-pass, numerically-robust mean+variance via Chan et al.'s parallel +// Welford: each thread folds its strided vec4 elements into a running +// (n, mean, M2), then this tree-reduces the WG_SIZE per-thread triples via +// pairwise merges — no naive E[x^2]-E[x]^2 cancellation risk, de-risked on +// CPU against large-mean/small-variance activations. +var shared_n: array; +var shared_mean: array; +var shared_m2: array; + +fn reduce_shared_welford(worker_id: u32) { + workgroupBarrier(); + var stride: u32 = WG_SIZE / 2u; + loop { + if (stride == 0u) { + break; + } + if (worker_id < stride) { + let na = shared_n[worker_id]; + let nb = shared_n[worker_id + stride]; + let n = na + nb; + if (n > 0.0) { + let delta = shared_mean[worker_id + stride] - shared_mean[worker_id]; + shared_mean[worker_id] = shared_mean[worker_id] + delta * (nb / n); + shared_m2[worker_id] = shared_m2[worker_id] + shared_m2[worker_id + stride] + + delta * delta * (na * nb / n); + shared_n[worker_id] = n; + } + } + workgroupBarrier(); + stride = stride >> 1u; + } +} + +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let row_idx = wid.y * stride_x + wid.x; + let worker_id = lid.x; + + if (row_idx >= params.num_rows) { + return; + } + + let row_width4 = params.row_width / 4u; + let base4 = row_idx * row_width4; + + var local_n: f32 = 0.0; + var local_mean: f32 = 0.0; + var local_m2: f32 = 0.0; + var x4: u32 = worker_id; + loop { + if (x4 >= row_width4) { + break; + } + let v4 = t_in[base4 + x4]; + for (var c: u32 = 0u; c < 4u; c = c + 1u) { + local_n = local_n + 1.0; + let v = v4[c]; + let delta = v - local_mean; + local_mean = local_mean + delta / local_n; + let delta2 = v - local_mean; + local_m2 = local_m2 + delta * delta2; + } + x4 = x4 + WG_SIZE; + } + shared_n[worker_id] = local_n; + shared_mean[worker_id] = local_mean; + shared_m2[worker_id] = local_m2; + reduce_shared_welford(worker_id); + let mean = shared_mean[0]; + let variance = shared_m2[0] / f32(params.row_width); + let rstd = inverseSqrt(variance + params.epsilon); + + if (worker_id == 0u) { + t_mean[row_idx] = mean; + t_rstd[row_idx] = rstd; + } + workgroupBarrier(); + + x4 = worker_id; + loop { + if (x4 >= row_width4) { + break; + } + let v4 = t_in[base4 + x4]; + let normed = (v4 - vec4(mean)) * rstd; + // weight/bias are optional (None when this is the group_norm LN-reframe); the + // dummy buffers are not bound then, so gate the affine on has_affine. + if (params.has_affine == 1u) { + t_out[base4 + x4] = normed * t_weight[x4] + t_bias[x4]; + } else { + t_out[base4 + x4] = normed; + } + x4 = x4 + WG_SIZE; + } +} +)"; + +inline constexpr uint32_t kNativeLayerNormWorkgroupSizeX = 64; +inline constexpr uint32_t kNativeLayerNormWorkgroupSizeY = 1; +inline constexpr uint32_t kNativeLayerNormWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 908e52cb856d8fc1750ac124ef50eff7b9ceb0e1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:15 -0700 Subject: [PATCH 05/41] [ExecuTorch][WebGPU] layer_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/20846 Splits the `layer_norm` op tests into their own diff, stacked directly above the `layer_norm` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_layer_norm.py` and registers the `layer_norm` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949713 @exported-using-ghexport Differential Revision: [D111072718](https://our.internmc.facebook.com/intern/diff/D111072718/) --- backends/webgpu/test/op_tests/cases.py | 49 +++++++++++++ backends/webgpu/test/ops/test_layer_norm.py | 78 +++++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 backends/webgpu/test/ops/test_layer_norm.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index d054a39ef4f..65ffdacff0c 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -325,3 +325,52 @@ def _gelu_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_layer_norm import ( + _ramp as _ln_ramp, + make_layer_norm, +) + + +@register_op_test("layer_norm") +def _layer_norm_suite() -> WebGPUTestSuite: + # LayerNorm over the last dim (BART + DaViT); affine + no-affine, widths + # below/equal/above the 64-wide workgroup reduction. + return WebGPUTestSuite( + module_factory=make_layer_norm, + cases=[ + Case( + name="affine_mat", + construct={"normalized_shape": 128}, + inputs=((4, 128),), + ), + Case( + name="affine_rank3", + construct={"normalized_shape": 768}, + inputs=((1, 16, 768),), + ), + Case( + name="no_affine", + construct={"normalized_shape": 128, "affine": False}, + inputs=((4, 128),), + ), + Case( + name="width_lt_wg", + construct={"normalized_shape": 32}, + inputs=((8, 32),), + ), + Case( + name="width_gt_wg", + construct={"normalized_shape": 132}, + inputs=((4, 132),), + ), + Case( + name="bart_hidden", + construct={"normalized_shape": 1024}, + inputs=(InputSpec(shape=(1, 8, 1024), gen=_ln_ramp),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_layer_norm.py b/backends/webgpu/test/ops/test_layer_norm.py new file mode 100644 index 00000000000..52640df03d4 --- /dev/null +++ b/backends/webgpu/test/ops/test_layer_norm.py @@ -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. + +"""`aten.native_layer_norm.default` module + inputs for the WebGPU op-test framework. + +`LayerNormModule`, `make_layer_norm`, and `_ramp` are imported by `cases.py` to +drive the declarative op-test suite. `LayerNormTest` is the export-delegation +smoke test. LayerNorm is pervasive in BART + the DaViT vision encoder +(Florence-2); both the affine (weight+bias) and the no-affine path are covered. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class LayerNormModule(torch.nn.Module): + """LayerNorm over the last dim; lowers to aten.native_layer_norm.default.""" + + def __init__(self, normalized_shape: int, affine: bool = True, eps: float = 1e-5): + super().__init__() + self.ln = torch.nn.LayerNorm( + normalized_shape, eps=eps, elementwise_affine=affine + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ln(x) + + +def make_layer_norm( + normalized_shape: int, affine: bool = True, eps: float = 1e-5 +) -> torch.nn.Module: + """Factory with deterministic non-trivial affine params (when affine).""" + m = LayerNormModule(normalized_shape, affine=affine, eps=eps) + if affine: + with torch.no_grad(): + m.ln.weight.copy_( + torch.linspace(0.5, 1.5, normalized_shape, dtype=torch.float32) + ) + m.ln.bias.copy_( + torch.linspace(-0.25, 0.25, normalized_shape, dtype=torch.float32) + ) + return m + + +def _ramp(shape) -> torch.Tensor: + """Deterministic linear ramp in [-1, 1] reshaped to `shape`.""" + n = 1 + for d in shape: + n *= d + return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) + + +def _export(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class LayerNormTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for affine in (True, False): + et = _export(make_layer_norm(128, affine=affine).eval(), _ramp((1, 4, 128))) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue( + found, f"Expected a VulkanBackend delegate (layer_norm affine={affine})" + ) From 7ebf5fd5b10d171f6c7c635e22776b76b181fa63 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:15 -0700 Subject: [PATCH 06/41] [ExecuTorch][WebGPU] linear_fp32 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/20848 Splits the `linear_fp32` op tests into their own diff, stacked directly above the `linear_fp32` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_linear_fp32.py` and registers the `linear_fp32` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949719 @exported-using-ghexport Differential Revision: [D111072710](https://our.internmc.facebook.com/intern/diff/D111072710/) --- backends/webgpu/test/op_tests/cases.py | 49 +++++++++++++ backends/webgpu/test/ops/test_linear_fp32.py | 77 ++++++++++++++++++++ 2 files changed, 126 insertions(+) create mode 100644 backends/webgpu/test/ops/test_linear_fp32.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 65ffdacff0c..bf96f694461 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -374,3 +374,52 @@ def _layer_norm_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_linear_fp32 import ( + _ramp as _lin_ramp, + make_linear, +) + + +@register_op_test("linear_fp32") +def _linear_fp32_suite() -> WebGPUTestSuite: + # fp32 linear (BART + DaViT projections); bias + no-bias, and shapes whose + # M*N exceeds the 65535 1D ceiling to exercise the 2D-dispatch spill. + return WebGPUTestSuite( + module_factory=make_linear, + cases=[ + Case( + name="bias_mat", + construct={"in_features": 64, "out_features": 32}, + inputs=((4, 64),), + ), + Case( + name="no_bias", + construct={"in_features": 64, "out_features": 32, "bias": False}, + inputs=((4, 64),), + ), + Case( + name="rank3", + construct={"in_features": 768, "out_features": 768}, + inputs=(InputSpec(shape=(1, 16, 768), gen=_lin_ramp),), + ), + Case( + name="tall_m", + construct={"in_features": 128, "out_features": 64}, + inputs=((256, 128),), + ), + Case( + name="bart_proj", + construct={"in_features": 1024, "out_features": 1024}, + inputs=(InputSpec(shape=(1, 8, 1024), gen=_lin_ramp),), + ), + Case( + name="odd_k", + construct={"in_features": 63, "out_features": 32}, + inputs=((4, 63),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_linear_fp32.py b/backends/webgpu/test/ops/test_linear_fp32.py new file mode 100644 index 00000000000..c635b4f9a3e --- /dev/null +++ b/backends/webgpu/test/ops/test_linear_fp32.py @@ -0,0 +1,77 @@ +# 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.linear.default` (fp32) module + inputs for the WebGPU op-test framework. + +`make_linear` and `_ramp` are imported by `cases.py` to drive the declarative +op-test suite. `LinearFp32Test` is the export-delegation smoke test. fp32 linear +is the projection used throughout BART + the DaViT vision encoder (Florence-2); +both the bias and no-bias paths are covered. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class LinearFp32Module(torch.nn.Module): + """fp32 linear; lowers to aten.linear.default with a prepacked [N, K] weight.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = True): + super().__init__() + self.fc = torch.nn.Linear(in_features, out_features, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc(x) + + +def make_linear( + in_features: int, out_features: int, bias: bool = True +) -> torch.nn.Module: + """Factory with deterministic weight (+ bias): a normalized ramp per row.""" + m = LinearFp32Module(in_features, out_features, bias=bias) + with torch.no_grad(): + w = torch.linspace( + -1.0, 1.0, out_features * in_features, dtype=torch.float32 + ).reshape(out_features, in_features) + m.fc.weight.copy_(w / in_features) + if bias: + m.fc.bias.copy_( + torch.linspace(-0.5, 0.5, out_features, dtype=torch.float32) + ) + return m + + +def _ramp(shape) -> torch.Tensor: + """Deterministic linear ramp in [-1, 1] reshaped to `shape`.""" + n = 1 + for d in shape: + n *= d + return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) + + +def _export(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class LinearFp32Test(unittest.TestCase): + def test_export_delegates(self) -> None: + for bias in (True, False): + et = _export(make_linear(64, 32, bias=bias).eval(), _ramp((4, 64))) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue( + found, f"Expected a VulkanBackend delegate (linear bias={bias})" + ) From 97d4c91470e667b14bbe822ca29e3ed14bb04943 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:16 -0700 Subject: [PATCH 07/41] [ExecuTorch][WebGPU] Add optimized conv2d 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/20849 **Adds a general fp32 2D convolution kernel — direct and transposed — to the WebGPU backend**, enabling the patch-embed / conv-stem and spatial-downsample layers of vision encoders (Florence-2 DaViT, SigLIP) to run GPU-accelerated in the browser. **Problem** — `aten.convolution.default` had no WebGPU kernel, so any model with a convolutional stem or downsample block (the DaViT patch-embed and its strided downsample convs) could not lower into the delegate and broke the graph. Both the direct and the transposed forms serialize to the same `aten.convolution.default` op (distinguished only by the `transposed` arg), so a single handler has to cover both. **Solution** - Before: no conv kernel — a convolution in the exported graph fell out of the delegate. - After: one `conv2d_impl` handler (registered once for `aten.convolution.default`) routes to three WGSL kernels by shape: `conv2d_vec4.wgsl` (vec4 fast path over input channels), `conv2d.wgsl` (scalar general path), and `conv_transpose2d.wgsl` (gather-form transposed conv). The `transposed` arg folds the transpose path into the same registration (a second `WEBGPU_REGISTER_OP(aten.convolution.default)` would be silently dropped), and among the non-transposed kernels the host picks vec4 vs scalar from the input-channels-per-group count. **Implementation** - Direct conv is one GPU thread per `(b, oc, oh, ow)` output element, looping over input-channel-per-group then `(kh, kw)`, with `continue` guards skipping out-of-bounds padded taps; `conv2d.wgsl` accumulates scalar `input[...] * weight[...]`. - The vec4 path is selected only when `icpg % 4 == 0` (`use_vec4` in the handler); because NCHW's channel dim has stride `IH*IW` (not memory-contiguous), `conv2d_vec4.wgsl` gathers four strided scalar loads into a `vec4` and uses `dot(in4, w4)` — a register-packing vec4, not a coalesced load — cutting the input-channel loop trip count 4x. A real RGB stem (`icpg=3`) correctly falls to the scalar path. - `conv_transpose2d.wgsl` implements the scatter-inversion as a gather: for each tap `(kh, kw)` an input row contributes only when `(oh + pH - kh*dH)` is divisible by `sH` and in range; weight is read in the un-flipped torch transposed layout `[IC, OC/groups, KH, KW]`. - All shape/stride/pad/dilation/groups constants ride in an 80-byte `ConvParams` uniform (16-byte-aligned); the handler parses `stride`/`padding`/`dilation` int-lists via `utils::parse_hw` (broadcasting a single value to both spatial dims). - Dispatch uses the shared adaptive-grid helper `utils::compute_dispatch_grid` (workgroup clamped to the device max, default 256) with a 2D spill past the 65535 per-dimension ceiling; the `stride_x` override constant lets the shader decode `i = gid.y*stride_x + gid.x`. Pipeline creation, the uniform upload, the grid override constants, and the optional-bias binding go through the shared `utils::make_compute_pipeline`, `utils::make_uniform`, `utils::make_grid_constants`, and `utils::make_optional_binding` helpers (a 4-byte dummy storage satisfies the bias binding when bias is `None`, gated in WGSL by `has_bias`). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Convolution.cpp` (its `Conv2dMethod` specializations `Depthwise` / `Pointwise` / `SlidingWindow` / `Transposed` collapse here into the vec4 / scalar / transpose routing). **Constraints** — fp32 only (bails if `nbytes != numel * sizeof(float)`); 4D input/weight/output in NCHW; `groups` must divide both `IC` and `OC` and `weight.dims[1]` must equal `IC/groups`; output element count must fit `uint32` for the (up to 2D) dispatch; non-zero `output_padding` is rejected on the non-transposed path (and `output_padding >= stride` on the transposed path); the fused `et_vk.conv_with_clamp` variant is not handled and would error clearly at load. Co-authored-with: Claude Code. ghstack-source-id: 405949715 @exported-using-ghexport Differential Revision: [D110836669](https://our.internmc.facebook.com/intern/diff/D110836669/) --- .../runtime/ops/et_vk_conv2d/Conv2d.cpp | 403 ++++++++++++++++++ .../runtime/ops/et_vk_conv2d/conv2d.wgsl | 81 ++++ .../runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl | 98 +++++ .../ops/et_vk_conv2d/conv2d_vec4_wgsl.h | 122 ++++++ .../runtime/ops/et_vk_conv2d/conv2d_wgsl.h | 105 +++++ .../ops/et_vk_conv2d/conv_transpose2d.wgsl | 94 ++++ .../ops/et_vk_conv2d/conv_transpose2d_wgsl.h | 118 +++++ 7 files changed, 1021 insertions(+) create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp new file mode 100644 index 00000000000..432fdb9de5b --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp @@ -0,0 +1,403 @@ +/* + * 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 ConvParams { + uint32_t B; + uint32_t IC; + uint32_t IH; + uint32_t IW; + uint32_t OC; + uint32_t OH; + uint32_t OW; + uint32_t KH; + uint32_t KW; + uint32_t sH; + uint32_t sW; + uint32_t pH; + uint32_t pW; + uint32_t dH; + uint32_t dW; + uint32_t groups; + uint32_t has_bias; + uint32_t _p0; + uint32_t _p1; + uint32_t _p2; +}; +static_assert( + sizeof(ConvParams) == 80, + "ConvParams must be 80 bytes (16-mult)"); + +// Transposed 2D convolution (gather form), folded into the convolution handler +// (C2: a 2nd WEBGPU_REGISTER_OP(aten.convolution.default) would be silently +// dropped). weight layout = torch convT [IC, OC/groups, KH, KW]. CPU-derisked +// vs torch.conv_transpose2d to fp64 round-off incl. non-square spatial + +// kernel. +void conv_transpose2d_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int stride_id = args.at(3); + const int padding_id = args.at(4); + const int dilation_id = args.at(5); + const int output_padding_id = args.at(7); + const int groups_id = args.at(8); + const int out_id = args.at(9); + + WGPUDevice device = graph.device(); + + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + if (in.dims.size() != 4 || weight.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error( + "WebGPU conv_transpose2d: expected 4D in/weight/out"); + } + + uint32_t sH, sW, pH, pW, dH, dW, opH, opW; + utils::parse_hw( + graph.get_int_list(stride_id), sH, sW, "conv_transpose2d", "stride"); + utils::parse_hw( + graph.get_int_list(padding_id), pH, pW, "conv_transpose2d", "padding"); + utils::parse_hw( + graph.get_int_list(dilation_id), dH, dW, "conv_transpose2d", "dilation"); + utils::parse_hw( + graph.get_int_list(output_padding_id), + opH, + opW, + "conv_transpose2d", + "output_padding"); + + const uint32_t B = static_cast(in.dims[0]); + const uint32_t IC = 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 groups = static_cast(graph.get_int(groups_id)); + // Transposed weight: [IC, OC/groups, KH, KW]. + const uint32_t KH = static_cast(weight.dims[2]); + const uint32_t KW = static_cast(weight.dims[3]); + const uint32_t OCpg = static_cast(weight.dims[1]); + const uint32_t OC = OCpg * groups; + + if (groups == 0 || IC % groups != 0 || OC % groups != 0) { + throw std::runtime_error( + "WebGPU conv_transpose2d: groups must divide IC/OC"); + } + if (static_cast(weight.dims[0]) != IC) { + throw std::runtime_error("WebGPU conv_transpose2d: weight dim0 != IC"); + } + if (sH == 0 || sW == 0) { + throw std::runtime_error("WebGPU conv_transpose2d: zero stride"); + } + // Guard a u32 underflow in (IH-1)/(IW-1) below, and a zero dilation/kernel + // that would mis-size the output geometry. + if (IH == 0 || IW == 0) { + throw std::runtime_error("WebGPU conv_transpose2d: zero input spatial dim"); + } + if (dH == 0 || dW == 0 || KH == 0 || KW == 0) { + throw std::runtime_error( + "WebGPU conv_transpose2d: zero dilation or kernel dim"); + } + if (opH >= sH || opW >= sW) { + throw std::runtime_error( + "WebGPU conv_transpose2d: output_padding >= stride"); + } + + // OH = (IH-1)*sH - 2*pH + dH*(KH-1) + output_padding + 1. + const int64_t oh_v = int64_t(IH - 1) * int64_t(sH) - 2 * int64_t(pH) + + int64_t(dH) * (int64_t(KH) - 1) + int64_t(opH) + 1; + const int64_t ow_v = int64_t(IW - 1) * int64_t(sW) - 2 * int64_t(pW) + + int64_t(dW) * (int64_t(KW) - 1) + int64_t(opW) + 1; + if (oh_v <= 0 || ow_v <= 0) { + throw std::runtime_error( + "WebGPU conv_transpose2d: invalid output geometry"); + } + const uint32_t OH = static_cast(oh_v); + const uint32_t OW = static_cast(ow_v); + + if (static_cast(out.dims[0]) != B || + static_cast(out.dims[1]) != OC || + static_cast(out.dims[2]) != OH || + static_cast(out.dims[3]) != OW) { + throw std::runtime_error("WebGPU conv_transpose2d: output shape mismatch"); + } + + uint64_t out_numel = utils::check_fp32(out, "conv_transpose2d", "output"); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + // 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 (FpnNeck @1008^2 out_numel exceeds it). 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, "conv_transpose2d"), + kConvTranspose2dWorkgroupSizeX, + "conv_transpose2d"); + + ConvParams params = {}; + params.B = B; + params.IC = IC; + params.IH = IH; + params.IW = IW; + params.OC = OC; + params.OH = OH; + params.OW = OW; + params.KH = KH; + params.KW = KW; + params.sH = sH; + params.sW = sW; + params.pH = pH; + params.pW = pW; + params.dH = dH; + params.dW = dW; + params.groups = groups; + params.has_bias = has_bias ? 1u : 0u; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvParams)); + + utils::OptionalBinding bias = 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, + kConvTranspose2dWGSL, + { + {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); + + wgpuBufferRelease(uniform_buffer); + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } +} + +// aten.convolution.default args: [input, weight, bias, stride, padding, +// dilation, transposed, output_padding, groups, out] (Vulkan Convolution.cpp). +// Direct 2D conv (non-transposed), NCHW fp32, general +// stride/pad/dilation/groups. The fused variant et_vk.conv_with_clamp.default +// is NOT handled (no post-conv activation in the SigLIP patch-embed); it would +// error clearly at load if hit. +void conv2d_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() != 10) { + throw std::runtime_error("WebGPU conv2d: expected 10 args (convolution)"); + } + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int stride_id = args.at(3); + const int padding_id = args.at(4); + const int dilation_id = args.at(5); + const int transposed_id = args.at(6); + const int groups_id = args.at(8); + const int out_id = args.at(9); + + WGPUDevice device = graph.device(); + + if (graph.get_bool(transposed_id)) { + // C2: fold — transposed conv is a distinct gather kernel, same + // registration. + conv_transpose2d_impl(graph, args); + return; + } + // output_padding (args[7]) is only meaningful for transposed conv; require it + // to be zero so a stale non-transposed export can't silently mis-size output. + for (int64_t v : graph.get_int_list(args.at(7))) { + if (v != 0) { + throw std::runtime_error( + "WebGPU conv2d: non-zero output_padding unsupported"); + } + } + + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + + if (in.dims.size() != 4 || weight.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU conv2d: expected 4D input/weight/out"); + } + + // stride/padding/dilation are int lists; PyTorch broadcasts a single value to + // both spatial dims (e.g. padding="valid" serializes as [0]). Accept size 1 + // (broadcast) or 2 (H, W). + uint32_t sH, sW, pH, pW, dH, dW; + utils::parse_hw(graph.get_int_list(stride_id), sH, sW, "conv2d", "stride"); + utils::parse_hw(graph.get_int_list(padding_id), pH, pW, "conv2d", "padding"); + utils::parse_hw( + graph.get_int_list(dilation_id), dH, dW, "conv2d", "dilation"); + + const uint32_t B = static_cast(in.dims[0]); + const uint32_t IC = 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 OC = static_cast(weight.dims[0]); + const uint32_t KH = static_cast(weight.dims[2]); + const uint32_t KW = static_cast(weight.dims[3]); + const uint32_t groups = static_cast(graph.get_int(groups_id)); + + if (groups == 0 || OC % groups != 0 || IC % groups != 0) { + throw std::runtime_error("WebGPU conv2d: groups must divide IC and OC"); + } + // weight is [OC, IC/groups, KH, KW]; verify the channel dim. + if (static_cast(weight.dims[1]) != IC / groups) { + throw std::runtime_error("WebGPU conv2d: weight in-channels != IC/groups"); + } + + if (sH == 0 || sW == 0) { + throw std::runtime_error("WebGPU conv2d: zero stride"); + } + if (dH == 0 || dW == 0 || KH == 0 || KW == 0) { + throw std::runtime_error("WebGPU conv2d: zero dilation or kernel dim"); + } + + // Compute in signed 64-bit: the numerator goes negative for invalid geometry + // (e.g. kernel larger than the padded input) and would wrap as unsigned. + const int64_t oh_num = + int64_t(IH) + 2 * int64_t(pH) - int64_t(dH) * (int64_t(KH) - 1) - 1; + const int64_t ow_num = + int64_t(IW) + 2 * int64_t(pW) - int64_t(dW) * (int64_t(KW) - 1) - 1; + if (oh_num < 0 || ow_num < 0) { + throw std::runtime_error( + "WebGPU conv2d: invalid geometry (kernel > input)"); + } + const uint32_t OH = static_cast(oh_num / int64_t(sH)) + 1; + const uint32_t OW = static_cast(ow_num / int64_t(sW)) + 1; + + // Validate against the serialized output tensor shape [B, OC, OH, OW]. + if (static_cast(out.dims[0]) != B || + static_cast(out.dims[1]) != OC || + static_cast(out.dims[2]) != OH || + static_cast(out.dims[3]) != OW) { + throw std::runtime_error("WebGPU conv2d: output shape mismatch"); + } + + uint64_t out_numel = utils::check_fp32(out, "conv2d", "output"); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + // NCHW's channel dim isn't memory-contiguous, so this is a register-packing + // vec4 (4 strided scalar loads gathered), not a coalesced load — still cuts + // the icg loop trip count 4x. Only valid when every group's input-channel + // count is a multiple of 4 (a real RGB stem's icpg=3 needs the scalar path). + 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"); + + ConvParams params = {}; + params.B = B; + params.IC = IC; + params.IH = IH; + params.IW = IW; + params.OC = OC; + params.OH = OH; + params.OW = OW; + params.KH = KH; + params.KW = KW; + params.sH = sH; + params.sW = sW; + params.pH = pH; + params.pW = pW; + params.dH = dH; + params.dW = dW; + params.groups = groups; + params.has_bias = has_bias ? 1u : 0u; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvParams)); + + // Dummy 4-byte storage to satisfy the bias binding when None (gated in WGSL). + utils::OptionalBinding bias = 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, + 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); + + wgpuBufferRelease(uniform_buffer); + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.convolution.default, conv2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl new file mode 100644 index 00000000000..d87c4d70005 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl @@ -0,0 +1,81 @@ +@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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution (non-transposed), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. Supports general stride/padding/dilation and +// groups. input [B,IC,IH,IW]; weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_c_base = (b * params.IC + ic) * params.IH; // *IW added per-row below + let w_c_base = (oc * icpg + icg) * params.KH; // *KW added per-row below + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + let w_row = (w_c_base + kh) * params.KW; + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + acc = acc + input[in_row + u32(iw)] * weight[w_row + kw]; + } + } + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl new file mode 100644 index 00000000000..48b38d5a3ad --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl @@ -0,0 +1,98 @@ +@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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution, vec4-gathered over IC (host selects this only when +// icpg % 4 == 0 — NCHW's channel dim has stride IH*IW, not contiguous, so this +// is a register-packing vec4 (4 strided scalar loads combined), not a +// coalesced-memory vec4 load; it still cuts the icg loop trip count 4x. ONE +// thread per (b, oc, oh, ow) output element. input [B,IC,IH,IW]; +// weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let icpg4 = icpg / 4u; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + let ic_stride = params.IH * params.IW; // channel stride in `input` + let w_ic_stride = params.KH * params.KW; // channel stride in `weight` + for (var icg4: u32 = 0u; icg4 < icpg4; icg4 = icg4 + 1u) { + let icg = icg4 * 4u; + let ic = ic0 + icg; + let in_c_base = b * params.IC * ic_stride + ic * ic_stride; + let w_c_base = oc * icpg * w_ic_stride + icg * w_ic_stride; + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let in_off = u32(ih) * params.IW + u32(iw); + let w_off = kh * params.KW + kw; + let in4 = vec4( + input[in_c_base + in_off], + input[in_c_base + ic_stride + in_off], + input[in_c_base + 2u * ic_stride + in_off], + input[in_c_base + 3u * ic_stride + in_off]); + let w4 = vec4( + weight[w_c_base + w_off], + weight[w_c_base + w_ic_stride + w_off], + weight[w_c_base + 2u * w_ic_stride + w_off], + weight[w_c_base + 3u * w_ic_stride + w_off]); + acc = acc + dot(in4, w4); + } + } + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h new file mode 100644 index 00000000000..19a124700bb --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h @@ -0,0 +1,122 @@ +/* + * 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_vec4.wgsl - DO NOT EDIT. +// wgsl-sha256: a60dbe62ac08849f9ed710700e6baaf3af9e618fb99852e6a82d9c5661440abe +inline constexpr const char* kConv2dVec4WGSL = 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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution, vec4-gathered over IC (host selects this only when +// icpg % 4 == 0 — NCHW's channel dim has stride IH*IW, not contiguous, so this +// is a register-packing vec4 (4 strided scalar loads combined), not a +// coalesced-memory vec4 load; it still cuts the icg loop trip count 4x. ONE +// thread per (b, oc, oh, ow) output element. input [B,IC,IH,IW]; +// weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let icpg4 = icpg / 4u; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + let ic_stride = params.IH * params.IW; // channel stride in `input` + let w_ic_stride = params.KH * params.KW; // channel stride in `weight` + for (var icg4: u32 = 0u; icg4 < icpg4; icg4 = icg4 + 1u) { + let icg = icg4 * 4u; + let ic = ic0 + icg; + let in_c_base = b * params.IC * ic_stride + ic * ic_stride; + let w_c_base = oc * icpg * w_ic_stride + icg * w_ic_stride; + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let in_off = u32(ih) * params.IW + u32(iw); + let w_off = kh * params.KW + kw; + let in4 = vec4( + input[in_c_base + in_off], + input[in_c_base + ic_stride + in_off], + input[in_c_base + 2u * ic_stride + in_off], + input[in_c_base + 3u * ic_stride + in_off]); + let w4 = vec4( + weight[w_c_base + w_off], + weight[w_c_base + w_ic_stride + w_off], + weight[w_c_base + 2u * w_ic_stride + w_off], + weight[w_c_base + 3u * w_ic_stride + w_off]); + acc = acc + dot(in4, w4); + } + } + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kConv2dVec4WorkgroupSizeX = 256; +inline constexpr uint32_t kConv2dVec4WorkgroupSizeY = 1; +inline constexpr uint32_t kConv2dVec4WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h new file mode 100644 index 00000000000..42e67d8bbb0 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h @@ -0,0 +1,105 @@ +/* + * 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.wgsl - DO NOT EDIT. +// wgsl-sha256: 8f675cbc2d43986af1c4592ee312984d1e9b02dee39f67442ae5d79539e4fd11 +inline constexpr const char* kConv2dWGSL = 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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution (non-transposed), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. Supports general stride/padding/dilation and +// groups. input [B,IC,IH,IW]; weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_c_base = (b * params.IC + ic) * params.IH; // *IW added per-row below + let w_c_base = (oc * icpg + icg) * params.KH; // *KW added per-row below + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + let w_row = (w_c_base + kh) * params.KW; + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + acc = acc + input[in_row + u32(iw)] * weight[w_row + kw]; + } + } + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kConv2dWorkgroupSizeX = 256; +inline constexpr uint32_t kConv2dWorkgroupSizeY = 1; +inline constexpr uint32_t kConv2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl new file mode 100644 index 00000000000..df2e64a15f9 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl @@ -0,0 +1,94 @@ +@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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Transposed 2D convolution (gather form), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. weight layout = torch convT [IC, OC/groups, +// KH, KW] (NOT flipped). For each kernel tap (kh,kw): an input row ih +// contributes iff (oh + pH - kh*dH) is divisible by sH and ih in range (the +// scatter-inversion). CPU-derisked vs torch.conv_transpose2d to fp64 round-off +// (/tmp/convtr_derisk.py), incl. non-square spatial + non-square kernel. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let oc_in_g = oc % ocpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let num_h = i32(oh) + i32(params.pH) - i32(kh) * i32(params.dH); + if (num_h % i32(params.sH) != 0) { + continue; + } + let ih = num_h / i32(params.sH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let num_w = i32(ow) + i32(params.pW) - i32(kw) * i32(params.dW); + if (num_w % i32(params.sW) != 0) { + continue; + } + let iw = num_w / i32(params.sW); + if (iw < 0 || iw >= iW) { + continue; + } + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_idx = + ((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw); + // weight [IC, OC/groups, KH, KW]: index (ic, oc_in_g, kh, kw) + let w_idx = + ((ic * ocpg + oc_in_g) * params.KH + kh) * params.KW + kw; + acc = acc + input[in_idx] * weight[w_idx]; + } + } + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h new file mode 100644 index 00000000000..471c7494bf0 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h @@ -0,0 +1,118 @@ +/* + * 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 conv_transpose2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 47bdb5dc2548e73c99e4f0f5b76c6fd996f9f3b0b5d6cfadbe675da806ab4a7b +inline constexpr const char* kConvTranspose2dWGSL = 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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Transposed 2D convolution (gather form), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. weight layout = torch convT [IC, OC/groups, +// KH, KW] (NOT flipped). For each kernel tap (kh,kw): an input row ih +// contributes iff (oh + pH - kh*dH) is divisible by sH and ih in range (the +// scatter-inversion). CPU-derisked vs torch.conv_transpose2d to fp64 round-off +// (/tmp/convtr_derisk.py), incl. non-square spatial + non-square kernel. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let oc_in_g = oc % ocpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let num_h = i32(oh) + i32(params.pH) - i32(kh) * i32(params.dH); + if (num_h % i32(params.sH) != 0) { + continue; + } + let ih = num_h / i32(params.sH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let num_w = i32(ow) + i32(params.pW) - i32(kw) * i32(params.dW); + if (num_w % i32(params.sW) != 0) { + continue; + } + let iw = num_w / i32(params.sW); + if (iw < 0 || iw >= iW) { + continue; + } + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_idx = + ((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw); + // weight [IC, OC/groups, KH, KW]: index (ic, oc_in_g, kh, kw) + let w_idx = + ((ic * ocpg + oc_in_g) * params.KH + kh) * params.KW + kw; + acc = acc + input[in_idx] * weight[w_idx]; + } + } + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kConvTranspose2dWorkgroupSizeX = 256; +inline constexpr uint32_t kConvTranspose2dWorkgroupSizeY = 1; +inline constexpr uint32_t kConvTranspose2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 805a52cc4ee7c67e1e237d9e92257c7f0014b15c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:16 -0700 Subject: [PATCH 08/41] [ExecuTorch][WebGPU] conv2d 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/20850 Splits the `conv2d` op tests into their own diff, stacked directly above the `conv2d` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_conv2d.py` and registers the `conv2d` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949711 @exported-using-ghexport Differential Revision: [D111072711](https://our.internmc.facebook.com/intern/diff/D111072711/) --- backends/webgpu/test/op_tests/cases.py | 59 +++++++++++++++ backends/webgpu/test/ops/test_conv2d.py | 95 +++++++++++++++++++++++++ 2 files changed, 154 insertions(+) create mode 100644 backends/webgpu/test/ops/test_conv2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index bf96f694461..f935b345a63 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -423,3 +423,62 @@ def _linear_fp32_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_conv2d import _chw_ramp, make_conv + + +@register_op_test("conv2d") +def _conv2d_suite() -> WebGPUTestSuite: + # DaViT patch-embed / downsample convs + conv_transpose2d (same registration, + # folded by the `transposed` arg). NCHW fp32. + return WebGPUTestSuite( + module_factory=make_conv, + cases=[ + Case( + name="conv3x3_pad1", + construct={"in_ch": 8, "out_ch": 16, "kernel": 3, "padding": 1}, + inputs=(InputSpec(shape=(1, 8, 16, 16), gen=_chw_ramp),), + ), + Case( + name="patch_embed", + construct={"in_ch": 3, "out_ch": 64, "kernel": 16, "stride": 16}, + inputs=(InputSpec(shape=(1, 3, 32, 32), gen=_chw_ramp),), + ), + Case( + name="strided", + construct={ + "in_ch": 3, + "out_ch": 8, + "kernel": 3, + "stride": 2, + "padding": 1, + }, + inputs=(InputSpec(shape=(1, 3, 16, 16), gen=_chw_ramp),), + ), + Case( + name="depthwise", + construct={ + "in_ch": 8, + "out_ch": 8, + "kernel": 3, + "padding": 1, + "groups": 8, + }, + inputs=(InputSpec(shape=(1, 8, 8, 8), gen=_chw_ramp),), + ), + Case( + name="transpose2x", + construct={ + "in_ch": 4, + "out_ch": 4, + "kernel": 2, + "stride": 2, + "transposed": True, + }, + inputs=(InputSpec(shape=(1, 4, 4, 4), gen=_chw_ramp),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_conv2d.py b/backends/webgpu/test/ops/test_conv2d.py new file mode 100644 index 00000000000..04ce2f2ef0a --- /dev/null +++ b/backends/webgpu/test/ops/test_conv2d.py @@ -0,0 +1,95 @@ +# 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.convolution.default` (conv2d + conv_transpose2d) module + inputs. + +`make_conv` and `_chw_ramp` are imported by `cases.py` to drive the declarative +op-test suite. `Conv2dTest` is the export-delegation smoke test. conv2d is the +DaViT patch-embed / downsample op in Florence-2; conv_transpose2d shares the +same `aten.convolution.default` registration (folded by the `transposed` arg). +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class Conv2dModule(torch.nn.Module): + def __init__( + self, + in_ch: int, + out_ch: int, + kernel: int, + stride: int = 1, + padding: int = 0, + groups: int = 1, + transposed: bool = False, + ): + super().__init__() + cls = torch.nn.ConvTranspose2d if transposed else torch.nn.Conv2d + self.conv = cls( + in_ch, out_ch, kernel, stride=stride, padding=padding, groups=groups + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +def make_conv( + in_ch: int, + out_ch: int, + kernel: int, + stride: int = 1, + padding: int = 0, + groups: int = 1, + transposed: bool = False, +) -> torch.nn.Module: + """Factory with deterministic small weights/bias for reproducible goldens.""" + m = Conv2dModule(in_ch, out_ch, kernel, stride, padding, groups, transposed) + with torch.no_grad(): + w = m.conv.weight + w.copy_( + torch.linspace(-1.0, 1.0, w.numel(), dtype=torch.float32).reshape(w.shape) + / (kernel * kernel) + ) + if m.conv.bias is not None: + b = m.conv.bias + b.copy_(torch.linspace(-0.5, 0.5, b.numel(), dtype=torch.float32)) + return m + + +def _chw_ramp(shape) -> torch.Tensor: + """Deterministic [N, C, H, W] ramp in [-1, 1].""" + n = 1 + for d in shape: + n *= d + return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) + + +def _export(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class Conv2dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + configs = [ + ("conv", make_conv(3, 8, 3, padding=1), (1, 3, 8, 8)), + ("transpose", make_conv(4, 4, 2, stride=2, transposed=True), (1, 4, 4, 4)), + ] + for name, model, shape in configs: + et = _export(model.eval(), _chw_ramp(shape)) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(found, f"Expected a VulkanBackend delegate (conv {name})") From 82b20ffdb44f4032c929d3a4ba566666c814f7f8 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:17 -0700 Subject: [PATCH 09/41] [ExecuTorch][WebGPU] Add optimized SDPA 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/20851 **Adds a non-causal fused scaled-dot-product-attention kernel to the WebGPU backend**, enabling the attention blocks of vision encoders (Florence-2 DaViT / SigLIP, SAM2) and non-causal cross-attention (BART) to run end-to-end on GPU. **Problem** — the delegate had no kernel for `et_vk.sdpa.default`, the plain non-causal attention `softmax(q @ kᵀ * scale + attn_mask) @ v` that the et_vk source transform plugs into every vision-encoder attention block. Without it the attention layers broke the graph. This is distinct from the causal KV-cache `sdpa_with_kv_cache`: no cache, an optional additive mask, and it must handle asymmetric sequence lengths (`S_q != S_kv`, e.g. a Hiera pooled query or cross-attention). **Solution** - Before: no non-causal fused-attention kernel — vision-encoder attention blocks could not lower into the delegate. - After: `et_vk_sdpa_impl` chains three compute dispatches over `[B, H, S, D]` (DSHB) row-major tensors — `et_vk_sdpa_qk.wgsl` computes the scaled `Q·Kᵀ` scores plus an optional additive mask, the reused `sdpa_softmax.wgsl` does the row-wise softmax, and `et_vk_sdpa_av.wgsl` computes `softmax · V` into the output. **Implementation** - QK phase (`et_vk_sdpa_qk.wgsl`): one GPU thread per `(b, h, s)` row of the `[B, H, S_q, S_kv]` attention-weight buffer; `q`/`k` are bound as `array>` over `D`, and the thread loops over `c` (key positions) and `d4` (`D/4`) accumulating `dot(q4, k4)`, then multiplies by `scale` and adds `mask[...]` when `has_mask`. Row count `B*H*S_q` stays well under the 1D dispatch limit for any ViT. - Softmax phase: reuses the existing `sdpa_softmax.wgsl` (one workgroup per row, hardcoded `@workgroup_size(64,1,1)`) dispatched on a near-square 2D workgroup grid past the 65535 ceiling; the shader recovers the flat row index from `@builtin(num_workgroups)`, so no override constant is needed. - AV phase (`et_vk_sdpa_av.wgsl`): one thread per `(b, h, s, d4)` computing a `vec4` of four output elements; `v`/`out` are bound as `array>` over `D` and the thread contracts over `c` scalar (`S_kv` is not guaranteed `% 4 == 0`), accumulating `sm[...] * v4`. - Two fp32 scratch buffers (`attn`, `softmax`, each `B*H*S_q*S_kv`) are allocated via `graph.create_scratch_buffer`; `scale` defaults to `1/sqrt(D)` when the arg is `None` or takes the `Double` value; the three uniforms are compact structs (`QkParams`/`AvParams` 32 bytes, `SoftmaxParams` 16 bytes). - Uses the shared runtime helpers: `utils::make_compute_pipeline`, `utils::make_uniform`, `utils::check_vec4_aligned` (guards `D % 4 == 0`), `utils::clamp_workgroup_size` + `utils::compute_1d_workgroup_count` (QK / AV grids), `utils::compute_2d_workgroup_count` (softmax grid), and `utils::make_optional_binding` (a 4-byte dummy satisfies the mask binding when absent; the shader never reads it under `has_mask == 0`). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/SDPA.cpp` `SDPAMode::FUSED` (general non-cache SDPA, `[B, H, S, D]` DSHB layout, optional additive `attn_mask`, optional `scale`, unpadded fp32 attention weights). **Constraints** — fp32 only (bails on `q`/`out` byte mismatch); non-causal (causality is expressed as a baked additive mask input, not a code path); `D % 4 == 0` for the vec4 QK/AV kernels (every model in scope uses `D=64` or `128`); `q` rank ≥ 3; `k.dims == v.dims`, `q`/`k`/`v` share `H` and `D` and all leading batch dims, and `out.dims == q.dims`; asymmetric `S_q != S_kv` is supported (reduces bit-identically to self-attention when equal); a supplied mask must be `[B, H, S_q, S_kv]` fp32. Co-authored-with: Claude Code. ghstack-source-id: 405949724 @exported-using-ghexport Differential Revision: [D110836679](https://our.internmc.facebook.com/intern/diff/D110836679/) --- .../runtime/ops/et_vk_sdpa/EtVkSdpa.cpp | 303 ++++++++++++++++++ .../runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl | 46 +++ .../ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h | 70 ++++ .../runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl | 55 ++++ .../ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h | 79 +++++ 5 files changed, 553 insertions(+) create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp new file mode 100644 index 00000000000..dbe6b03f30c --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp @@ -0,0 +1,303 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct QkParams { + uint32_t B; + uint32_t H; + uint32_t S_q; + uint32_t S_kv; + uint32_t D; + uint32_t has_mask; + uint32_t _pad0; + float scale; +}; +static_assert(sizeof(QkParams) == 32, "QkParams must be 32 bytes"); + +struct AvParams { + uint32_t B; + uint32_t H; + uint32_t S_q; + uint32_t S_kv; + uint32_t D; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; +}; +static_assert(sizeof(AvParams) == 32, "AvParams must be 32 bytes"); + +// Mirrors the Params struct in sdpa_softmax.wgsl (file-local in Sdpa.cpp, so +// re-declared here for the reuse). +struct SoftmaxParams { + uint32_t num_rows; + uint32_t row_width; + uint32_t _pad0; + uint32_t _pad1; +}; +static_assert(sizeof(SoftmaxParams) == 16, "SoftmaxParams must be 16 bytes"); + +// aten op: et_vk.sdpa.default. Args: [q, k, v, attn_mask, scale, out] (mirrors +// Vulkan fused_sdpa_impl, SDPA.cpp). Non-causal, no KV-cache; all tensors are +// DSHB [B, H, S, D], row-major. Three dispatches: QK (scaled, optional additive +// mask) -> softmax (reused sdpa_softmax.wgsl) -> AV. +void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() != 6) { + throw std::runtime_error("WebGPU et_vk.sdpa: expected 6 args"); + } + const int q_id = args.at(0); + const int k_id = args.at(1); + const int v_id = args.at(2); + const int mask_id = args.at(3); + const int scale_id = args.at(4); + const int out_id = args.at(5); + + WGPUDevice device = graph.device(); + + const auto& q = graph.get_tensor(q_id); + const auto& k = graph.get_tensor(k_id); + const auto& v = graph.get_tensor(v_id); + const auto& out = graph.get_tensor(out_id); + + if (q.dims.size() < 3) { + throw std::runtime_error("WebGPU et_vk.sdpa: q rank < 3"); + } + const int rank = static_cast(q.dims.size()); + const uint32_t D = static_cast(q.dims[rank - 1]); + const uint32_t S_q = static_cast(q.dims[rank - 2]); + const uint32_t H = static_cast(q.dims[rank - 3]); + if (D == 0 || S_q == 0 || H == 0) { + throw std::runtime_error("WebGPU et_vk.sdpa: zero D/S_q/H"); + } + // QK/AV kernels view q/k/v/out as vec4 over D; every model in scope + // (Whisper/Voxtral/DaViT/BART/Hiera) uses D=64 or 128, always %4==0. + utils::check_vec4_aligned(D, "et_vk.sdpa", "D"); + const uint64_t q_numel = utils::check_fp32(q, "et_vk.sdpa", "q"); + const uint32_t B = static_cast(q_numel / (uint64_t(H) * S_q * D)); + + // Asymmetric seq supported (S_q != S_kv, e.g. Hiera pooled query): k/v carry + // S_kv, q carries S_q; B/H/D must match across q/k/v. out is [B, H, S_q, D]. + // When S_q == S_kv this is plain self-attention (bit-identical to before). + if (k.dims != v.dims) { + throw std::runtime_error("WebGPU et_vk.sdpa: k/v shape mismatch"); + } + if (k.dims.size() != q.dims.size()) { + throw std::runtime_error("WebGPU et_vk.sdpa: q/k rank mismatch"); + } + if (static_cast(k.dims[rank - 1]) != D || + static_cast(k.dims[rank - 3]) != H) { + throw std::runtime_error("WebGPU et_vk.sdpa: q/k/v must share H and D"); + } + const uint32_t S_kv = static_cast(k.dims[rank - 2]); + if (S_kv == 0) { + throw std::runtime_error("WebGPU et_vk.sdpa: zero S_kv"); + } + // Leading (batch) dims must agree across q/k/v. + for (int d = 0; d < rank - 3; ++d) { + if (k.dims[d] != q.dims[d]) { + throw std::runtime_error("WebGPU et_vk.sdpa: q/k batch dims mismatch"); + } + } + // out must be [B, H, S_q, D] (same as q's shape). + if (out.dims != q.dims) { + throw std::runtime_error( + "WebGPU et_vk.sdpa: out shape must match q [B, H, S_q, D]"); + } + + const bool has_mask = + graph.get_value_type(mask_id) == WebGPUGraph::ValueType::Tensor; + if (has_mask) { + // The QK shader indexes mask as [B, H, S_q, S_kv] row-major; require it. + const auto& mask = graph.get_tensor(mask_id); + if (mask.nbytes != uint64_t(B) * H * S_q * S_kv * sizeof(float)) { + throw std::runtime_error( + "WebGPU et_vk.sdpa: attn_mask must be [B, H, S_q, S_kv] fp32"); + } + } + + float scale = 1.0f / std::sqrt(static_cast(D)); + const auto scale_type = graph.get_value_type(scale_id); + if (scale_type == WebGPUGraph::ValueType::Double) { + scale = static_cast(graph.get_double(scale_id)); + } else if (scale_type != WebGPUGraph::ValueType::Null) { + throw std::runtime_error("WebGPU et_vk.sdpa: scale must be Double or None"); + } + + const uint64_t num_rows = uint64_t(B) * H * S_q; // attn_weights rows + const uint64_t aw_numel = num_rows * S_kv; // [B, H, S_q, S_kv] + const uint64_t out_numel = uint64_t(B) * H * S_q * D; + 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"); + const uint32_t av_wg_size = + utils::clamp_workgroup_size(device, kEtVkSdpaAvWorkgroupSizeX); + const uint64_t out_numel4 = + out_numel / 4; // exact: D % 4 == 0 (checked above) + const uint32_t av_wg_count = utils::compute_1d_workgroup_count( + device, static_cast(out_numel4), av_wg_size, "et_vk_sdpa_av"); + // Near-square 2D grid of workgroups (1 workgroup = 1 row) past the 65535 + // per-dimension ceiling; sdpa_softmax.wgsl recovers the flat row index from + // @builtin(num_workgroups), so no override constant is needed here. + utils::WgCount softmax_grid = utils::compute_2d_workgroup_count( + device, + static_cast(num_rows), + /*workgroup_size=*/1, + "et_vk_sdpa_softmax"); + utils::check_fp32(out, "et_vk.sdpa", "out"); + + // NOTE: graph.create_scratch_buffer allocates a fresh buffer per call (no + // pooling), so a multi-layer graph holds 2 × num_layers × aw_bytes live, not + // 2 × max(aw_bytes) (e.g. S=1024 × 12 layers ≈ 1.1 GB) — a memory-headroom + // gap in the shared allocator, not this op; not a correctness issue. + 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) ---- + { + QkParams p = {}; + p.B = B; + p.H = H; + p.S_q = S_q; + p.S_kv = S_kv; + p.D = D; + p.has_mask = has_mask ? 1u : 0u; + p.scale = scale; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, &p, sizeof(QkParams)); + graph.add_uniform_buffer_bytes(sizeof(QkParams)); + + // 4-byte dummy storage to satisfy the mask binding when absent (shader + // never reads it under has_mask == 0). + utils::OptionalBinding mask = utils::make_optional_binding( + device, + has_mask, + has_mask ? graph.get_tensor(mask_id).buffer : nullptr, + has_mask ? graph.get_tensor(mask_id).nbytes : 0); + + WGPUConstantEntry wg_const = utils::make_wg_size_constant(qk_wg_size); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEtVkSdpaQkWGSL, + { + {0, WGPUBufferBindingType_Storage, attn_buf, aw_bytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, q.buffer, q.nbytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, k.buffer, k.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + mask.buffer, + mask.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(QkParams)}, + }, + &wg_const, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count}); + + wgpuBufferRelease(uniform_buffer); + if (mask.owned_dummy != nullptr) { + wgpuBufferRelease(mask.owned_dummy); + } + } + + // ---- Dispatch 2: softmax over the last dim (reuse sdpa_softmax.wgsl) ---- + { + SoftmaxParams p = {}; + p.num_rows = static_cast(num_rows); + p.row_width = S_kv; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, &p, sizeof(SoftmaxParams)); + graph.add_uniform_buffer_bytes(sizeof(SoftmaxParams)); + + // sdpa_softmax.wgsl hardcodes @workgroup_size(64,1,1); no override + // constant is needed to decode the near-square 2D grid. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSdpaSoftmaxWGSL, + { + {0, WGPUBufferBindingType_Storage, softmax_buf, aw_bytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, attn_buf, aw_bytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(SoftmaxParams)}, + }); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, softmax_grid.x, softmax_grid.y); + + wgpuBufferRelease(uniform_buffer); + } + + // ---- Dispatch 3: AV (one thread per (b,h,s,d4) vec4 output element) ---- + { + AvParams p = {}; + p.B = B; + p.H = H; + p.S_q = S_q; + p.S_kv = S_kv; + p.D = D; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, &p, sizeof(AvParams)); + graph.add_uniform_buffer_bytes(sizeof(AvParams)); + + WGPUConstantEntry wg_const = utils::make_wg_size_constant(av_wg_size); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEtVkSdpaAvWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, softmax_buf, aw_bytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, v.buffer, v.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(AvParams)}, + }, + &wg_const, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, av_wg_count}); + + wgpuBufferRelease(uniform_buffer); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.sdpa.default, et_vk_sdpa_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl new file mode 100644 index 00000000000..3ccd1d200a2 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl @@ -0,0 +1,46 @@ +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var sm: array; +@group(0) @binding(2) var v: array>; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, AV phase. out[b,h,s,d]=sum_c sm[b,h,s,c]*v[b,h,c,d]. +// DSHB layout, row-major: out/sm rows over S_q, v rows over S_kv; v/out viewed +// as vec4 over D (caller guarantees D % 4 == 0 for every model in scope). +// Supports asymmetric seq (S_q != S_kv); reduces to self-attention when +// S_q == S_kv. ONE thread per (b, h, s, d4) computing 4 output elements; the +// thread contracts over c (0..S_kv, scalar — S_kv isn't guaranteed % 4 == 0). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let d4_count = params.D / 4u; + let total = params.B * params.H * params.S_q * d4_count; + let i = gid.x; + if (i >= total) { + return; + } + let d4 = i % d4_count; + let s = (i / d4_count) % params.S_q; + let h = (i / (d4_count * params.S_q)) % params.H; + let b = i / (d4_count * params.S_q * params.H); + + let smbase = ((b * params.H + h) * params.S_q + s) * params.S_kv; + let vblock4 = (b * params.H + h) * params.S_kv; // first V row of this (b, h) + + var acc: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + acc = acc + sm[smbase + c] * v[(vblock4 + c) * d4_count + d4]; + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h new file mode 100644 index 00000000000..6f625437589 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h @@ -0,0 +1,70 @@ +/* + * 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_av.wgsl - DO NOT EDIT. +// wgsl-sha256: 0f91462226b666d2fc05381866caed709a6ef3ed2f2562addfe9a23ac81f173b +inline constexpr const char* kEtVkSdpaAvWGSL = R"( +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var sm: array; +@group(0) @binding(2) var v: array>; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, AV phase. out[b,h,s,d]=sum_c sm[b,h,s,c]*v[b,h,c,d]. +// DSHB layout, row-major: out/sm rows over S_q, v rows over S_kv; v/out viewed +// as vec4 over D (caller guarantees D % 4 == 0 for every model in scope). +// Supports asymmetric seq (S_q != S_kv); reduces to self-attention when +// S_q == S_kv. ONE thread per (b, h, s, d4) computing 4 output elements; the +// thread contracts over c (0..S_kv, scalar — S_kv isn't guaranteed % 4 == 0). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let d4_count = params.D / 4u; + let total = params.B * params.H * params.S_q * d4_count; + let i = gid.x; + if (i >= total) { + return; + } + let d4 = i % d4_count; + let s = (i / d4_count) % params.S_q; + let h = (i / (d4_count * params.S_q)) % params.H; + let b = i / (d4_count * params.S_q * params.H); + + let smbase = ((b * params.H + h) * params.S_q + s) * params.S_kv; + let vblock4 = (b * params.H + h) * params.S_kv; // first V row of this (b, h) + + var acc: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + acc = acc + sm[smbase + c] * v[(vblock4 + c) * d4_count + d4]; + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kEtVkSdpaAvWorkgroupSizeX = 64; +inline constexpr uint32_t kEtVkSdpaAvWorkgroupSizeY = 1; +inline constexpr uint32_t kEtVkSdpaAvWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl new file mode 100644 index 00000000000..f1c4f29ed8c --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl @@ -0,0 +1,55 @@ +@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]; q/k viewed as vec4 over D (caller guarantees +// D % 4 == 0 for every model in scope). Supports asymmetric seq (S_q != S_kv, +// e.g. Hiera pooled query); when S_q == S_kv this reduces to plain +// self-attention. ONE thread per (b, h, s) ROW of attn_weights +// [B, H, S_q, S_kv]; the thread loops over c (0..S_kv) and d4 (0..D/4). Row +// count = B*H*S_q stays well under the 65535 1D dispatch limit for any ViT. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let num_rows = params.B * params.H * params.S_q; + let row = gid.x; + if (row >= num_rows) { + return; + } + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let d4_count = params.D / 4u; + let qbase4 = ((b * params.H + h) * params.S_q + s) * d4_count; + let kblock4 = (b * params.H + h) * params.S_kv; // first K row of this (b, h) + let arow = ((b * params.H + h) * params.S_q + s) * params.S_kv; + + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + let kbase4 = (kblock4 + c) * d4_count; + var acc: f32 = 0.0; + for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) { + acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]); + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[arow + c]; + } + attn[arow + c] = acc; + } +} diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h new file mode 100644 index 00000000000..030e9dec24e --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h @@ -0,0 +1,79 @@ +/* + * 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.wgsl - DO NOT EDIT. +// wgsl-sha256: cdaa33d5b652395821e1d74222b113e72391b9d618358199abd0772cb821ef40 +inline constexpr const char* kEtVkSdpaQkWGSL = 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]; q/k viewed as vec4 over D (caller guarantees +// D % 4 == 0 for every model in scope). Supports asymmetric seq (S_q != S_kv, +// e.g. Hiera pooled query); when S_q == S_kv this reduces to plain +// self-attention. ONE thread per (b, h, s) ROW of attn_weights +// [B, H, S_q, S_kv]; the thread loops over c (0..S_kv) and d4 (0..D/4). Row +// count = B*H*S_q stays well under the 65535 1D dispatch limit for any ViT. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let num_rows = params.B * params.H * params.S_q; + let row = gid.x; + if (row >= num_rows) { + return; + } + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let d4_count = params.D / 4u; + let qbase4 = ((b * params.H + h) * params.S_q + s) * d4_count; + let kblock4 = (b * params.H + h) * params.S_kv; // first K row of this (b, h) + let arow = ((b * params.H + h) * params.S_q + s) * params.S_kv; + + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + let kbase4 = (kblock4 + c) * d4_count; + var acc: f32 = 0.0; + for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) { + acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]); + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[arow + c]; + } + attn[arow + c] = acc; + } +} +)"; + +inline constexpr uint32_t kEtVkSdpaQkWorkgroupSizeX = 64; +inline constexpr uint32_t kEtVkSdpaQkWorkgroupSizeY = 1; +inline constexpr uint32_t kEtVkSdpaQkWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 5a984c66dc9289258de46529539aafa7d7b048f6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:17 -0700 Subject: [PATCH 10/41] [ExecuTorch][WebGPU] SDPA 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/20852 Splits the `SDPA` op tests into their own diff, stacked directly above the `SDPA` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_et_vk_sdpa.py` and registers the `et_vk_sdpa` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949730 @exported-using-ghexport Differential Revision: [D111072726](https://our.internmc.facebook.com/intern/diff/D111072726/) --- backends/webgpu/test/op_tests/cases.py | 44 +++++++ backends/webgpu/test/ops/test_et_vk_sdpa.py | 131 ++++++++++++++++++++ 2 files changed, 175 insertions(+) create mode 100644 backends/webgpu/test/ops/test_et_vk_sdpa.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index f935b345a63..55b43851b61 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -482,3 +482,47 @@ def _conv2d_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_et_vk_sdpa import SdpaModule + + +def _sdpa_randn(shape): + g = torch.Generator().manual_seed(sum(int(x) for x in shape)) + return torch.randn(*shape, generator=g) + + +def _sdpa_mask(b, h, sq, skv): + g = torch.Generator().manual_seed(7) + return torch.randn(b, h, sq, skv, generator=g).clamp(-1.0, 0.0) + + +@register_op_test("et_vk_sdpa") +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. + def qkv(b, h, sq, skv, d): + return ( + InputSpec(shape=(b, h, sq, d), gen=_sdpa_randn), + InputSpec(shape=(b, h, skv, d), gen=_sdpa_randn), + InputSpec(shape=(b, h, skv, d), gen=_sdpa_randn), + ) + + return WebGPUTestSuite( + module_factory=lambda mask=None: SdpaModule(mask), + cases=[ + Case(name="selfattn_small", inputs=qkv(1, 4, 8, 8, 16)), + Case(name="selfattn_siglip", inputs=qkv(1, 12, 576, 576, 64)), + Case(name="asym_qpool", inputs=qkv(1, 8, 4, 16, 16)), + Case( + name="masked_bart", + construct={"mask": _sdpa_mask(1, 4, 8, 8)}, + inputs=qkv(1, 4, 8, 8, 16), + ), + Case(name="d128_voxtral", inputs=qkv(1, 4, 6, 6, 128)), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_et_vk_sdpa.py b/backends/webgpu/test/ops/test_et_vk_sdpa.py new file mode 100644 index 00000000000..c999fd01a1f --- /dev/null +++ b/backends/webgpu/test/ops/test_et_vk_sdpa.py @@ -0,0 +1,131 @@ +# 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. + +"""Non-causal fused SDPA (`et_vk.sdpa.default`) export + reference checks. + +This is the SAM2/SigLIP/DaViT/BART attention op (NOT the causal KV-cache +`sdpa_with_kv_cache` covered by `test/ops/test_sdpa.py`). The et_vk source +transform `_et_vk_sdpa_attn` plugs `torch.ops.et_vk.sdpa.default(q, k, v, +attn_mask, scale)` (q/k/v `[B, H, S, D]`) into every Florence-2 attention block; +it is plain non-causal attention `softmax(q @ kᵀ * scale + attn_mask) @ v`. + +`test_export_delegates` checks the op lowers into the VulkanBackend delegate. +`test_golden_matches_eager` checks the custom op's eager math matches +`F.scaled_dot_product_attention` (so the on-device golden can't be self- +fulfilling). On-device GPU numerics run via the op-test framework on a GPU rig. +""" + +import math +import unittest +from dataclasses import dataclass +from typing import Optional + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk.sdpa + +import torch +import torch.nn.functional as F +from executorch.backends.vulkan import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +NEG_INF = -1e30 + + +@dataclass(frozen=True) +class SdpaConfig: + name: str + b: int + h: int + s_q: int + s_kv: int + d: int + masked: bool = False + causal: bool = False + + +# Shapes from real SAM2/SigLIP/DaViT encoders + a cheap correctness case + an +# asymmetric (S_q != S_kv) pooled-query case + the BART causal-mask path. +CONFIGS = [ + SdpaConfig("selfattn_small", 1, 4, 8, 8, 16), + SdpaConfig("selfattn_siglip", 1, 12, 576, 576, 64), + SdpaConfig("asym_qpool", 1, 8, 4, 16, 16), + SdpaConfig("masked", 1, 4, 8, 8, 16, masked=True), + SdpaConfig("causal", 1, 4, 8, 8, 16, causal=True), +] + + +def _qkv(cfg: SdpaConfig): + g = torch.Generator().manual_seed(0) + q = torch.randn(cfg.b, cfg.h, cfg.s_q, cfg.d, generator=g) + k = torch.randn(cfg.b, cfg.h, cfg.s_kv, cfg.d, generator=g) + v = torch.randn(cfg.b, cfg.h, cfg.s_kv, cfg.d, generator=g) + return q, k, v + + +def _mask(cfg: SdpaConfig) -> Optional[torch.Tensor]: + if cfg.causal: + assert cfg.s_q == cfg.s_kv, "causal mask requires S_q == S_kv" + m = torch.triu( + torch.full((cfg.s_q, cfg.s_kv), NEG_INF, dtype=torch.float32), diagonal=1 + ) + return ( + m.reshape(1, 1, cfg.s_q, cfg.s_kv) + .expand(cfg.b, cfg.h, cfg.s_q, cfg.s_kv) + .contiguous() + ) + if not cfg.masked: + return None + g = torch.Generator().manual_seed(1) + return torch.randn(cfg.b, cfg.h, cfg.s_q, cfg.s_kv, generator=g).clamp(-1.0, 0.0) + + +class SdpaModule(torch.nn.Module): + """Wraps the registered et_vk.sdpa op. A baked mask is held as a buffer + (constant) so the partitioner prepacks it; the runner forwards only q/k/v.""" + + def __init__(self, mask: Optional[torch.Tensor] = None): + super().__init__() + # Always register as a buffer (even when None) so named_buffers() stays + # consistent across the masked and unmasked configs. + self.register_buffer("mask", mask) + + def forward(self, q, k, v): + return torch.ops.et_vk.sdpa.default(q, k, v, self.mask, None) + + +def _lower(cfg: SdpaConfig, q, k, v): + module = SdpaModule(_mask(cfg)).eval() + ep = torch.export.export(module, (q, k, v)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +class TestEtVkSdpa(unittest.TestCase): + def test_export_delegates(self) -> None: + for cfg in CONFIGS: + with self.subTest(config=cfg.name): + q, k, v = _qkv(cfg) + et = _lower(cfg, q, k, v).to_executorch() + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue( + found, f"Expected a VulkanBackend delegate (et_vk.sdpa {cfg.name})" + ) + + def test_golden_matches_eager(self) -> None: + # The custom op's eager math must equal F.scaled_dot_product_attention, + # so the on-device golden (computed from the op) is not self-fulfilling. + for cfg in CONFIGS: + with self.subTest(config=cfg.name): + q, k, v = _qkv(cfg) + mask = _mask(cfg) + scale = 1.0 / math.sqrt(cfg.d) + got = torch.ops.et_vk.sdpa.default(q, k, v, mask, scale) + ref = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, scale=scale + ) + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-3) From 1fe89ed2a3fb6b18ca21742f65a706fce43ee825 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:17 -0700 Subject: [PATCH 11/41] [ExecuTorch][WebGPU] embedding 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/20854 Splits the `embedding` op tests into their own diff, stacked directly above the `embedding` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_embedding.py` and registers the `embedding` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949731 @exported-using-ghexport Differential Revision: [D111072723](https://our.internmc.facebook.com/intern/diff/D111072723/) --- backends/webgpu/test/op_tests/cases.py | 36 ++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 55b43851b61..faeb12ac42d 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -526,3 +526,39 @@ def qkv(b, h, sq, skv, d): atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_embedding import EmbeddingModule + + +def _emb_idx_small(_shape): + return torch.tensor([0, 3, 15, 7], dtype=torch.long) + + +def _emb_idx_bart(_shape): + return torch.tensor([[1, 5, 1023, 0, 42]], dtype=torch.long) + + +@register_op_test("embedding") +def _embedding_suite() -> WebGPUTestSuite: + # fp32 token/pos embedding lookup (BART). int32 indices via the op-test + # framework's int-input path. + return WebGPUTestSuite( + module_factory=lambda num_embeddings, embed_dim: EmbeddingModule( + num_embeddings, embed_dim + ), + cases=[ + Case( + name="small", + construct={"num_embeddings": 16, "embed_dim": 8}, + inputs=(InputSpec(shape=(4,), gen=_emb_idx_small),), + ), + Case( + name="bart_tok", + construct={"num_embeddings": 1024, "embed_dim": 768}, + inputs=(InputSpec(shape=(1, 5), gen=_emb_idx_bart),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) From 2535a2fc88f9153320928af7fa30a91a116d284a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:18 -0700 Subject: [PATCH 12/41] [ExecuTorch][WebGPU] Add optimized addmm 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/20855 **fp32 `aten.addmm.default` now runs on the WebGPU backend as a shared-memory-tiled GEMM, so the fused bias-plus-matmul that HuggingFace `Linear` lowers to executes on-device.** In the Florence-2 BART and DaViT graphs the HF `nn.Linear` layers lower to `aten.addmm` (not `aten.linear`), making addmm the dominant dense GEMM in those stacks. **Problem** — The WebGPU backend had no `aten.addmm.default` kernel, so any graph whose linears lowered through addmm delegated at export but failed at runtime load. As with plain matmul, a naive per-output kernel would re-stream both operands from global memory with no reuse. **Solution** — Before: `aten.addmm.default` had no runtime kernel (load-time failure). After: the handler binds `self` (bias), `mat1 [M,K]`, `mat2 [K,N]`, and the output, and dispatches a tiled GEMM computing `out = beta*self + alpha*(mat1 @ mat2)`. Each workgroup cooperatively stages 32x32 slabs of `mat1` and `mat2` into workgroup shared memory, `workgroupBarrier`-syncs, and accumulates across `ceil(K/32)` k-tiles so the coalesced tile load is amortized over the whole output tile. The `beta*self + alpha*acc` epilogue broadcasts `self` from either `[N]` (the common HF bias case) or a full `[M,N]`. **Implementation** — Shares the tiled-GEMM approach with `linear_fp32`: `TILE = 32`, `RPT = 4`, `@workgroup_size(8, 8, 1)`, two `var` `a_sub`/`b_sub` arrays, each thread accumulating a 4x4 register sub-tile over `ceil(K/32)` tiles. It re-derives the B-side read for `mat2`'s actual `[K,N]` layout: `read_b` reads `t_mat2[krow*N + col]` (N contiguous), unlike `linear`'s transposed `[N,K]` weight read. No vec4 variant is provided: with `mat2 [K,N]` the N dimension is contiguous, so a vec4-over-K view would require a strided gather on the mat2 side that erodes the benefit, and the cooperative shared-memory tile load already closes the coalescing gap. Epilogue selects `t_self[r*N + c]` when `self_2d`, else the broadcast `t_self[c]`, and writes `beta*self_val + alpha*acc[ir][ic]`. `beta`/`alpha` are read with the shared `utils::scalar_or` (Double/Int/Null-permissive, default `1.0`). Dispatch is a 2D grid `ceil(N/32) x ceil(M/32)`, throwing if either dimension exceeds 65535. Uses the shared `utils::make_compute_pipeline` and `utils::make_uniform` (32-byte `AddmmParams {M,N,K,self_2d,beta,alpha,_pad[2]}`). Validates non-null buffers, 2D `mat1`/`mat2`/`out`, `mat2` shape `== [K,N]`, and `self` broadcastable from `[N]` or `[M,N]`. **Constraints** — fp32 only (output byte-size check); `mat1`/`mat2`/`out` must be 2D with `mat2` exactly `[K,N]`; `self` must be `[N]` or `[M,N]`; scalar-only tiled path (no vec4); 2D dispatch capped at 65535 workgroups per dimension; fixed `@workgroup_size(8, 8, 1)`. Co-authored-with: Claude Code. ghstack-source-id: 405949737 @exported-using-ghexport Differential Revision: [D110836673](https://our.internmc.facebook.com/intern/diff/D110836673/) --- backends/webgpu/runtime/ops/addmm/Addmm.cpp | 139 ++++++++++++++++++ .../webgpu/runtime/ops/addmm/addmm_tiled.wgsl | 99 +++++++++++++ .../runtime/ops/addmm/addmm_tiled_wgsl.h | 123 ++++++++++++++++ 3 files changed, 361 insertions(+) create mode 100644 backends/webgpu/runtime/ops/addmm/Addmm.cpp create mode 100644 backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl create mode 100644 backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h diff --git a/backends/webgpu/runtime/ops/addmm/Addmm.cpp b/backends/webgpu/runtime/ops/addmm/Addmm.cpp new file mode 100644 index 00000000000..06522ba9481 --- /dev/null +++ b/backends/webgpu/runtime/ops/addmm/Addmm.cpp @@ -0,0 +1,139 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct AddmmParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t self_2d; + float beta; + float alpha; + uint32_t _pad[2]; +}; +static_assert(sizeof(AddmmParams) == 32, "AddmmParams must be 32 bytes"); + +constexpr uint32_t kTile = 32u; + +// aten.addmm.default args: [self, mat1, mat2, beta, alpha, out]. +// out = beta*self + alpha*(mat1 @ mat2); mat1 [M,K], mat2 [K,N], self [N] or +// [M,N] (HF Linear lowers to addmm with a [N] bias). Shared-memory-tiled GEMM, +// re-derived for mat2's [K,N] layout (mirrors the sibling `linear` op's +// linear_tiled.wgsl skeleton, NOT its [N,K]-shaped read_b). No vec4 variant: +// unlike linear's weight [N,K] (K contiguous, vec4-over-K natural on both +// sides), mat2 [K,N] has N contiguous — vec4-over-K would need a strided +// gather on the mat2 side, eroding the benefit; the tiled kernel alone already +// solves the coalescing gap via the cooperative shared-memory tile load. +void addmm_impl(WebGPUGraph& graph, const std::vector& args) { + const int self_id = args.at(0); + const int mat1_id = args.at(1); + const int mat2_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& self_t = graph.get_tensor(self_id); + const auto& mat1 = graph.get_tensor(mat1_id); + const auto& mat2 = graph.get_tensor(mat2_id); + const auto& out = graph.get_tensor(out_id); + + if (self_t.buffer == nullptr || mat1.buffer == nullptr || + mat2.buffer == nullptr || out.buffer == nullptr) { + throw std::runtime_error("WebGPU addmm: null buffer binding"); + } + if (mat1.dims.size() != 2 || mat2.dims.size() != 2 || out.dims.size() != 2) { + throw std::runtime_error("WebGPU addmm: expected 2D mat1/mat2/out"); + } + + const uint32_t M = static_cast(out.dims[0]); + const uint32_t N = static_cast(out.dims[1]); + const uint32_t K = static_cast(mat1.dims[1]); + if (M == 0 || N == 0 || K == 0) { + throw std::runtime_error("WebGPU addmm: zero M/N/K"); + } + if (static_cast(mat2.dims[0]) != K || + static_cast(mat2.dims[1]) != N) { + throw std::runtime_error("WebGPU addmm: mat2 shape != [K, N]"); + } + utils::check_fp32(out, "addmm", "output"); + + // self is either a [N] bias (broadcast over rows) or a full [M, N]; validate + // the actual dims, not just the element count (a coincidental-numel tensor + // with the wrong rank/shape would otherwise be silently accepted). + const bool self_2d = self_t.dims.size() == 2 && + static_cast(self_t.dims[0]) == M && + static_cast(self_t.dims[1]) == N; + const bool self_1d = + self_t.dims.size() == 1 && static_cast(self_t.dims[0]) == N; + if (!self_2d && !self_1d) { + throw std::runtime_error("WebGPU addmm: self must be [N] or [M, N]"); + } + + const float beta = utils::scalar_or(graph, args.at(3), 1.0f); + const float alpha = utils::scalar_or(graph, args.at(4), 1.0f); + + // A genuinely-2D tile dispatch doesn't need the 1D-flat stride_x recovery + // trick; throw before any allocation if it can't fit. + utils::WgCount tile_grid = + utils::compute_tile_grid_2d(device, N, M, kTile, "addmm"); + + AddmmParams params = {}; + params.M = M; + params.N = N; + params.K = K; + params.self_2d = self_2d ? 1u : 0u; + params.beta = beta; + params.alpha = alpha; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(AddmmParams)); + graph.add_uniform_buffer_bytes(sizeof(AddmmParams)); + + // Tiled kernel has a fixed @workgroup_size(8, 8, 1) — no override constant. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kAddmmTiledWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + self_t.buffer, + self_t.nbytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, mat1.buffer, mat1.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, mat2.buffer, mat2.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(AddmmParams)}, + }); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, tile_grid.x, tile_grid.y); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.addmm.default, addmm_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl b/backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl new file mode 100644 index 00000000000..a2e9a24ecdc --- /dev/null +++ b/backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl @@ -0,0 +1,99 @@ +struct Params { + M: u32, + N: u32, + K: u32, + self_2d: u32, + beta: f32, + alpha: f32, + _pad0: u32, + _pad1: u32, +}; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_self: array; +@group(0) @binding(2) var t_mat1: array; +@group(0) @binding(3) var t_mat2: array; +@group(0) @binding(4) var params: Params; + +// Shared-memory tiled addmm (same skeleton as the sibling `linear` op's +// linear_tiled.wgsl), re-derived for mat2's actual [K,N] layout (NOT [N,K] +// like linear's weight — mat2[k,n] is contiguous over n, so the tile-load's +// per-thread-row reads are cacheline-coalesced, unlike linear's weight; still +// RPT=4-float-strided across threads within a row, not 1-float-strided). +// out = beta*self + alpha*(mat1 @ mat2); mat1 [M,K], mat2 [K,N]. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(row: u32, col: u32) -> f32 { + if (row < params.M && col < params.K) { + return t_mat1[row * params.K + col]; + } + return 0.0; +} + +fn read_b(krow: u32, col: u32) -> f32 { + if (krow < params.K && col < params.N) { + return t_mat2[krow * params.N + col]; + } + return 0.0; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + 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 = (params.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 r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + var self_val: f32; + if (params.self_2d == 1u) { + self_val = t_self[r * params.N + c]; + } else { + self_val = t_self[c]; + } + t_out[r * params.N + c] = params.beta * self_val + params.alpha * acc[ir][ic]; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h b/backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h new file mode 100644 index 00000000000..47e5ac60527 --- /dev/null +++ b/backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h @@ -0,0 +1,123 @@ +/* + * 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 addmm_tiled.wgsl - DO NOT EDIT. +// wgsl-sha256: 70d1e4fe13aea0d042fa657c8a9d6da81f3ea7fa8e35de98209a5a140308d364 +inline constexpr const char* kAddmmTiledWGSL = R"( +struct Params { + M: u32, + N: u32, + K: u32, + self_2d: u32, + beta: f32, + alpha: f32, + _pad0: u32, + _pad1: u32, +}; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_self: array; +@group(0) @binding(2) var t_mat1: array; +@group(0) @binding(3) var t_mat2: array; +@group(0) @binding(4) var params: Params; + +// Shared-memory tiled addmm (same skeleton as the sibling `linear` op's +// linear_tiled.wgsl), re-derived for mat2's actual [K,N] layout (NOT [N,K] +// like linear's weight — mat2[k,n] is contiguous over n, so the tile-load's +// per-thread-row reads are cacheline-coalesced, unlike linear's weight; still +// RPT=4-float-strided across threads within a row, not 1-float-strided). +// out = beta*self + alpha*(mat1 @ mat2); mat1 [M,K], mat2 [K,N]. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(row: u32, col: u32) -> f32 { + if (row < params.M && col < params.K) { + return t_mat1[row * params.K + col]; + } + return 0.0; +} + +fn read_b(krow: u32, col: u32) -> f32 { + if (krow < params.K && col < params.N) { + return t_mat2[krow * params.N + col]; + } + return 0.0; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + 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 = (params.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 r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + var self_val: f32; + if (params.self_2d == 1u) { + self_val = t_self[r * params.N + c]; + } else { + self_val = t_self[c]; + } + t_out[r * params.N + c] = params.beta * self_val + params.alpha * acc[ir][ic]; + } + } + } +} +)"; + +inline constexpr uint32_t kAddmmTiledWorkgroupSizeX = 8; +inline constexpr uint32_t kAddmmTiledWorkgroupSizeY = 8; +inline constexpr uint32_t kAddmmTiledWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9bfc36b2580ba66e2cb1ff90314c4efffcaec5ec Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:18 -0700 Subject: [PATCH 13/41] [ExecuTorch][WebGPU] addmm 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/20856 Splits the `addmm` op tests into their own diff, stacked directly above the `addmm` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_addmm.py` and registers the `addmm` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949733 @exported-using-ghexport Differential Revision: [D111072709](https://our.internmc.facebook.com/intern/diff/D111072709/) --- backends/webgpu/test/op_tests/cases.py | 42 ++++++++++++++++++++++++++ backends/webgpu/test/ops/test_addmm.py | 28 +++++++++++++++++ 2 files changed, 70 insertions(+) create mode 100644 backends/webgpu/test/ops/test_addmm.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index faeb12ac42d..ef79748bba7 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -562,3 +562,45 @@ def _embedding_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_addmm import ( + _randn as _addmm_randn, + AddmmModule, +) + + +@register_op_test("addmm") +def _addmm_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=lambda n: AddmmModule(n), + cases=[ + Case( + name="small", + construct={"n": 32}, + inputs=( + InputSpec(shape=(4, 16), gen=_addmm_randn), + InputSpec(shape=(16, 32), gen=_addmm_randn), + ), + ), + Case( + name="bart", + construct={"n": 768}, + inputs=( + InputSpec(shape=(16, 768), gen=_addmm_randn), + InputSpec(shape=(768, 768), gen=_addmm_randn), + ), + ), + Case( + name="odd_k", + construct={"n": 32}, + inputs=( + InputSpec(shape=(4, 15), gen=_addmm_randn), + InputSpec(shape=(15, 32), gen=_addmm_randn), + ), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_addmm.py b/backends/webgpu/test/ops/test_addmm.py new file mode 100644 index 00000000000..118ccaa1b92 --- /dev/null +++ b/backends/webgpu/test/ops/test_addmm.py @@ -0,0 +1,28 @@ +# 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.addmm.default` module for the WebGPU op-test framework. + +HF Linear lowers to addmm (not aten.linear), so this is the dominant GEMM in the +Florence-2 BART + DaViT graphs: out = beta*self + alpha*(mat1 @ mat2). +""" + +import torch + + +class AddmmModule(torch.nn.Module): + def __init__(self, n: int, beta: float = 1.0, alpha: float = 1.0): + super().__init__() + self.bias = torch.nn.Parameter(torch.linspace(-0.5, 0.5, n)) + self.beta, self.alpha = beta, alpha + + def forward(self, mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor: + return torch.addmm(self.bias, mat1, mat2, beta=self.beta, alpha=self.alpha) + + +def _randn(shape) -> torch.Tensor: + g = torch.Generator().manual_seed(sum(int(x) for x in shape)) + return torch.randn(*shape, generator=g) * 0.1 From dbca29faa5ac52661d40b9ccdf413b166374de50 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:19 -0700 Subject: [PATCH 14/41] [ExecuTorch][WebGPU] Add constant_pad_nd 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/20857 **Adds `aten.constant_pad_nd.default` to the WebGPU backend, unblocking the DaViT window-padding path in vision models.** **Problem** — the backend had no `constant_pad_nd` handler, so any graph that pads a tensor (e.g. DaViT's window partitioning) could not fully delegate to WebGPU and threw at runtime. **Solution** — a single gather/fill compute kernel: Before — no handler; `aten.constant_pad_nd.default` unsupported at runtime. After — one thread per output element gathers the source element when its coordinates land inside the input, otherwise writes the constant fill `value`. **Implementation**: - The handler right-aligns the (rank 1..4) dims into fixed `vec4` params (`out_dims`, `in_dims`, `left`); leading slots get extent 1 / pad 0, so the WGSL is rank-agnostic and always iterates 4 dims. - The `pad` `IntList` is reversed-dim (innermost-first `(left, right)` pairs); the handler expands it to per-dim `left`/`right`, then validates `out.dims[d] == in.dims[d] + left[d] + right[d]` before any buffer allocation (loud-fail, no leak-on-throw). - The kernel decodes each output element's 4D coords (last dim fastest), subtracts each dim's `left` pad as an unsigned wrapping subtract (a negative coord wraps to a huge value and is rejected by the `< in_dims` bound check); if all four coords are in-bounds it copies `inp[flat_in]`, else it writes `value` — a pure copy/fill, so bit-exact. - The fill `value` is read via `utils::scalar_or` (a `Scalar` may serialize as `Int` or `Double`), defaulting to `0`. - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid` (workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 workgroup-count ceiling; the `stride_x` override lets the shader decode `i = gid.y*stride_x + gid.x`). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Pad.cpp` (same reversed-dim `(before, after)` pad convention and `constant_pad_nd` resize logic). **Constraints** — fp32 only (`nbytes == numel*4` guard); rank 1..4; `pad` must be even-length and no longer than the rank; the output element count must fit `u32` (`<= 2^32`). Co-authored-with: Claude Code. ghstack-source-id: 405949740 @exported-using-ghexport Differential Revision: [D110836672](https://our.internmc.facebook.com/intern/diff/D110836672/) --- .../ops/constant_pad_nd/ConstantPadNd.cpp | 156 ++++++++++++++++++ .../ops/constant_pad_nd/constant_pad_nd.wgsl | 63 +++++++ .../constant_pad_nd/constant_pad_nd_wgsl.h | 87 ++++++++++ 3 files changed, 306 insertions(+) create mode 100644 backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp create mode 100644 backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl create mode 100644 backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp b/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp new file mode 100644 index 00000000000..2dc5a0907f1 --- /dev/null +++ b/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp @@ -0,0 +1,156 @@ +/* + * 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 PadParams { + uint32_t out_dims[4]; + uint32_t in_dims[4]; + uint32_t left[4]; + uint32_t out_numel; + float value; + uint32_t _p0; + uint32_t _p1; +}; +static_assert(sizeof(PadParams) == 64, "PadParams must be 64 bytes"); + +// `pad` is reversed-dim (last-dim-first pairs); output right-aligned into 4D. +void constant_pad_nd_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 3) { + throw std::runtime_error("WebGPU constant_pad_nd: expected >=3 args"); + } + const int in_id = args.at(0); + const int pad_id = args.at(1); + 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); + + const size_t nd = in.dims.size(); + if (nd == 0 || nd > 4) { + throw std::runtime_error("WebGPU constant_pad_nd: rank must be 1..4"); + } + if (out.dims.size() != nd) { + throw std::runtime_error("WebGPU constant_pad_nd: in/out rank mismatch"); + } + + if (graph.get_value_type(pad_id) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error("WebGPU constant_pad_nd: pad is not an IntList"); + } + const std::vector& pad = graph.get_int_list(pad_id); + if (pad.size() % 2 != 0) { + throw std::runtime_error("WebGPU constant_pad_nd: pad must be even-length"); + } + const size_t npad = pad.size() / 2; + if (npad > nd) { + throw std::runtime_error("WebGPU constant_pad_nd: pad longer than rank"); + } + + // value scalar (default 0). Vulkan serializes a Scalar as Int or Double. + float value = 0.0f; + if (args.size() >= 4) { + value = utils::scalar_or(graph, args.at(2), 0.0f); + } + + // Per-dim left/right pad (pad list is reversed-dim, from the LAST dim). + std::array left = {0, 0, 0, 0}; + std::array right = {0, 0, 0, 0}; + for (size_t k = 0; k < npad; k++) { + const size_t d = nd - 1 - k; // k-th pad entry -> dim (nd-1-k) + left[d] = pad[2 * k]; + right[d] = pad[2 * k + 1]; + } + + // Validate output dims == in + left + right per dim (loud-fail on a wrong + // pad-list interpretation), before any buffer alloc -> no leak-on-throw. + for (size_t d = 0; d < nd; d++) { + // The kernel's pad params are u32; a negative pad (cropping) would wrap + // into a huge offset, so reject it loudly rather than gather out of bounds. + if (left[d] < 0 || right[d] < 0) { + throw std::runtime_error( + "WebGPU constant_pad_nd: negative pad (cropping) not supported"); + } + const int64_t expect = in.dims[d] + left[d] + right[d]; + if (expect < 0 || static_cast(out.dims[d]) != expect) { + throw std::runtime_error("WebGPU constant_pad_nd: output shape mismatch"); + } + } + + const uint64_t out_numel = + utils::check_fp32(out, "constant_pad_nd", "output"); + utils::check_fp32(in, "constant_pad_nd", "input"); + + // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 + // ceiling. stride_x lets the shader decode idx = gid.y*stride_x + gid.x. + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + utils::checked_u32(out_numel, "constant_pad_nd"), + kConstantPadNdWorkgroupSizeX, + "constant_pad_nd"); + + // Right-align dims into [4]: leading (4-nd) entries get extent 1, pad 0. + PadParams params = {}; + for (int s = 0; s < 4; s++) { + params.out_dims[s] = 1; + params.in_dims[s] = 1; + params.left[s] = 0; + } + const size_t off = 4 - nd; + for (size_t d = 0; d < nd; d++) { + params.out_dims[off + d] = static_cast(out.dims[d]); + params.in_dims[off + d] = static_cast(in.dims[d]); + params.left[off + d] = static_cast(left[d]); + } + params.out_numel = static_cast(out_numel); + params.value = value; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(PadParams)); + graph.add_uniform_buffer_bytes(sizeof(PadParams)); + + auto constants = utils::make_grid_constants(grid); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConstantPadNdWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(PadParams)}, + }, + constants.data(), + constants.size()); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.constant_pad_nd.default, constant_pad_nd_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl new file mode 100644 index 00000000000..9d7d4b76142 --- /dev/null +++ b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl @@ -0,0 +1,63 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +// Up to 4D. The handler right-aligns dims into [4] (leading entries = 1, left/ +// right pad = 0 for unpadded/leading dims), so the shader is rank-agnostic and +// always iterates 4 dims. in_dims[d] = input extent, left[d] = that dim's +// left-pad, out_dims[d] = in_dims[d] + left[d] + right[d]. +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: 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 + +// constant_pad_nd, gather form, NCHW row-major fp32. One thread per OUTPUT +// element: decode its 4D coords, subtract each dim's left-pad to get the input +// coord; if ALL input coords are in-bounds -> copy inp[flat_in], else write +// `value`. Pure copy/fill -> bit-exact. (CPU-derisked == torch at 0.) +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.out_numel) { + return; + } + + // decode out coords (last dim fastest) + var rem = i; + let o3 = rem % params.out_dims.w; + rem = rem / params.out_dims.w; + let o2 = rem % params.out_dims.z; + rem = rem / params.out_dims.z; + let o1 = rem % params.out_dims.y; + rem = rem / params.out_dims.y; + let o0 = rem % params.out_dims.x; + + // subtract left pad -> input coord (wrapping subtract; check via < in_dim on + // the unsigned result catches negatives because they wrap to huge values) + let c0 = o0 - params.left.x; + let c1 = o1 - params.left.y; + let c2 = o2 - params.left.z; + let c3 = o3 - params.left.w; + + let in0 = o0 >= params.left.x && c0 < params.in_dims.x; + let in1 = o1 >= params.left.y && c1 < params.in_dims.y; + let in2 = o2 >= params.left.z && c2 < params.in_dims.z; + let in3 = o3 >= params.left.w && c3 < params.in_dims.w; + + if (in0 && in1 && in2 && in3) { + let in_idx = + ((c0 * params.in_dims.y + c1) * params.in_dims.z + c2) * params.in_dims.w + + c3; + out[i] = inp[in_idx]; + } else { + out[i] = params.value; + } +} diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h new file mode 100644 index 00000000000..a8380304890 --- /dev/null +++ b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h @@ -0,0 +1,87 @@ +/* + * 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 constant_pad_nd.wgsl - DO NOT EDIT. +// wgsl-sha256: 67496c3b851bbe69c64d26bcca415cd4c0de20d555b0ed5cccbe66499a757592 +inline constexpr const char* kConstantPadNdWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +// Up to 4D. The handler right-aligns dims into [4] (leading entries = 1, left/ +// right pad = 0 for unpadded/leading dims), so the shader is rank-agnostic and +// always iterates 4 dims. in_dims[d] = input extent, left[d] = that dim's +// left-pad, out_dims[d] = in_dims[d] + left[d] + right[d]. +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: 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 + +// constant_pad_nd, gather form, NCHW row-major fp32. One thread per OUTPUT +// element: decode its 4D coords, subtract each dim's left-pad to get the input +// coord; if ALL input coords are in-bounds -> copy inp[flat_in], else write +// `value`. Pure copy/fill -> bit-exact. (CPU-derisked == torch at 0.) +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.out_numel) { + return; + } + + // decode out coords (last dim fastest) + var rem = i; + let o3 = rem % params.out_dims.w; + rem = rem / params.out_dims.w; + let o2 = rem % params.out_dims.z; + rem = rem / params.out_dims.z; + let o1 = rem % params.out_dims.y; + rem = rem / params.out_dims.y; + let o0 = rem % params.out_dims.x; + + // subtract left pad -> input coord (wrapping subtract; check via < in_dim on + // the unsigned result catches negatives because they wrap to huge values) + let c0 = o0 - params.left.x; + let c1 = o1 - params.left.y; + let c2 = o2 - params.left.z; + let c3 = o3 - params.left.w; + + let in0 = o0 >= params.left.x && c0 < params.in_dims.x; + let in1 = o1 >= params.left.y && c1 < params.in_dims.y; + let in2 = o2 >= params.left.z && c2 < params.in_dims.z; + let in3 = o3 >= params.left.w && c3 < params.in_dims.w; + + if (in0 && in1 && in2 && in3) { + let in_idx = + ((c0 * params.in_dims.y + c1) * params.in_dims.z + c2) * params.in_dims.w + + c3; + out[i] = inp[in_idx]; + } else { + out[i] = params.value; + } +} +)"; + +inline constexpr uint32_t kConstantPadNdWorkgroupSizeX = 256; +inline constexpr uint32_t kConstantPadNdWorkgroupSizeY = 1; +inline constexpr uint32_t kConstantPadNdWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 61edf4287efc53fca5049c41e2aafe067e42933a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:19 -0700 Subject: [PATCH 15/41] [ExecuTorch][WebGPU] constant_pad_nd 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/20858 Splits the `constant_pad_nd` op tests into their own diff, stacked directly above the `constant_pad_nd` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_constant_pad_nd.py` and registers the `constant_pad_nd` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949732 @exported-using-ghexport Differential Revision: [D111072712](https://our.internmc.facebook.com/intern/diff/D111072712/) --- backends/webgpu/test/op_tests/cases.py | 28 +++++++++++++++++++ .../webgpu/test/ops/test_constant_pad_nd.py | 23 +++++++++++++++ 2 files changed, 51 insertions(+) create mode 100644 backends/webgpu/test/ops/test_constant_pad_nd.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index ef79748bba7..7dc224191e0 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -604,3 +604,31 @@ def _addmm_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_constant_pad_nd import ( + _randn as _pad_randn, + PadModule, +) + + +@register_op_test("constant_pad_nd") +def _constant_pad_nd_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=lambda pad: PadModule(pad), + cases=[ + Case( + name="last2", + construct={"pad": [1, 2]}, + inputs=(InputSpec(shape=(3, 8), gen=_pad_randn),), + ), + Case( + name="rank3", + construct={"pad": [1, 1, 2, 0]}, + inputs=(InputSpec(shape=(2, 4, 8), gen=_pad_randn),), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_constant_pad_nd.py b/backends/webgpu/test/ops/test_constant_pad_nd.py new file mode 100644 index 00000000000..746cdb7b10e --- /dev/null +++ b/backends/webgpu/test/ops/test_constant_pad_nd.py @@ -0,0 +1,23 @@ +# 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.constant_pad_nd.default` module (DaViT window padding) for op-tests.""" + +import torch + + +class PadModule(torch.nn.Module): + def __init__(self, pad): + super().__init__() + self.pad = tuple(pad) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.pad(x, self.pad, value=0.0) + + +def _randn(shape) -> torch.Tensor: + g = torch.Generator().manual_seed(sum(int(x) for x in shape)) + return torch.randn(*shape, generator=g) * 0.1 From 5f4576477508e4cd86dc83ebeb076b5babe49ae1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:20 -0700 Subject: [PATCH 16/41] [ExecuTorch][WebGPU] Add upsample_nearest2d 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/20859 **Adds `aten.upsample_nearest2d.vec` to the WebGPU backend, enabling the SAM2/SAM3 pixel-decoder / FPN nearest-neighbour upsample path.** **Problem** — nearest-neighbour 2D upsample (`F.interpolate(..., mode="nearest")`) had no WebGPU kernel, blocking the FPN 2x upsample chain in the SAM2/SAM3 pixel decoder. **Solution** — Before — no handler; `aten.upsample_nearest2d.vec` unsupported. After — one thread per output element `(n, c, oh, ow)` maps back to its nearest source pixel and copies it. **Implementation**: - The kernel computes the source index with ATen's legacy nearest formula `ih = floor(oh*IH/OH)`, `iw = floor(ow*IW/OW)` (integer division), then copies `inp[((n*C+c)*IH+ih)*IW+iw]` — a plain NCHW row-major gather, bit-exact. - `OH`/`OW` are taken from the output tensor's own dims (not re-derived from the `size`/`scale` args), and the handler validates that `N`/`C` match between input and output. - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid` (workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 ceiling; `stride_x` decode). - Divergence from Vulkan (intentional): this does NOT mirror Vulkan `backends/vulkan/runtime/graph/ops/impl/Upsample.cpp`, which computes the source index with the half-pixel-center reciprocal-scale formula. Half-pixel-center is correct for bilinear but wrong for non-exact nearest; the ATen `nearest_neighbor_compute_source_index` (`UpSample.h`) `floor(oh*IH/OH)` form is the one that matches PyTorch's `mode="nearest"`. The two agree on exact-integer ratios (e.g. 2x) but diverge on non-integer ratios (e.g. 5->8), which the op-test's `non_2x_ratio` case pins down. **Constraints** — fp32 only; 4D NCHW input and output; `nearest` mode only; non-zero spatial dims; the output element count must fit `u32`. Co-authored-with: Claude Code. ghstack-source-id: 405949741 @exported-using-ghexport Differential Revision: [D110836668](https://our.internmc.facebook.com/intern/diff/D110836668/) --- .../upsample_nearest2d/UpsampleNearest2d.cpp | 118 ++++++++++++++++++ .../upsample_nearest2d.wgsl | 40 ++++++ .../upsample_nearest2d_wgsl.h | 64 ++++++++++ 3 files changed, 222 insertions(+) create mode 100644 backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp create mode 100644 backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl create mode 100644 backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp b/backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp new file mode 100644 index 00000000000..0664e53a0ae --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp @@ -0,0 +1,118 @@ +/* + * 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 _p0; + uint32_t _p1; +}; +static_assert(sizeof(UpsampleParams) == 32, "UpsampleParams must be 32 bytes"); + +// OH/OW come from the output tensor's own dims, not the size/scale args. +void upsample_nearest2d_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 2) { + throw std::runtime_error("WebGPU upsample_nearest2d: expected >=2 args"); + } + const int in_id = args.at(0); + 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_nearest2d: 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_nearest2d: N/C mismatch"); + } + if (IH == 0 || IW == 0 || OH == 0 || OW == 0) { + throw std::runtime_error("WebGPU upsample_nearest2d: zero spatial dim"); + } + + uint64_t out_numel = utils::check_fp32(out, "upsample_nearest2d", "output"); + + // 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. stride_x lets the shader decode idx = gid.y*stride_x + gid.x. + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + utils::checked_u32(out_numel, "upsample_nearest2d"), + kUpsampleNearest2dWorkgroupSizeX, + "upsample_nearest2d"); + + UpsampleParams params = {}; + params.N = N; + params.C = C; + params.IH = IH; + params.IW = IW; + params.OH = OH; + params.OW = OW; + + 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, + kUpsampleNearest2dWGSL, + { + {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()); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.upsample_nearest2d.vec, upsample_nearest2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl new file mode 100644 index 00000000000..7f5310597ca --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl @@ -0,0 +1,40 @@ +@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, + _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 + +// Nearest-neighbour 2D upsample, NCHW row-major, fp32. One thread per output +// element (n, c, oh, ow). oh -> floor(oh*IH/OH), matching real ATen's +// nearest_neighbor_compute_source_index (UpSample.h) — intentionally NOT +// mirroring Vulkan's upsample_2d.glsl, which uses the half-pixel-center +// formula correct for bilinear but not for non-exact nearest (see diff summary). +@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 ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = ((n * params.C + c) * params.IH + ih) * params.IW + iw; + out[i] = inp[in_idx]; +} diff --git a/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h new file mode 100644 index 00000000000..89da4f3474c --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h @@ -0,0 +1,64 @@ +/* + * 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_nearest2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 3eb5bb5604a34371e493ecc150f02995374ab183fed7e7c4a4e036b5b45563c0 +inline constexpr const char* kUpsampleNearest2dWGSL = 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, + _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 + +// Nearest-neighbour 2D upsample, NCHW row-major, fp32. One thread per output +// element (n, c, oh, ow). oh -> floor(oh*IH/OH), matching real ATen's +// nearest_neighbor_compute_source_index (UpSample.h) — intentionally NOT +// mirroring Vulkan's upsample_2d.glsl, which uses the half-pixel-center +// formula correct for bilinear but not for non-exact nearest (see diff summary). +@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 ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = ((n * params.C + c) * params.IH + ih) * params.IW + iw; + out[i] = inp[in_idx]; +} +)"; + +inline constexpr uint32_t kUpsampleNearest2dWorkgroupSizeX = 256; +inline constexpr uint32_t kUpsampleNearest2dWorkgroupSizeY = 1; +inline constexpr uint32_t kUpsampleNearest2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 0cd69f3992bd79dc085fcca201f22a0af821b6ef Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:20 -0700 Subject: [PATCH 17/41] [ExecuTorch][WebGPU] upsample_nearest2d 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/20860 Splits the `upsample_nearest2d` op tests into their own diff, stacked directly above the `upsample_nearest2d` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_upsample_nearest2d.py` and registers the `upsample_nearest2d` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949736 @exported-using-ghexport Differential Revision: [D111072707](https://our.internmc.facebook.com/intern/diff/D111072707/) --- backends/webgpu/test/op_tests/cases.py | 45 +++++++++++++++++++ .../test/ops/test_upsample_nearest2d.py | 33 ++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 backends/webgpu/test/ops/test_upsample_nearest2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 7dc224191e0..e3fb19f0e90 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -632,3 +632,48 @@ def _constant_pad_nd_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_upsample_nearest2d import ( + _det_input as _upsample_det_input, + UpsampleNearest2dModule, +) + + +@register_op_test("upsample_nearest2d") +def _upsample_nearest2d_suite() -> WebGPUTestSuite: + # SAM2 FPN 2x upsample chain (36->72->144) + a cheap tiny eyeball case. + return WebGPUTestSuite( + module_factory=lambda scale_factor: UpsampleNearest2dModule(scale_factor), + cases=[ + Case( + name="fpn_36_72", + construct={"scale_factor": 2.0}, + inputs=(InputSpec(shape=(1, 8, 36, 36), gen=_upsample_det_input),), + ), + Case( + name="fpn_72_144", + construct={"scale_factor": 2.0}, + inputs=(InputSpec(shape=(1, 8, 72, 72), gen=_upsample_det_input),), + ), + Case( + name="tiny", + construct={"scale_factor": 2.0}, + inputs=(InputSpec(shape=(1, 4, 5, 7), gen=_upsample_det_input),), + ), + Case( + # Non-integer, non-2x ratio (5->8): floor(oh*5/8) and the + # half-pixel-center formula round((oh+0.5)*5/8-0.5) diverge at + # oh=3,6 (1 vs 2, 3 vs 4) — locks in the legacy "nearest" + # formula (see upsample_nearest2d.wgsl) against a genuinely + # discriminating ratio, not just the 2x cases above where both + # formulas happen to agree. + name="non_2x_ratio", + construct={"scale_factor": 1.6}, + inputs=(InputSpec(shape=(1, 2, 5, 5), gen=_upsample_det_input),), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_upsample_nearest2d.py b/backends/webgpu/test/ops/test_upsample_nearest2d.py new file mode 100644 index 00000000000..c082a6487f5 --- /dev/null +++ b/backends/webgpu/test/ops/test_upsample_nearest2d.py @@ -0,0 +1,33 @@ +# 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_nearest2d.vec` module for the WebGPU op-test framework. + +Upsample is on the SAM2/SAM3 pixel-decoder / FPN path +(`F.interpolate(..., mode="nearest")`). +""" + +import torch + + +class UpsampleNearest2dModule(torch.nn.Module): + def __init__(self, scale_factor: float) -> None: + super().__init__() + self.scale_factor = scale_factor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.interpolate( + x, scale_factor=self.scale_factor, mode="nearest" + ) + + +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 e15feb7c70542e35a2fe6da9c46e88e954ce5bd9 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:20 -0700 Subject: [PATCH 18/41] [ExecuTorch][WebGPU] Add max_pool2d 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/20861 **Adds `aten.max_pool2d_with_indices.default` to the WebGPU backend, enabling the SAM2 Hiera q_pool path (values, with optional indices).** **Problem** — max pooling (`F.max_pool2d`, which decomposes to `max_pool2d_with_indices`) had no WebGPU kernel, blocking the SAM2 Hiera `q_pool` downsampling stages. **Solution** — Before — no handler; the multi-output `max_pool2d_with_indices` unsupported. After — one thread per output element `(n, c, oh, ow)` gathers its pooling window and writes the max (and, when the graph requests them, the argmax indices). **Implementation**: - The kernel iterates the `kH x kW` window with general `stride`/`padding`/`dilation`, skips out-of-range (padding) cells, and tracks the running max; `best` initialises to `-3.4e38` (a large finite negative, because Dawn/Tint rejects the exact `-FLT_MAX` literal). - Argmax is ALWAYS tracked; a `write_indices` override (a compile-time spec constant) gates only the final `out_idx` store, so the values-only and with-indices paths run one shader and stay bit-identical on `out_vals`. Indices are the flat `ih*IW+iw` spatial-plane offset (matching `torch.nn.functional.max_pool2d(return_indices=True)`, not an absolute NCHW offset). - The out is a `ValueList` `[values, indices]`; when indices are not requested the handler binds a tiny dummy storage buffer to slot 3 (the shader never writes it) rather than allocating a real indices tensor — mirroring the `dummy_affine` pattern in `NativeLayerNorm.cpp`. When indices ARE requested it validates that the indices tensor is `int32` (4 bytes/elem) and shares the values shape, else throws (the kernel writes `i32`, so an `int64` indices tensor would mis-stride the write). - The handler parses `kernel_size`/`stride`/`padding`/`dilation` via `utils::parse_hw` (a single value broadcasts to both spatial dims; `stride` defaults to `kernel_size` when empty), derives `OH`/`OW` from the pooling formula, and validates them against the serialized values output (loud-fail on a `ceil_mode` / layout mismatch). - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Pool.cpp` (the `ValueList[values, indices]` shape, the `write_indices` spec constant, unconditional argmax tracking, and 32-bit indices despite ATen's int64 schema). **Constraints** — fp32 values; 4D NCHW; general `kernel`/`stride`/`padding`/`dilation`; `ceil_mode` unsupported (the output-shape validation assumes floor); the optional indices output must be `int32`; the output element count must fit `u32`. Co-authored-with: Claude Code. ghstack-source-id: 405949738 @exported-using-ghexport Differential Revision: [D110836680](https://our.internmc.facebook.com/intern/diff/D110836680/) --- .../runtime/ops/max_pool2d/MaxPool2d.cpp | 198 ++++++++++++++++++ .../runtime/ops/max_pool2d/max_pool2d.wgsl | 79 +++++++ .../runtime/ops/max_pool2d/max_pool2d_wgsl.h | 103 +++++++++ 3 files changed, 380 insertions(+) create mode 100644 backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp create mode 100644 backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl create mode 100644 backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp new file mode 100644 index 00000000000..682eff61197 --- /dev/null +++ b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp @@ -0,0 +1,198 @@ +/* + * 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 PoolParams { + uint32_t N, C, IH, IW, OH, OW; + uint32_t kH, kW, sH, sW, pH, pW, dH, dW; + uint32_t _p0, _p1; +}; +static_assert(sizeof(PoolParams) == 64, "PoolParams must be 64 bytes"); + +// out_list = ValueList[values, indices] (mirrors Vulkan Pool.cpp). +void max_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 6) { + throw std::runtime_error("WebGPU max_pool2d: expected >=6 args"); + } + const int in_id = args.at(0); + const int kernel_id = args.at(1); + const int stride_id = args.at(2); + const int padding_id = args.at(3); + const int dilation_id = args.at(4); + const int out_list_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + if (graph.get_value_type(out_list_id) != WebGPUGraph::ValueType::ValueList) { + throw std::runtime_error("WebGPU max_pool2d: out is not a ValueList"); + } + const std::vector& outs = graph.get_value_list(out_list_id); + if (outs.empty()) { + throw std::runtime_error("WebGPU max_pool2d: empty out ValueList"); + } + const int values_id = outs.at(0); // [0]=values, [1]=indices + const bool has_indices = outs.size() > 1 && + graph.get_value_type(outs.at(1)) == WebGPUGraph::ValueType::Tensor; + if (has_indices) { + const auto& idx_t = graph.get_tensor(outs.at(1)); + if (idx_t.dims != graph.get_tensor(values_id).dims) { + throw std::runtime_error( + "WebGPU max_pool2d: indices output shape must match values output"); + } + // The WGSL kernel writes int32 indices (mirrors Vulkan's Pool.cpp, which + // also computes/stores 32-bit indices despite ATen's int64 schema) — throw + // rather than silently mis-stride the buffer if this graph's indices + // tensor ever turns out to be int64 (8 bytes/elem) instead. + const uint64_t idx_numel = utils::numel_of(idx_t.dims); + if (idx_t.nbytes != idx_numel * sizeof(int32_t)) { + throw std::runtime_error( + "WebGPU max_pool2d: indices output must be int32 (4 bytes/elem); " + "got a different byte width, would mis-stride the i32 kernel write"); + } + } + + const auto& in = graph.get_tensor(in_id); + const auto& out = graph.get_tensor(values_id); + if (in.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU max_pool2d: expected 4D in/out"); + } + + // kernel/stride/padding/dilation are int lists; PyTorch broadcasts a single + // value to both spatial dims. stride defaults to kernel_size when empty. + uint32_t kH, kW, sH, sW, pH, pW, dH, dW; + utils::parse_hw( + graph.get_int_list(kernel_id), kH, kW, "max_pool2d", "kernel_size"); + const std::vector stride_v = graph.get_int_list(stride_id); + if (stride_v.empty()) { + sH = kH; + sW = kW; + } else { + utils::parse_hw(stride_v, sH, sW, "max_pool2d", "stride"); + } + utils::parse_hw( + graph.get_int_list(padding_id), pH, pW, "max_pool2d", "padding"); + utils::parse_hw( + graph.get_int_list(dilation_id), dH, dW, "max_pool2d", "dilation"); + + 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]); + if (sH == 0 || sW == 0) { + throw std::runtime_error("WebGPU max_pool2d: zero stride"); + } + + const int64_t oh_num = + int64_t(IH) + 2 * int64_t(pH) - int64_t(dH) * (int64_t(kH) - 1) - 1; + const int64_t ow_num = + int64_t(IW) + 2 * int64_t(pW) - int64_t(dW) * (int64_t(kW) - 1) - 1; + if (oh_num < 0 || ow_num < 0) { + throw std::runtime_error("WebGPU max_pool2d: kernel larger than input"); + } + const uint32_t OH = static_cast(oh_num / int64_t(sH)) + 1; + const uint32_t OW = static_cast(ow_num / int64_t(sW)) + 1; + + // Validate against the serialized values output [N, C, OH, OW] (loud-fail if + // the arg interpretation is wrong, e.g. ceil_mode or a different layout). + if (static_cast(out.dims[0]) != N || + static_cast(out.dims[1]) != C || + static_cast(out.dims[2]) != OH || + static_cast(out.dims[3]) != OW) { + throw std::runtime_error("WebGPU max_pool2d: output shape mismatch"); + } + + uint64_t out_numel = utils::check_fp32(out, "max_pool2d", "output"); + + // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 + // ceiling. stride_x lets the shader decode idx = gid.y*stride_x + gid.x. + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + utils::checked_u32(out_numel, "max_pool2d"), + kMaxPool2dWorkgroupSizeX, + "max_pool2d"); + + PoolParams params = {}; + params.N = N; + params.C = C; + params.IH = IH; + params.IW = IW; + params.OH = OH; + params.OW = OW; + params.kH = kH; + params.kW = kW; + params.sH = sH; + params.sW = sW; + params.pH = pH; + params.pW = pW; + params.dH = dH; + params.dW = dW; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(PoolParams)); + graph.add_uniform_buffer_bytes(sizeof(PoolParams)); + + auto grid_constants = utils::make_grid_constants(grid); + WGPUConstantEntry constants[3] = {}; + constants[0] = grid_constants[0]; + constants[1] = grid_constants[1]; + constants[2].key = {"write_indices", WGPU_STRLEN}; + constants[2].value = has_indices ? 1.0 : 0.0; + + // write_indices==0 -> the shader never stores to out_idx, so a tiny dummy + // buffer is safe (mirrors NativeLayerNorm.cpp's dummy_affine pattern). + utils::OptionalBinding idx = utils::make_optional_binding( + device, + has_indices, + has_indices ? graph.get_tensor(outs.at(1)).buffer : nullptr, + has_indices ? graph.get_tensor(outs.at(1)).nbytes : 0); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kMaxPool2dWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(PoolParams)}, + {3, WGPUBufferBindingType_Storage, idx.buffer, idx.nbytes}, + }, + constants, + 3); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + + wgpuBufferRelease(uniform_buffer); + if (idx.owned_dummy != nullptr) { + wgpuBufferRelease(idx.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.max_pool2d_with_indices.default, max_pool2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl new file mode 100644 index 00000000000..3d7c67ddf9a --- /dev/null +++ b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl @@ -0,0 +1,79 @@ +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill +override write_indices: u32 = 0u; // 1 = also write out_idx (compile-time gate, mirrors Vulkan) + +// max_pool2d (values [+ optional indices]), NCHW row-major, fp32. One thread per +// output element (n, c, oh, ow); gather the window, take the max. General +// stride/pad/dilation. Argmax is ALWAYS tracked (mirrors Vulkan Pool.cpp, which +// computes indices unconditionally and only gates the final write) so the +// values-only and with-indices paths share one kernel; write_indices==0 skips +// only the out_idx store, not the tracking, keeping both paths bit-identical on +// out_vals. Indices are the flat (ih*IW+iw) spatial-plane offset (matches +// torch.nn.functional.max_pool2d(return_indices=True)'s documented convention, +// NOT an absolute offset into the full NCHW tensor). Pad cells are skipped +// (init -inf); out_idx is left at its bound (dummy, when unused) value. +@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 iH = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; // * IW added per-row below + + var best: f32 = -3.4e38; // large finite negative max-init (Dawn/Tint rejects -3.40282347e38: > f32 max); matches q4gsw_dq8ca convention + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + out_vals[i] = best; + if (write_indices != 0u) { + out_idx[i] = best_idx; + } +} diff --git a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h new file mode 100644 index 00000000000..1dbe0571e40 --- /dev/null +++ b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h @@ -0,0 +1,103 @@ +/* + * 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 max_pool2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 2215b7ba725face441baac5eca8e2133458b414377868354c3fda9046a88109b +inline constexpr const char* kMaxPool2dWGSL = R"( +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill +override write_indices: u32 = 0u; // 1 = also write out_idx (compile-time gate, mirrors Vulkan) + +// max_pool2d (values [+ optional indices]), NCHW row-major, fp32. One thread per +// output element (n, c, oh, ow); gather the window, take the max. General +// stride/pad/dilation. Argmax is ALWAYS tracked (mirrors Vulkan Pool.cpp, which +// computes indices unconditionally and only gates the final write) so the +// values-only and with-indices paths share one kernel; write_indices==0 skips +// only the out_idx store, not the tracking, keeping both paths bit-identical on +// out_vals. Indices are the flat (ih*IW+iw) spatial-plane offset (matches +// torch.nn.functional.max_pool2d(return_indices=True)'s documented convention, +// NOT an absolute offset into the full NCHW tensor). Pad cells are skipped +// (init -inf); out_idx is left at its bound (dummy, when unused) value. +@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 iH = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; // * IW added per-row below + + var best: f32 = -3.4e38; // large finite negative max-init (Dawn/Tint rejects -3.40282347e38: > f32 max); matches q4gsw_dq8ca convention + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + out_vals[i] = best; + if (write_indices != 0u) { + out_idx[i] = best_idx; + } +} +)"; + +inline constexpr uint32_t kMaxPool2dWorkgroupSizeX = 256; +inline constexpr uint32_t kMaxPool2dWorkgroupSizeY = 1; +inline constexpr uint32_t kMaxPool2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From e498661c069ce324a2d2b47006b62d4848c2c936 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:21 -0700 Subject: [PATCH 19/41] [ExecuTorch][WebGPU] max_pool2d 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/20862 Splits the `max_pool2d` op tests into their own diff, stacked directly above the `max_pool2d` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_max_pool2d.py` and registers the `max_pool2d` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949735 @exported-using-ghexport Differential Revision: [D111072724](https://our.internmc.facebook.com/intern/diff/D111072724/) --- backends/webgpu/test/op_tests/cases.py | 42 +++++++++++++++++++++ backends/webgpu/test/ops/test_max_pool2d.py | 42 +++++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 backends/webgpu/test/ops/test_max_pool2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index e3fb19f0e90..4ca730428ce 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -677,3 +677,45 @@ def _upsample_nearest2d_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) + + +from executorch.backends.webgpu.test.ops.test_max_pool2d import ( + _det_input as _maxpool_det_input, + MaxPool2dModule, +) + + +@register_op_test("max_pool2d") +def _max_pool2d_suite() -> WebGPUTestSuite: + # SAM2 Hiera q_pool (native shape) + real Hiera channel count (ch768) + + # a padding case + a tiny eyeball case. + return WebGPUTestSuite( + module_factory=lambda kernel_size, stride, padding: MaxPool2dModule( + kernel_size, stride, padding + ), + cases=[ + Case( + name="q_pool", + construct={"kernel_size": 2, "stride": 2, "padding": 0}, + inputs=(InputSpec(shape=(1, 8, 12, 12), gen=_maxpool_det_input),), + ), + Case( + name="ch768", + construct={"kernel_size": 2, "stride": 2, "padding": 0}, + inputs=(InputSpec(shape=(1, 768, 14, 14), gen=_maxpool_det_input),), + ), + Case( + name="pad1", + construct={"kernel_size": 3, "stride": 2, "padding": 1}, + inputs=(InputSpec(shape=(1, 8, 7, 7), gen=_maxpool_det_input),), + ), + Case( + name="tiny", + construct={"kernel_size": 2, "stride": 2, "padding": 0}, + inputs=(InputSpec(shape=(1, 4, 5, 5), gen=_maxpool_det_input),), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_max_pool2d.py b/backends/webgpu/test/ops/test_max_pool2d.py new file mode 100644 index 00000000000..523ed75d7ee --- /dev/null +++ b/backends/webgpu/test/ops/test_max_pool2d.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.max_pool2d_with_indices.default` module for the WebGPU op-test +framework. + +`F.max_pool2d` decomposes to `aten.max_pool2d_with_indices.default`, a +multi-output op whose out is a ValueList `[values, indices]`; the handler +writes the VALUES only and never the int64 indices +(`runtime/ops/max_pool2d/MaxPool2d.cpp`). Max-pool is on the SAM2 Hiera +q_pool path. +""" + +import torch + + +class MaxPool2dModule(torch.nn.Module): + def __init__(self, kernel_size: int, stride: int, padding: int) -> None: + super().__init__() + self.kernel_size = kernel_size + self.stride = stride + self.padding = padding + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.max_pool2d( + x, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + ) + + +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 2986c6777ef96523bbc38d010a079a503afcf124 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:21 -0700 Subject: [PATCH 20/41] [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 ab6eec5475854e81779bf123a867d20674b324d4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:22 -0700 Subject: [PATCH 21/41] [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 164f549d2556da11dff355ae8126c5207c372d44 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:22 -0700 Subject: [PATCH 22/41] [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 ad334e24a1ed9fec33fff877d68e5618c2af0f42 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:22 -0700 Subject: [PATCH 23/41] [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 2bfd02875ec12b26965d47473bcc7d895ac73535 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:23 -0700 Subject: [PATCH 24/41] [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 a57413bae02b03470f16e557bb213cd7ac9f41c4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:23 -0700 Subject: [PATCH 25/41] [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 0a414209d9d06ca8acde198b9431b5d904aa0b84 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:24 -0700 Subject: [PATCH 26/41] [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 a506d03f7ba9e6f37fa6dac6c7d2caade9fa66c3 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:24 -0700 Subject: [PATCH 27/41] [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 5ff79aa2a2d48ee6bbc2fd168e8517040c08b508 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:24 -0700 Subject: [PATCH 28/41] [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 a14237b56795889419f57c02f8722769fa02f5db Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:25 -0700 Subject: [PATCH 29/41] [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 30/41] [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 c740ba7ba576aa3b3b0da0c2811383abaf9e8caf Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:26 -0700 Subject: [PATCH 31/41] [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 35e1e4c7cd84ce989f087e6e839887cf4132d245 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:26 -0700 Subject: [PATCH 32/41] [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 51f30645b13c13d9a51b8a2a81a359a8c32474b1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:26 -0700 Subject: [PATCH 33/41] [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 3696742a81fb3ea8c55b91d5b89772257c17d015 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:27 -0700 Subject: [PATCH 34/41] [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 4c7e000f7df982d63d61a031b344a8b0407c70a2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:27 -0700 Subject: [PATCH 35/41] [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 7d54ee1c07edac65a25b6297c80ec7e40107e5b9 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:28 -0700 Subject: [PATCH 36/41] [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 30abd20145cbe7e3d5940e8b903311e44074bf34 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:28 -0700 Subject: [PATCH 37/41] [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 6c89f5fd5fecb6580b9b39edc24b52c723799155 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:28 -0700 Subject: [PATCH 38/41] [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 707d0376287c0ebecbb13e7f952c00c1eb16a508 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:29 -0700 Subject: [PATCH 39/41] [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 b1911f94bd0e41642247d4e3c044f927ed090059 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:29 -0700 Subject: [PATCH 40/41] [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 044f3abb7b92e7627c340dfc0015817895210517 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 22:05:30 -0700 Subject: [PATCH 41/41] [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);