From 9755c5dd2a37dd0f3775761a55381a843fb447e0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:53 -0700 Subject: [PATCH 01/14] [ExecuTorch][WebGPU] Dynamic tensor-shape resize engine core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20574 The WebGPU backend baked static tensor shapes at build time, so a dynamic `.pte` needed a separate graph for each shape (prefill vs. decode). This adds a tensor-shape resize engine mirroring Vulkan: tensors carry live `cur_dims` ≤ max, inputs resize per call, and a bounded-fixpoint propagates tensor-level resize hooks. **Key changes:** - `WebGPUTensor`: add `cur_dims`/`cur_nbytes` (live sizes ≤ max allocation), initialized to max at build - `WebGPUGraph`: `resize_input`/`set_cur_dims` validate live dims fit max, `propagate_resize` runs tensor hooks for dirty shapes - `update_symints_from_inputs` reads live `cur_dims`; adds `sym_size.int` dim source path - `copy_inputs` uploads only live bytes; `WebGPUBackend::execute` shrinks inputs and resizes outputs to live shapes Static graphs stay byte-identical: `cur == max` forever, no hooks fire, no reallocations. ghstack-source-id: 399812823 @exported-using-ghexport Differential Revision: [D109906091](https://our.internmc.facebook.com/intern/diff/D109906091/) --- backends/webgpu/runtime/WebGPUBackend.cpp | 38 +++++- backends/webgpu/runtime/WebGPUGraph.cpp | 125 +++++++++++++++--- backends/webgpu/runtime/WebGPUGraph.h | 56 +++++++- backends/webgpu/runtime/WebGPUUtils.h | 13 ++ .../ops/select_as_symint/SelectAsSymint.cpp | 23 ++++ 5 files changed, 225 insertions(+), 30 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index ceca89d1710..ba7f6f13e64 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -14,8 +14,11 @@ #include #include +#include #include +#include + #include namespace executorch { @@ -35,6 +38,7 @@ using executorch::runtime::Error; using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::register_backend; +using executorch::runtime::resize_tensor; using executorch::runtime::Result; using executorch::runtime::Span; @@ -100,19 +104,39 @@ Error WebGPUBackend::execute( // Copy inputs from EValue tensors to GPU buffers std::vector inputs; inputs.reserve(num_inputs); - for (size_t i = 0; i < num_inputs; i++) { - const auto& tensor = args[i]->toTensor(); - const bool host_is_int64 = - tensor.scalar_type() == executorch::aten::ScalarType::Long; - inputs.push_back({tensor.const_data_ptr(), tensor.nbytes(), host_is_int64}); - } // Fail loud as a runtime Error so a throw never crosses the backend boundary. try { + // Build the input list and, for dynamic shapes, shrink each input to its + // live sizes before upload (mirrors Vulkan maybe_resize_input). No-op when + // unchanged, so a static graph is byte-identical. + for (size_t i = 0; i < num_inputs; i++) { + const auto& tensor = args[i]->toTensor(); + const bool host_is_int64 = + tensor.scalar_type() == executorch::aten::ScalarType::Long; + inputs.push_back( + {tensor.const_data_ptr(), tensor.nbytes(), host_is_int64}); + const auto sizes = tensor.sizes(); + std::vector new_dims(sizes.begin(), sizes.end()); + graph->resize_input(graph->input_ids()[i], new_dims); + } graph->copy_inputs(inputs); graph->update_symints_from_inputs(inputs); graph->propagate_resize(); + // Resize each output EValue to its live shape so the readback length is + // correct (mirrors Vulkan maybe_resize_output). + for (size_t i = 0; i < num_outputs; i++) { + const auto& cd = graph->cur_dims(graph->output_ids()[i]); + std::vector osizes(cd.begin(), cd.end()); + Error e = resize_tensor( + args[num_inputs + i]->toTensor(), + ArrayRef(osizes.data(), osizes.size())); + if (e != Error::Ok) { + ET_LOG(Error, "WebGPU: output %zu resize failed", i); + return Error::Internal; + } + } } catch (const std::exception& e) { - ET_LOG(Error, "WebGPU input copy / symint refresh failed: %s", e.what()); + ET_LOG(Error, "WebGPU input/output resize / copy failed: %s", e.what()); return Error::Internal; } diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index b72c5256e69..bc3883399a1 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -62,6 +63,18 @@ bool vk_datatype_is_int(vkgraph::VkDataType dtype) { } } +// Normalize a possibly-negative dim against rank; throws (fail-loud) if OOR. +int normalize_dim(int dim, int rank, const char* op) { + if (dim < 0) { + dim += rank; + } + if (dim < 0 || dim >= rank) { + throw std::runtime_error( + std::string("WebGPU ") + op + ": dim out of range"); + } + return dim; +} + } // namespace WebGPUGraph::WebGPUGraph() = default; @@ -104,11 +117,10 @@ void WebGPUGraph::update_symints_from_inputs( throw std::runtime_error( "select_as_symint: source tensor is not a graph input"); } - const auto& dims = tensors_[src.input_tensor_id].dims; - int dim = src.dim < 0 ? src.dim + static_cast(dims.size()) : src.dim; - if (dim < 0 || dim >= static_cast(dims.size())) { - throw std::runtime_error("select_as_symint: dim out of range"); - } + // Live cur_dims: the source may be a dynamic-shape input. + const auto& dims = tensors_[src.input_tensor_id].cur_dims; + int dim = normalize_dim( + src.dim, static_cast(dims.size()), "select_as_symint"); int index = src.index; if (index < 0) { index += static_cast(dims[dim]); @@ -129,20 +141,26 @@ void WebGPUGraph::update_symints_from_inputs( } // Reads the [0,..,index,..,0] element; symint sources are scalar-ish. const int64_t offset = static_cast(index) * stride; - // elem_size back-derived from build-time numel (sources are static-shaped). const void* host = inputs[pos].data; - const size_t elem_size = inputs[pos].nbytes / static_cast(numel); + // Interpret the HOST buffer by its scalar type, not the tensor's serialized + // elem_size: copy_inputs narrows an int64 host input to an int32 buffer, so + // elem_size (buffer-derived) would misread int64 host data as int32. int32_t val; - if (elem_size == sizeof(int64_t)) { + if (inputs[pos].host_is_int64) { val = static_cast(static_cast(host)[offset]); - } else if (elem_size == sizeof(int32_t)) { - val = static_cast(host)[offset]; } else { - throw std::runtime_error( - "select_as_symint: unsupported input element size"); + val = static_cast(host)[offset]; } set_symint(src.symint_id, val); } + // sym_size.int: SymInt = a tensor's live dim (cur_dims). Usually unused (ops + // read cur_dims directly); for an intermediate source cur_dims is the build + // max here (hooks run later in propagate_resize), which is fine while unused. + for (const auto& s : symint_dim_sources_) { + const auto& d = tensors_[s.tensor_id].cur_dims; + int dim = normalize_dim(s.dim, static_cast(d.size()), "sym_size"); + set_symint(s.symint_id, static_cast(d[dim])); + } } void WebGPUGraph::set_symint(int id, int32_t val) { @@ -158,16 +176,78 @@ void WebGPUGraph::set_symint(int id, int32_t val) { } } +void WebGPUGraph::set_cur_dims( + int value_id, + const std::vector& new_dims) { + auto& t = tensors_[value_id]; + if (new_dims.size() != t.dims.size()) { + throw std::runtime_error("WebGPU resize: tensor rank changed"); + } + size_t numel = 1; + for (size_t d = 0; d < new_dims.size(); d++) { + // 0-sized dims unsupported: live shapes are always in [1, max] per dim. + if (new_dims[d] <= 0) { + throw std::runtime_error("WebGPU resize: new dim must be positive"); + } + if (new_dims[d] > t.dims[d]) { + throw std::runtime_error( + "WebGPU resize: new dim exceeds the max (serialized) allocation"); + } + numel *= static_cast(new_dims[d]); + } + const size_t new_nbytes = numel * t.elem_size; + if (t.cur_dims != new_dims) { + t.cur_dims = new_dims; + t.cur_nbytes = new_nbytes; + dirty_tensors_.insert(value_id); + } +} + +void WebGPUGraph::resize_input( + int value_id, + const std::vector& new_dims) { + if (std::find(input_ids_.begin(), input_ids_.end(), value_id) == + input_ids_.end()) { + throw std::runtime_error( + "WebGPUGraph::resize_input: value_id is not a graph input"); + } + set_cur_dims(value_id, new_dims); +} + void WebGPUGraph::propagate_resize() { - if (dirty_symints_.empty()) { + if (dirty_symints_.empty() && dirty_tensors_.empty()) { return; } + // Hooks fire in registration (topological) order: operands update first. for (auto& hook : resize_hooks_) { if (dirty_symints_.count(hook.symint_id) != 0) { hook.fn(*this); } } dirty_symints_.clear(); + // Tensor hooks: bounded fixpoint. A hook may dirty its output (cascading to a + // consumer); each pass handles the currently-dirty set. A forward DAG + // converges in <= depth passes (set_cur_dims re-dirties only on a change). + for (size_t pass = 0; + !dirty_tensors_.empty() && pass <= tensor_resize_hooks_.size(); + pass++) { + std::unordered_set processing; + processing.swap(dirty_tensors_); + for (auto& hook : tensor_resize_hooks_) { + if (processing.count(hook.trigger_tensor_id) != 0) { + hook.fn(*this); + } + } + } + if (!dirty_tensors_.empty()) { + throw std::runtime_error( + "WebGPU resize: tensor resize hooks did not converge"); + } + // Tensor hooks must not set_symint (dirty_symints_ already drained above). + if (!dirty_symints_.empty()) { + throw std::runtime_error( + "WebGPU resize: a tensor resize hook set a SymInt; not supported"); + } } WebGPUGraph::~WebGPUGraph() { @@ -322,6 +402,10 @@ void WebGPUGraph::build( tensor.elem_size = vk_datatype_size(vk_tensor->datatype()); tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); tensor.nbytes = numel * tensor.elem_size; + // Live dims start == max (serialized upper bound); resize_input shrinks + // them per call. Static graphs keep cur == max forever. + tensor.cur_dims = tensor.dims; + tensor.cur_nbytes = tensor.nbytes; int constant_id = vk_tensor->constant_id(); int mem_obj_id = vk_tensor->mem_obj_id(); @@ -624,17 +708,20 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { } int tid = input_ids_[i]; const auto& tensor = tensors_[tid]; + // Upload only the live (cur) bytes, not the max allocation; cur_nbytes == + // nbytes on a static graph, so this is byte-identical there. + const size_t live_nbytes = tensor.cur_nbytes; // Fast path: host and GPU element types match byte-for-byte. - if (in.nbytes == tensor.nbytes) { - wgpuQueueWriteBuffer(queue_, tensor.buffer, 0, in.data, tensor.nbytes); + if (in.nbytes == live_nbytes) { + wgpuQueueWriteBuffer(queue_, tensor.buffer, 0, in.data, live_nbytes); continue; } // Narrow int64 host indices into the int32 buffer (mirrors Vulkan). const bool buffer_is_int32 = tensor.is_int && tensor.elem_size == 4; - if (in.host_is_int64 && buffer_is_int32 && in.nbytes == tensor.nbytes * 2) { - const size_t numel = tensor.nbytes / 4; + if (in.host_is_int64 && buffer_is_int32 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / 4; const int64_t* src = static_cast(in.data); std::vector narrowed(numel); for (size_t e = 0; e < numel; e++) { @@ -648,7 +735,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { narrowed[e] = static_cast(src[e]); } wgpuQueueWriteBuffer( - queue_, tensor.buffer, 0, narrowed.data(), tensor.nbytes); + queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); continue; } @@ -656,7 +743,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { "WebGPU: unsupported input copy for input " + std::to_string(i) + " (host " + std::to_string(in.nbytes) + " bytes" + (in.host_is_int64 ? " int64" : "") + " vs buffer " + - std::to_string(tensor.nbytes) + " bytes)"); + std::to_string(live_nbytes) + " bytes)"); } } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 87f1576a5f3..5474ee4667a 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -23,8 +23,14 @@ namespace executorch::backends::webgpu { struct WebGPUTensor { WGPUBuffer buffer = nullptr; + // Max (allocation) dims/nbytes: the serialized upper-bound shape. The GPU + // buffer is sized from these and never reallocated (Vulkan allocate-at-max). std::vector dims; size_t nbytes = 0; + // Live dims/nbytes for dynamic shapes; always <= the max. == max on a static + // graph, so dynamic-resize logic keyed off these is inert there. + std::vector cur_dims; + size_t cur_nbytes = 0; // Serialized (GPU-side) element type, used to narrow wider host inputs. size_t elem_size = 0; bool is_int = false; @@ -171,6 +177,17 @@ class WebGPUGraph { return symint_sources_; } + // Records that a SymInt is a tensor's live dim size (sym_size.int), read from + // cur_dims at execute; distinct from SymIntSource (a scalar data element). + struct SymIntDimSource { + int symint_id; + int tensor_id; + int dim; + }; + void add_symint_dim_source(int symint_id, int tensor_id, int dim) { + symint_dim_sources_.push_back({symint_id, tensor_id, dim}); + } + // Execute-time select_as_symint read; mirrors Vulkan select_as_symint_impl. void update_symints_from_inputs(const std::vector& inputs); @@ -178,7 +195,25 @@ class WebGPUGraph { void add_resize_hook(int symint_id, std::function fn) { resize_hooks_.push_back({symint_id, std::move(fn)}); } - // Run hooks for changed SymInts then clear; call before execute(). + + // Set a graph input's live dims (<= max) + dirty it; static path stays inert. + void resize_input(int value_id, const std::vector& new_dims); + // Set a tensor's live dims (an op resize hook calls this for its output to + // cascade to consumers); validates the new dims fit the max, never reallocs. + void set_cur_dims(int value_id, const std::vector& new_dims); + const std::vector& cur_dims(int value_id) const { + return tensors_[value_id].cur_dims; + } + + // Per-tensor resize hook; mirrors Vulkan ExecuteNode::resize_fn. Runs in + // propagate_resize when trigger_tensor_id is dirty. + void add_tensor_resize_hook( + int trigger_tensor_id, + std::function fn) { + tensor_resize_hooks_.push_back({trigger_tensor_id, std::move(fn)}); + } + + // Run hooks for changed SymInts and tensors, then clear; call before execute. void propagate_resize(); // Mutable dispatch access for resize hooks (to rewrite workgroup_count_x). @@ -196,12 +231,14 @@ class WebGPUGraph { return queue_; } - void add_dispatch(WebGPUDispatch dispatch) { + // Returns the new dispatch's index (resize hooks rewrite it via dispatch_at). + size_t add_dispatch(WebGPUDispatch dispatch) { dispatches_.push_back(dispatch); + return dispatches_.size() - 1; } - // Record an in-graph-order buffer-to-buffer DMA (e.g. a flat copy). - void add_buffer_copy(WGPUBuffer src, WGPUBuffer dst, size_t nbytes) { + // 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; d.kind = WebGPUDispatch::Kind::Copy; d.copy_src = src; @@ -209,6 +246,7 @@ class WebGPUGraph { d.copy_nbytes = nbytes; d.kernel_name = "flat_copy"; dispatches_.push_back(d); + return dispatches_.size() - 1; } // Materialize a recorded prepack-routed constant into dst via one CPU->GPU @@ -297,6 +335,7 @@ class WebGPUGraph { }; std::unordered_map symints_; std::vector symint_sources_; + std::vector symint_dim_sources_; // Resize hooks + the set of SymInts changed since the last propagate_resize. struct ResizeHook { @@ -306,6 +345,15 @@ class WebGPUGraph { std::vector resize_hooks_; std::unordered_set dirty_symints_; + // Tensor-shape resize hooks + the set of tensors changed since the last + // propagate_resize (mirrors the SymInt pair above, for dynamic shapes). + struct TensorResizeHook { + int trigger_tensor_id; + std::function fn; + }; + std::vector tensor_resize_hooks_; + std::unordered_set dirty_tensors_; + std::vector input_ids_; std::vector output_ids_; diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index c5c779ffd5e..afa90d54aec 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -15,6 +15,7 @@ #include #include #include +#include namespace executorch::backends::webgpu::utils { @@ -24,6 +25,18 @@ inline T div_up(T a, T b) { return (a + b - 1) / b; } +// Product of dims (live element count); used by dynamic-resize hooks. +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; +} + // Clamp workgroup size to device limit (SwiftShader caps at 128). inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) { WGPULimits limits = {}; diff --git a/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp b/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp index 59f4d687a52..be333c99d53 100644 --- a/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp +++ b/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp @@ -40,6 +40,28 @@ void select_as_symint_impl(WebGPUGraph& graph, const std::vector& args) { static_cast(graph.get_int(index_id))); } +// aten.sym_size.int(self, dim) -> SymInt = self.size(dim). The WebGPU ops read +// live sizes from cur_dims directly, so this SymInt is usually unused. +void sym_size_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 3) { + throw std::runtime_error("sym_size.int: expected [self, dim, out] args"); + } + const int self_id = args.at(0); + const int dim_id = args.at(1); + const int out_id = args.at(2); + if (graph.get_value_type(out_id) != WebGPUGraph::ValueType::SymInt) { + return; // folded to a static Int -> nothing live to source + } + if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error("sym_size.int: dim arg is not an Int"); + } + if (graph.get_value_type(self_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("sym_size.int: self arg is not a Tensor"); + } + graph.add_symint_dim_source( + out_id, self_id, static_cast(graph.get_int(dim_id))); +} + // An operand is a live SymInt or a static Int constant. int32_t read_scalar(WebGPUGraph& graph, int id) { if (graph.get_value_type(id) == WebGPUGraph::ValueType::SymInt) { @@ -109,6 +131,7 @@ void sym_floordiv_impl(WebGPUGraph& graph, const std::vector& args) { WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(et_vk.select_as_symint.default, select_as_symint_impl); + WEBGPU_REGISTER_OP(sym_size.int, sym_size_impl); WEBGPU_REGISTER_OP(add, sym_add_impl); WEBGPU_REGISTER_OP(sub, sym_sub_impl); WEBGPU_REGISTER_OP(mul, sym_mul_impl); From 83e1cd73bc7fde46c4ac40ad1ad3b4088351c359 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:54 -0700 Subject: [PATCH 02/14] [ExecuTorch][WebGPU] Dynamic resize hooks for rms_norm, embedding, rope Pull Request resolved: https://github.com/pytorch/executorch/pull/20575 These ops baked their dispatch count, param UBO, and output dims at `build()` for the max seq-len. On a dynamic-shape graph at a smaller live S they would over-dispatch and leave the output sized at the max, so the resize engine could not actually shrink them. This adds tensor resize hooks to rms_norm, embedding_q4gsw, and apply_rotary_emb. When an input is resized, each hook recomputes the live row/token count, rewrites the param UBO, updates the dispatch `workgroup_count_x`, and sets the output's `cur_dims`. The hook is inert until a resize happens, so static graphs are byte-identical. Implementation: - `rms_norm`: recompute `num_rows` from live `cur_dims`; out dims follow the input. - `embedding_q4gsw`: recompute `num_indices`/`total_blocks`; out dims = indices dims + `[embed_dim]`. - `apply_rotary_emb`: `add_rope_dispatch` now returns its uniform handle; one hook rewrites both the xq and xk dispatches/UBOs for the live S and sets both outputs. - Each keeps its uniform buffer alive via `own_uniform_buffer` (the hook rewrites it) instead of releasing it at build. Mirrors Vulkan per-op `resize_*_node` (recompute sizes + dispatch each execute). No kernel/WGSL/numerics change. Behavior-neutral on static graphs (hook only fires when live dims differ from max). `quantized_linear` and SDPA resize hooks land in following diffs; `prepack` needs none (constants are fixed-size). ghstack-source-id: 399812824 @exported-using-ghexport Differential Revision: [D109906096](https://our.internmc.facebook.com/intern/diff/D109906096/) --- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 89 ++++++++++-- .../webgpu/runtime/ops/rms_norm/RmsNorm.cpp | 51 ++++++- .../runtime/ops/rope/RotaryEmbedding.cpp | 131 +++++++++++++++--- 3 files changed, 241 insertions(+), 30 deletions(-) diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index d24c693e486..6f1febf77d8 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -16,6 +16,7 @@ #include #include #include +#include namespace executorch::backends::webgpu { @@ -36,12 +37,48 @@ static_assert( sizeof(EmbeddingParams) == 32, "EmbeddingParams must be 32 bytes"); -uint64_t numel_of(const std::vector& dims) { - uint64_t n = 1; - for (int64_t d : dims) { - n *= static_cast(d); +// Resize hook body: recompute counts/dispatch; out = indices dims + +// [embed_dim]. +void resize_embedding_q4gsw( + WebGPUGraph& g, + int indices_id, + int out_id, + uint32_t embed_dim, + uint32_t blocks_per_row, + uint32_t gs_u, + uint32_t groups_per_row, + uint32_t bytes_per_row, + uint32_t wg_size, + size_t dispatch_idx, + WGPUBuffer params_buf) { + const auto& id = g.cur_dims(indices_id); + const uint64_t ni = utils::numel_of(id); + if (ni == 0) { + throw std::runtime_error("WebGPU embedding_q4gsw: zero indices"); } - return n; + const uint64_t total_blocks = ni * blocks_per_row; + if (total_blocks > UINT32_MAX) { + throw std::runtime_error( + "WebGPU embedding_q4gsw: total_blocks exceeds uint32"); + } + std::vector od = id; + od.push_back(static_cast(embed_dim)); + g.set_cur_dims(out_id, od); + EmbeddingParams p = {}; + p.embed_dim = embed_dim; + p.blocks_per_row = blocks_per_row; + p.num_indices = static_cast(ni); + p.group_size = gs_u; + p.groups_per_row = groups_per_row; + p.bytes_per_row = bytes_per_row; + p.total_blocks = static_cast(total_blocks); + 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(total_blocks), + wg_size, + "embedding_q4gsw(resize)"); } // arg order mirrors Vulkan EmbeddingQ4gsw.cpp. @@ -99,7 +136,7 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { } // Leading index dims flatten row-major (mirrors Vulkan num_indices). - const uint64_t out_numel = numel_of(out.dims); + const uint64_t out_numel = utils::numel_of(out.dims); const uint32_t num_indices = static_cast(out_numel / embed_dim); const uint32_t groups_per_row = static_cast(scales.dims[1]); const uint32_t blocks_per_row = embed_dim / 32u; @@ -116,9 +153,9 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { } // Per-type byte guards (no runtime dtype): indices i32, weight u8, fp32 rest. - const uint64_t indices_numel = numel_of(indices.dims); - const uint64_t weight_numel = numel_of(weight.dims); - const uint64_t scales_numel = numel_of(scales.dims); + const uint64_t indices_numel = utils::numel_of(indices.dims); + const uint64_t weight_numel = utils::numel_of(weight.dims); + const uint64_t scales_numel = utils::numel_of(scales.dims); if (indices_numel != num_indices || indices.nbytes != indices_numel * sizeof(int32_t) || weight.nbytes != weight_numel || @@ -228,13 +265,43 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch( + const size_t dispatch_idx = graph.add_dispatch( {pipeline, bind_group, workgroup_count, "embedding_q4gsw"}); + // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim]. + const uint32_t gs_u = static_cast(group_size); + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + indices_id, + [indices_id, + out_id, + embed_dim, + blocks_per_row, + gs_u, + groups_per_row, + bytes_per_row, + wg_size, + dispatch_idx, + params_buf](WebGPUGraph& g) { + resize_embedding_q4gsw( + g, + indices_id, + out_id, + embed_dim, + blocks_per_row, + gs_u, + groups_per_row, + bytes_per_row, + wg_size, + dispatch_idx, + params_buf); + }); + wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - wgpuBufferRelease(uniform_buffer); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } } // namespace diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp index e73c6e23a88..c07e58e8ec8 100644 --- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp +++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include #include @@ -31,6 +32,39 @@ struct RmsNormParams { }; static_assert(sizeof(RmsNormParams) == 16, "RmsNormParams must be 16 bytes"); +// Resize hook body: recompute num_rows + rewrite the UBO for the live input. +void resize_rms_norm( + WebGPUGraph& g, + int in_id, + int out_id, + uint32_t row_width, + float epsilon, + size_t dispatch_idx, + WGPUBuffer params_buf) { + const auto& d = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(d); + if (numel % static_cast(row_width) != 0) { + throw std::runtime_error( + "WebGPU rms_norm: numel not a multiple of row_width"); + } + const uint32_t rows = + static_cast(numel / static_cast(row_width)); + if (rows == 0) { + throw std::runtime_error("WebGPU rms_norm: zero rows"); + } + if (rows > 65535u) { + throw std::runtime_error( + "WebGPU rms_norm: num_rows exceeds the 1D dispatch limit (65535)"); + } + RmsNormParams p = {}; + p.num_rows = rows; + p.row_width = row_width; + p.epsilon = epsilon; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = rows; + g.set_cur_dims(out_id, d); +} + void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { // et_vk.rms_norm.default args: [in, weight, eps, out] const int in_id = args.at(0); @@ -187,14 +221,25 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { static_assert( kRmsNormVec4WorkgroupSizeX == 64, "must match @workgroup_size and WG_SIZE in rms_norm_vec4.wgsl"); - graph.add_dispatch({pipeline, bind_group, num_rows}); + const size_t dispatch_idx = + graph.add_dispatch({pipeline, bind_group, num_rows}); + + // Dynamic shapes: recompute num_rows + rewrite the UBO for the live input. + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, row_width, epsilon, dispatch_idx, params_buf]( + WebGPUGraph& g) { + resize_rms_norm( + g, in_id, out_id, row_width, epsilon, dispatch_idx, params_buf); + }); // Release intermediate objects (pipeline and bind_group are kept by dispatch) wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } } // namespace diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index cf4fa0a1ca2..76eafb0e738 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -34,16 +34,15 @@ struct RotaryParams { }; static_assert(sizeof(RotaryParams) == 32, "RotaryParams must be 32 bytes"); -uint64_t numel_of(const std::vector& dims) { - uint64_t n = 1; - for (int64_t d : dims) { - n *= static_cast(d); - } - return n; -} +// A rope dispatch: its param-uniform (rewritten on resize) and its index in the +// graph's dispatch list (so a resize hook can update the workgroup count). +struct RopeDispatch { + WGPUBuffer uniform; + size_t dispatch_index; +}; // Rotate one (x->out) with the shared shader; freqs shared between xq and xk. -void add_rope_dispatch( +RopeDispatch add_rope_dispatch( WebGPUGraph& graph, WGPUDevice device, WGPUComputePipeline pipeline, @@ -58,7 +57,8 @@ void add_rope_dispatch( uint32_t workgroup_count) { const uint32_t half_dim = head_dim / 2u; // out.dims == in.dims (asserted in impl), so this matches the caller's wgc. - const uint32_t num_pairs = static_cast(numel_of(out.dims) / 2u); + const uint32_t num_pairs = + static_cast(utils::numel_of(out.dims) / 2u); RotaryParams params = {}; params.n_heads = n_heads; @@ -101,10 +101,67 @@ void add_rope_dispatch( bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch( + const size_t dispatch_index = graph.add_dispatch( {pipeline, bind_group, workgroup_count, "apply_rotary_emb"}); - wgpuBufferRelease(uniform_buffer); + // Graph owns it so a resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); + return {uniform_buffer, dispatch_index}; +} + +// Resize hook body: recompute S/num_pairs + both dispatches; out follows xq/xk. +void resize_rope( + WebGPUGraph& g, + int xq_id, + int xk_id, + int xq_out_id, + int xk_out_id, + uint32_t n_heads_q, + uint32_t n_heads_k, + uint32_t head_dim, + uint32_t half_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(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); + // pk = pq (seq=s); require k's seq == s, not silently q's. + if (static_cast(kd[kd.size() - 3]) != s) { + throw std::runtime_error( + "apply_rotary_emb(resize): q and k seq lengths differ"); + } + // freqs stay max-allocated; shader indexes by position (S = prefix). + RotaryParams 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); + RotaryParams 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(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(resize)"); + g.set_cur_dims(xq_out_id, qd); + g.set_cur_dims(xk_out_id, kd); } // args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. @@ -164,9 +221,9 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { } // All tensors are fp32; output shapes equal their inputs. - const uint64_t xq_numel = numel_of(xq.dims); - const uint64_t xk_numel = numel_of(xk.dims); - const uint64_t freqs_numel = numel_of(freqs_cos.dims); + 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(seq) * half_dim || xq.nbytes != xq_numel * sizeof(float) || xk.nbytes != xk_numel * sizeof(float) || @@ -246,7 +303,7 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { WGPUComputePipeline pipeline_k = wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - add_rope_dispatch( + RopeDispatch q_disp = add_rope_dispatch( graph, device, pipeline_q, @@ -259,7 +316,7 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { seq, head_dim, xq_wgc); - add_rope_dispatch( + RopeDispatch k_disp = add_rope_dispatch( graph, device, pipeline_k, @@ -272,6 +329,48 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { seq, head_dim, 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; + + // Dynamic shapes: recompute S/num_pairs + both dispatches; out follows xq/xk. + const int xq_out_id = out_list[0]; + const int xk_out_id = out_list[1]; + // Register on both xq and xk so the recompute fires whichever is marked dirty + // (q and k co-resize on S; resize_rope is idempotent, so a double-fire when + // both are dirty is harmless). + auto rope_hook = [xq_id, + xk_id, + xq_out_id, + xk_out_id, + n_heads_q, + n_heads_k, + head_dim, + half_dim, + wg_size, + q_idx, + k_idx, + q_ubuf, + k_ubuf](WebGPUGraph& g) { + resize_rope( + g, + xq_id, + xk_id, + xq_out_id, + xk_out_id, + n_heads_q, + n_heads_k, + head_dim, + half_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); wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); From ae2238921144d9b0d18e8e20884e235ddfeac185 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:54 -0700 Subject: [PATCH 03/14] [ExecuTorch][WebGPU] Dynamic resize hook for linear_q4gsw MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20576 **Make the 4-bit quantized linear serve any live M (rows) from one graph, so a dynamic prefill+decode graph computes correct-size outputs.** **Problem:** `linear_q4gsw` baked its dispatch count, `params.M`, and output shape at `build()` for the max M. On a dynamic-shape graph at a smaller live M (e.g. decode M=1 vs prefill M=S) it would over-dispatch and leave the output sized at the max. **Solution:** - Before: one fixed dispatch sized for the build-time M. - After: a tensor resize hook on the input recomputes the live M from `cur_dims`, rewrites `params.M`, updates the dispatch `workgroup_count_x` for the SAME kernel chosen at build (bicol GEMV / shmem GEMM / register-tiled), and sets the output `cur_dims` (= input dims with the last dim replaced by N). Inert until the input is resized. **Implementation:** - The build-time kernel select (bicol GEMV for M==1, else shmem GEMM for large K/N, else register-tiled) is fixed at build; the hook re-runs `compute_q4gsw_workgroup_count` for whichever of the three the build chose and rewrites the param UBO + output dims for the live M — it does not switch kernels (runtime M-switching is a separate optimization). - `own_uniform_buffer` keeps the param UBO alive so the hook can rewrite it. - Mirrors Vulkan `resize_q4gsw_linear_node` (recompute M-derived dispatch each execute). **Constraints:** Behavior-neutral on static graphs (hook fires only when the input's live M differs from the max). No kernel/WGSL/numerics change. Runtime M-based kernel switching is deliberately out of scope (a later opt diff). Co-authored-with: Claude Code. ghstack-source-id: 399812825 @exported-using-ghexport Differential Revision: [D109906094](https://our.internmc.facebook.com/intern/diff/D109906094/) --- .../ops/quantized_linear/QuantizedLinear.cpp | 168 +++++++++++++----- 1 file changed, 128 insertions(+), 40 deletions(-) diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 44525b54afc..130d81e600f 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include namespace executorch::backends::webgpu { @@ -50,6 +52,62 @@ constexpr int64_t kQ4gswShmemTileN = 32; constexpr uint32_t kQ4gswShmemMinDim = 4096u; constexpr uint32_t kQ4gswShmemNMinDim = 2048u; +// Workgroup count for a linear_q4gsw dispatch (bicol GEMV / shmem GEMM / tiled +// GEMM), with the range/limit guards shared by the build-time path and the +// resize hook. use_gemv/use_shmem_gemm are the build-time routing decision (the +// shader/pipeline is fixed at build); the resize hook re-runs this with live m. +uint32_t compute_q4gsw_workgroup_count( + WGPUDevice device, + bool use_gemv, + bool use_shmem_gemm, + uint32_t m, + uint32_t n, + uint32_t wg_size, + const char* op_name) { + if (use_gemv) { + // bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over + // ceil(N/2) column-pairs (M == 1 on this decode path). + const uint64_t pairs = (static_cast(n) + 1u) / 2u; + if (pairs == 0u || pairs > UINT32_MAX) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": N/2 out of range"); + } + const uint32_t wgc = + utils::clamp_workgroup_count(device, static_cast(pairs)); + if (wgc == 0u) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": zero GEMV dispatch"); + } + return wgc; + } + if (use_shmem_gemm) { + // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit. + const int64_t total_wgs = utils::div_up(m, kQ4gswShmemTileM) * + utils::div_up(n, kQ4gswShmemTileN); + WGPULimits limits = {}; + const uint32_t max_wgs = + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupsPerDimension > 0 + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; + if (total_wgs > static_cast(max_wgs)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": shmem GEMM tile count exceeds the 1D dispatch limit"); + } + return static_cast(total_wgs); + } + const int64_t total_tiles = utils::div_up(m, kQ4gswTileM) * + utils::div_up(n, kQ4gswTileN); + if (total_tiles > static_cast(UINT32_MAX)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": tile count exceeds the 1D dispatch limit"); + } + return utils::compute_1d_workgroup_count( + device, static_cast(total_tiles), wg_size, op_name); +} + // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out]. void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); @@ -137,44 +195,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL : kQ4gswLinearWGSL; - uint32_t workgroup_count; - if (use_gemv) { - // bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over - // ceil(N/2) column-pairs (M == 1 on this decode path). - const uint64_t pairs = (static_cast(N) + 1u) / 2u; - if (pairs == 0u || pairs > UINT32_MAX) { - throw std::runtime_error("WebGPU linear_q4gsw: N/2 out of range"); - } - workgroup_count = - utils::clamp_workgroup_count(device, static_cast(pairs)); - if (workgroup_count == 0u) { - throw std::runtime_error("WebGPU linear_q4gsw: zero GEMV dispatch"); - } - } else if (use_shmem_gemm) { - // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit. - const int64_t total_wgs = utils::div_up(M, kQ4gswShmemTileM) * - utils::div_up(N, kQ4gswShmemTileN); - WGPULimits limits = {}; - const uint32_t max_wgs = - wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && - limits.maxComputeWorkgroupsPerDimension > 0 - ? limits.maxComputeWorkgroupsPerDimension - : 65535u; - if (total_wgs > static_cast(max_wgs)) { - throw std::runtime_error( - "WebGPU linear_q4gsw: shmem GEMM tile count exceeds the 1D dispatch limit"); - } - workgroup_count = static_cast(total_wgs); - } else { - const int64_t total_tiles = utils::div_up(M, kQ4gswTileM) * - utils::div_up(N, kQ4gswTileN); - if (total_tiles > static_cast(UINT32_MAX)) { - throw std::runtime_error( - "WebGPU linear_q4gsw: tile count exceeds the 1D dispatch limit"); - } - workgroup_count = utils::compute_1d_workgroup_count( - device, static_cast(total_tiles), wg_size, "linear_q4gsw"); - } + const uint32_t workgroup_count = compute_q4gsw_workgroup_count( + device, use_gemv, use_shmem_gemm, M, N, wg_size, "linear_q4gsw"); // Optional bias: real buffer if present, else a dummy for the fixed layout. uint32_t has_bias = 0; @@ -287,12 +309,78 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count, "linear_q4gsw"}); + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count, "linear_q4gsw"}); + + // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and + // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs + // the same path's workgroup-count formula with the live m. + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + M, + K, + N, + K_packed, + gs, + padded_N, + has_bias, + wg_size, + use_gemv, + use_shmem_gemm, + dispatch_idx, + uniform_buffer](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.empty()) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): empty input dims"); + } + const uint64_t numel = utils::numel_of(d); + if (numel % static_cast(K) != 0u) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live input numel not a multiple " + "of K"); + } + const uint32_t m = + static_cast(numel / static_cast(K)); + if (m == 0u) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); + } + // Buffers/bind-groups were sized for the build-time max M; a larger + // live M would write out of bounds. + if (m > M) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); + } + const uint32_t wgc = compute_q4gsw_workgroup_count( + g.device(), + use_gemv, + use_shmem_gemm, + m, + N, + wg_size, + "linear_q4gsw(resize)"); + Q4gswParams p = {}; + p.M = m; + p.N = N; + p.K = K; + p.K_packed = K_packed; + p.group_size = gs; + p.padded_N = padded_N; + p.has_bias = has_bias; + wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc; + std::vector od(d.begin(), d.end()); + od.back() = static_cast(N); + g.set_cur_dims(out_id, od); + }); wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - wgpuBufferRelease(uniform_buffer); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } } // namespace From fd16aade5b1a00115a993d2a90fe145f80a5ea2c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:54 -0700 Subject: [PATCH 04/14] [ExecuTorch][WebGPU] Dynamic resize hooks for add and mul Pull Request resolved: https://github.com/pytorch/executorch/pull/20577 **Make the elementwise add and mul ops serve any live shape from one graph.** **Problem:** `aten.add.Tensor` and `aten.mul.Tensor` baked their element count + param UBO(s) + output shape at `build()` for the max shape. On a dynamic-shape graph at a smaller live shape they would over-dispatch and leave the output sized at the max. **Solution:** - Before: one fixed dispatch sized for the build-time shape. - After: each registers a resize hook on BOTH operands (the dynamic one may be either operand by arg order). The hook recomputes the live element count, rewrites the param UBO(s), updates the dispatch `workgroup_count_x`, and sets the output `cur_dims`. Inert until an operand is resized. **Implementation:** - `add`: out follows the larger operand (robust when one input is a static residual and the other is the dynamic-S tensor); rewrites `AddParams`. - `mul`: recomputes the broadcast output shape and rebuilds all three `TensorMeta` UBOs via `fill_tensor_meta_broadcast`. - Each keeps its uniform buffer(s) alive via `own_uniform_buffer` instead of releasing at build. - Mirrors Vulkan per-op `resize_*_node` (recompute sizes + dispatch each execute). **Constraints:** Behavior-neutral on static graphs (the hook fires only when an operand's live shape differs from the max). No kernel/WGSL/numerics change. Co-authored-with: Claude Code. ghstack-source-id: 399812828 @exported-using-ghexport Differential Revision: [D109906093](https://our.internmc.facebook.com/intern/diff/D109906093/) --- backends/webgpu/runtime/ops/add/BinaryOp.cpp | 39 ++++++++++++++- backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 50 ++++++++++++++++++-- 2 files changed, 82 insertions(+), 7 deletions(-) diff --git a/backends/webgpu/runtime/ops/add/BinaryOp.cpp b/backends/webgpu/runtime/ops/add/BinaryOp.cpp index 578799a9c38..8c56ad6c15d 100644 --- a/backends/webgpu/runtime/ops/add/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/add/BinaryOp.cpp @@ -159,13 +159,48 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); graph.add_dispatch({pipeline, bind_group, workgroup_count}); + const size_t dispatch_idx = graph.num_dispatches() - 1; + + // Dynamic shapes: recompute numel/dispatch; out follows the larger operand. + WGPUBuffer params_buf = uniform_buffer; + auto add_resize = [in1_id, + in2_id, + out_id, + alpha, + wg_size, + dispatch_idx, + params_buf](WebGPUGraph& g) { + const auto& d1 = g.cur_dims(in1_id); + const auto& d2 = g.cur_dims(in2_id); + const uint64_t n1 = utils::numel_of(d1); + const uint64_t n2 = utils::numel_of(d2); + const uint64_t numel = n2 > n1 ? n2 : n1; + const uint64_t n_min = n2 > n1 ? n1 : n2; + // The flat add follows the larger operand and broadcasts the smaller; valid + // only when the smaller tiles evenly into it (rejects e.g. [4,1] vs [1,3], + // whose true [4,3] result this flat kernel cannot produce). + if (n_min == 0u || numel % n_min != 0u) { + throw std::runtime_error( + "add(resize): operands are not broadcast-compatible by numel"); + } + g.set_cur_dims(out_id, n2 > n1 ? d2 : d1); + AddParams p = {}; + p.num_elements = static_cast(numel); + p.alpha = alpha; + 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, "add(resize)"); + }; + graph.add_tensor_resize_hook(in1_id, add_resize); + graph.add_tensor_resize_hook(in2_id, add_resize); // Release intermediate objects (pipeline and bind_group are kept by dispatch) wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + // Graph owns it so a resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } } // namespace diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index 007b7b2d8da..2ccb6c0e1bf 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -14,6 +14,7 @@ #include +#include #include #include @@ -164,15 +165,54 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + const size_t dispatch_idx = + graph.add_dispatch({pipeline, bind_group, workgroup_count}); + + // 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; + auto mul_resize = + [in1_id, in2_id, out_id, wg_size, dispatch_idx, o_buf, a_buf, b_buf]( + WebGPUGraph& g) { + const auto& a = g.cur_dims(in1_id); + const auto& b = g.cur_dims(in2_id); + const size_t r = std::max(a.size(), b.size()); + std::vector out_d(r, 1); + for (size_t i = 0; i < r; i++) { + const int64_t av = (i + a.size() < r) ? 1 : a[i - (r - a.size())]; + const int64_t bv = (i + b.size() < r) ? 1 : b[i - (r - b.size())]; + if (av != bv && av != 1 && bv != 1) { + throw std::runtime_error( + "mul(resize): operands are not broadcast-compatible"); + } + out_d[i] = av > bv ? av : bv; + } + g.set_cur_dims(out_id, out_d); + const uint32_t out_ndim = static_cast(r); + WebGPUTensor ta, tb, to; + ta.dims = a; + tb.dims = b; + to.dims = out_d; + TensorMeta om, am, bm; + fill_tensor_meta_broadcast(to, out_ndim, &om); + fill_tensor_meta_broadcast(ta, out_ndim, &am); + fill_tensor_meta_broadcast(tb, out_ndim, &bm); + wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om)); + wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); + wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), om.numel, wg_size, "mul(resize)"); + }; + graph.add_tensor_resize_hook(in1_id, mul_resize); + graph.add_tensor_resize_hook(in2_id, mul_resize); wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - // Drop our refs; the bind group keeps the uniforms alive until release. - wgpuBufferRelease(out_meta_buf); - wgpuBufferRelease(in1_meta_buf); - wgpuBufferRelease(in2_meta_buf); + // 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); + graph.own_uniform_buffer(in2_meta_buf); } } // namespace From 50663c2438c07e864bf315714d0ea4ad9e3f84c5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:55 -0700 Subject: [PATCH 05/14] [ExecuTorch][WebGPU] Dynamic resize hooks for sigmoid and select_copy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20578 **Make sigmoid and select_copy serve any live shape from one graph; fix select's last-token index under dynamic shapes.** **Problem:** Both ops baked their dispatch/params/output shape at `build()` for the max shape. `select_copy` was worse: a negative index (e.g. `-1` for the last token) was normalized against the build-time MAX dim, so at a smaller live S it selected a stale/zero position past the live data — producing wrong (often zero) output. **Solution:** - `sigmoid` (generic `add_unary_op`): a resize hook recomputes `num_elements`/dispatch and sets the output `cur_dims` (shape-preserving). - `select_copy`: KEEP the raw (possibly negative) index at build; a resize hook re-resolves it against the LIVE dim, recomputes the output dims (= input minus `dim`), rebuilds the out/in `TensorMeta` UBOs and the dispatch. - Both keep their uniform buffer(s) alive via `own_uniform_buffer`. **Implementation:** - The select out/in meta is rebuilt from synthetic `WebGPUTensor{dims}` via `fill_tensor_meta` (reads only `.dims`). - Mirrors Vulkan per-op `resize_*_node`. **Constraints:** Behavior-neutral on static graphs (hooks fire only when an input's live shape differs from the max). No kernel/WGSL/numerics change. Co-authored-with: Claude Code. ghstack-source-id: 399812832 @exported-using-ghexport Differential Revision: [D109906095](https://our.internmc.facebook.com/intern/diff/D109906095/) --- backends/webgpu/runtime/ops/select/Select.cpp | 71 ++++++++++++++++--- .../webgpu/runtime/ops/sigmoid/UnaryOp.cpp | 26 ++++++- 2 files changed, 85 insertions(+), 12 deletions(-) diff --git a/backends/webgpu/runtime/ops/select/Select.cpp b/backends/webgpu/runtime/ops/select/Select.cpp index 5686bbc79c0..5d3e8103b7d 100644 --- a/backends/webgpu/runtime/ops/select/Select.cpp +++ b/backends/webgpu/runtime/ops/select/Select.cpp @@ -37,6 +37,19 @@ int64_t read_scalar(WebGPUGraph& graph, int id, const char* what) { throw std::runtime_error(std::string("select: dynamic/unsupported ") + what); } +// Build a TensorMeta from live dims, write it to buf, return numel. +uint32_t write_meta_from_dims( + WebGPUGraph& g, + WGPUBuffer buf, + const std::vector& dims) { + WebGPUTensor t; + t.dims = dims; + TensorMeta m; + fill_tensor_meta(t, &m); + wgpuQueueWriteBuffer(g.queue(), buf, 0, &m, sizeof(m)); + return m.numel; +} + void select_impl(WebGPUGraph& graph, const std::vector& args) { // args: [self, dim, index, out]; output rank = in rank - 1. const int in_id = args.at(0); @@ -58,10 +71,9 @@ void select_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("select: dim out of range"); } const int64_t in_size = in_tensor.dims[dim]; - int64_t index = read_scalar(graph, args.at(2), "index"); - if (index < 0) { - index += in_size; - } + // Keep the RAW index: -1 normalizes against the LIVE dim (the resize hook). + const int64_t raw_index = read_scalar(graph, args.at(2), "index"); + int64_t index = raw_index < 0 ? raw_index + in_size : raw_index; if (index < 0 || index >= in_size) { throw std::runtime_error("select: index out of range"); } @@ -164,15 +176,56 @@ void select_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + const size_t dispatch_idx = + graph.add_dispatch({pipeline, bind_group, workgroup_count}); + + // Dynamic shapes: out = in minus `dim`; re-resolve index, meta, dispatch. + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + dim, + raw_index, + out_meta_buf, + in_meta_buf, + params_buf, + wg_size, + dispatch_idx](WebGPUGraph& g) { + const auto& ind = g.cur_dims(in_id); + if (dim < 0 || dim >= static_cast(ind.size())) { + throw std::runtime_error("select(resize): dim out of range"); + } + const int64_t live_in_size = ind[dim]; + int64_t idx = raw_index < 0 ? raw_index + live_in_size : raw_index; + if (idx < 0 || idx >= live_in_size) { + throw std::runtime_error("select(resize): index out of range"); + } + std::vector od; + od.reserve(ind.size() - 1); + for (size_t k = 0; k < ind.size(); k++) { + if (static_cast(k) != dim) { + od.push_back(ind[k]); + } + } + g.set_cur_dims(out_id, od); + const uint32_t out_numel = write_meta_from_dims(g, out_meta_buf, od); + write_meta_from_dims(g, in_meta_buf, ind); + SelectParams p = {}; + p.dim = static_cast(dim); + p.index = static_cast(idx); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), out_numel, wg_size, "select(resize)"); + }); 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); - wgpuBufferRelease(params_buf); + // 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); + graph.own_uniform_buffer(params_buf); } } // namespace diff --git a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp index 4d1a087cae5..7c8b8c0e9ce 100644 --- a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp +++ b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp @@ -135,14 +135,34 @@ void add_unary_op( bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + const size_t dispatch_idx = + graph.add_dispatch({pipeline, 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); + 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)"); + }); // Release intermediates (pipeline + bind_group are kept by dispatch). wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } void sigmoid_impl(WebGPUGraph& graph, const std::vector& args) { From 72933977a858db202ea67c311f6264cf9aebf233 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:55 -0700 Subject: [PATCH 06/14] [ExecuTorch][WebGPU] Dynamic resize hook for view_copy Pull Request resolved: https://github.com/pytorch/executorch/pull/20579 **Make `view_copy` track the live sequence length under dynamic shapes.** **Problem:** `view_copy` lowers to a flat DMA buffer copy (`add_buffer_copy`) sized at the build-time max shape. With one dynamic graph serving any seq-len S (prefill S=K, decode S=1), the copy moved the full max-S byte count and the output kept its max dims, so a downstream consumer read a live shape that was too large. **Solution:** register a tensor resize hook on the input so the copy follows the live input numel (a view preserves numel). - Before: `copy_nbytes` and the output dims are fixed at the serialized max. - After: the hook recomputes the live numel from `cur_dims(in)`, scales the single dynamic output dim to preserve numel, sets the output `cur_dims`, and rewrites the Copy dispatch's `copy_nbytes`. **Implementation:** - Keep the existing DMA path (`Kind::Copy`); the hook only rewrites `copy_nbytes` via `dispatch_at`, no new kernel. - Handle the aliased in/out fast path (no copy emitted) by still setting the output `cur_dims` so the resize cascade reaches consumers. - Mirrors Vulkan's `view_buffer` contiguous fast path; numel-preserving like the other dynamic-shape op hooks. **Constraints:** inert on a static graph (`cur_dims == dims`), so byte-identical to the prior behavior; fp32-only and numel-preserving invariants unchanged. Co-authored-with: Claude Code. ghstack-source-id: 399812833 @exported-using-ghexport Differential Revision: [D109906098](https://our.internmc.facebook.com/intern/diff/D109906098/) --- .../webgpu/runtime/ops/view_copy/ViewCopy.cpp | 47 +++++++++++++++++-- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp index 67119472643..eb274481893 100644 --- a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp +++ b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include @@ -38,11 +39,49 @@ void add_flat_copy(WebGPUGraph& graph, int in_id, int out_id) { } // Aliased in/out already in place; CopyBufferToBuffer rejects src == dst. - if (in_tensor.buffer == out_tensor.buffer) { - return; - } + const bool aliased = in_tensor.buffer == out_tensor.buffer; + const size_t dispatch_idx = aliased + ? 0 + : graph.add_buffer_copy( + in_tensor.buffer, out_tensor.buffer, out_tensor.nbytes); - graph.add_buffer_copy(in_tensor.buffer, out_tensor.buffer, out_tensor.nbytes); + // Dynamic shapes: view preserves numel; copy_nbytes + out dims track live in. + std::vector out_max = out_tensor.dims; + graph.add_tensor_resize_hook( + in_id, [in_id, out_id, out_max, dispatch_idx, aliased](WebGPUGraph& g) { + const uint64_t target = utils::numel_of(g.cur_dims(in_id)); + std::vector od = out_max; + const uint64_t maxnumel = utils::numel_of(out_max); + if (maxnumel != target) { + bool resolved = false; + // Assumes one dynamic dim; picks the leftmost numel-divisible. + for (size_t d = 0; d < od.size(); d++) { + if (out_max[d] <= 0) { + continue; + } + const uint64_t rest = maxnumel / static_cast(out_max[d]); + if (rest != 0 && target % rest == 0) { + const uint64_t nd = target / rest; + if (nd <= static_cast(out_max[d])) { + od[d] = static_cast(nd); + resolved = true; + break; + } + } + } + // Fail loud: a silent miss would leave od at max while copy_nbytes + // shrinks to the live size, desyncing consumers from the real copy. + if (!resolved) { + throw std::runtime_error( + "view_copy(resize): could not resolve live output shape"); + } + } + g.set_cur_dims(out_id, od); + if (!aliased) { + g.dispatch_at(dispatch_idx).copy_nbytes = + static_cast(target) * sizeof(float); + } + }); } namespace { From 2b2d6326cb2b4538f61acc4d857e1fffb484f8d5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:56 -0700 Subject: [PATCH 07/14] [ExecuTorch][WebGPU] Dynamic resize hook for SDPA (live seq-len S) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20580 **Make `sdpa_with_kv_cache` serve any live seq-len S from one graph (batched prefill S=K and decode S=1).** **Problem:** the existing dynamic path only reacted to a live `input_pos` (decode), with S captured at build time. It rewrote the QK dispatch (which depends on `context_len`) but left `update_cache`, softmax, and AV sized for the build-time S. Under a dynamic seq-len S (one graph serving prefill and decode), `kv_numel`, the QK/AV tile grids, and the softmax row count all depend on S and were stale. **Solution:** a single recompute hook driven by either a live S (q tensor resize) or a live `input_pos` (SymInt), recomputing every per-step quantity from the live shape. - Before: hook keyed only on `input_pos`; recomputes ctx + QK count; S fixed. - After: hook keyed on q (always) and `input_pos` (when SymInt); reads live S from `cur_dims(q)` and live pos, recomputes all five dispatches' counts + UBOs (`update_cache` K/V, QK, softmax, AV), and sets the output `cur_dims` to q's. **Implementation:** - Capture the `update_cache`/softmax/AV dispatch indices (previously only QK) so their workgroup counts can be rewritten per step. - QK/AV workgroup counts use the landed register-tiled grids (`Hq*ceil(S/TM)*ceil(ctx-or-D/TN)`); softmax is one workgroup per `Hq*S` row. - Register the hook on q unconditionally — inert until q is resized, so a static graph is byte-identical. - Mirrors Vulkan `DynamicDispatchNode` (recompute workgroups per execute); scratch is sized at build (S=max, ctx=Cmax) so buffers never move and bind groups stay valid. **Constraints:** fp32-only, batch=1, GQA, `is_causal=true`, `D%4==0` invariants unchanged; the static / decode-only paths are unaffected (the q hook never fires without a resize). Co-authored-with: Claude Code. ghstack-source-id: 399812834 @exported-using-ghexport Differential Revision: [D109906097](https://our.internmc.facebook.com/intern/diff/D109906097/) --- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 190 +++++++++++++--------- 1 file changed, 117 insertions(+), 73 deletions(-) diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 74e9e265162..c55f5cfa22f 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -240,7 +240,7 @@ static WGPUBuffer record_update_cache_dispatch( uint32_t kv_dst_offset, uint64_t cache_numel, uint32_t uc_wg, - bool dynamic_pos, + bool retain_uniform, const char* label) { const uint32_t wgc = utils::compute_1d_workgroup_count( device, static_cast(kv_numel), uc_wg, label); @@ -258,7 +258,7 @@ static WGPUBuffer record_update_cache_dispatch( sizeof(uc), wgc, uc_wg, - dynamic_pos, + retain_uniform, "update_cache"); return ubuf; } @@ -417,7 +417,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { // Dynamic input_pos: the resize hook rewrites these per step. WGPUBuffer uc_k_buf = nullptr, uc_v_buf = nullptr, qk_buf = nullptr, softmax_buf = nullptr, av_buf = nullptr; - size_t qk_idx = 0; + size_t qk_idx = 0, uc_k_idx = 0, uc_v_idx = 0, softmax_idx = 0, av_idx = 0; const WGPUDevice device = graph.device(); const uint32_t uc_wg = @@ -442,7 +442,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { kv_dst_offset, numel(k_cache), uc_wg, - dynamic_pos, + true, "update_cache(K)"); uc_v_buf = record_update_cache_dispatch( graph, @@ -453,8 +453,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { kv_dst_offset, numel(v_cache), uc_wg, - dynamic_pos, + true, "update_cache(V)"); + uc_k_idx = graph.num_dispatches() - 2; + uc_v_idx = graph.num_dispatches() - 1; // FlashDecoding decode (S==1, static pos). Shapes FD can't handle (head dim // > kSdpaFdMaxHeadDim) fall through to the materialized path below. @@ -494,7 +496,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { sizeof(p), wgc, qk_wg, - dynamic_pos, + true, "sdpa_compute_attn_weights"); qk_buf = ubuf; qk_idx = graph.num_dispatches() - 1; @@ -518,9 +520,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { sizeof(p), wgc, 0, - dynamic_pos, + true, "sdpa_softmax"); softmax_buf = ubuf; + softmax_idx = graph.num_dispatches() - 1; } // --- Dispatch 5: AV -> out. One thread per TM x TN tile. @@ -544,77 +547,118 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { sizeof(p), wgc, av_wg, - dynamic_pos, + true, "sdpa_compute_out"); av_buf = ubuf; + av_idx = graph.num_dispatches() - 1; } - // Per-step recompute hook; mirrors Vulkan DynamicDispatchNode. + // Per-step recompute: live S (q resize) or input_pos (SymInt); inert if + // static. + const int64_t pos_const = input_pos; + auto sdpa_resize = [q_id, + qn, + S, + out_id, + dynamic_pos, + input_pos_id, + pos_const, + Hq, + Hkv, + D, + Cmax, + g, + scale, + qk_idx, + uc_k_idx, + uc_v_idx, + softmax_idx, + av_idx, + uc_wg, + qk_wg, + av_wg, + uc_k_buf, + uc_v_buf, + qk_buf, + softmax_buf, + av_buf](WebGPUGraph& gr) { + const int64_t s = gr.cur_dims(q_id)[qn - 3]; + const int64_t pos = dynamic_pos + ? static_cast(gr.read_symint(input_pos_id)) + : pos_const; + if (s <= 0 || pos < 0) { + throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); + } + // Scratch (attn_weights/softmax) is sized at build for S=max; a larger live + // S would overrun it. Make that invariant load-bearing. + if (s > S) { + throw std::runtime_error( + "WebGPU sdpa: live S exceeds the build-time max (scratch capacity)"); + } + const int64_t ctx = s + pos; + if (ctx <= 0 || ctx > Cmax) { + throw std::runtime_error( + "WebGPU sdpa: context_len exceeds cache capacity"); + } + const uint32_t kv_off = static_cast( + static_cast(pos) * static_cast(Hkv) * + static_cast(D)); + const uint64_t aw_floats = static_cast(Hq) * + static_cast(s) * static_cast(ctx); + if (aw_floats > UINT32_MAX) { + throw std::runtime_error("WebGPU sdpa: Hq*S*context_len exceeds uint32"); + } + const uint64_t kv_numel = static_cast(s) * + static_cast(Hkv) * static_cast(D); + if (kv_numel > UINT32_MAX) { + throw std::runtime_error("WebGPU sdpa: S*Hkv*D exceeds uint32"); + } + const uint64_t k_cache_numel = static_cast(Cmax) * + static_cast(Hkv) * static_cast(D); + + // update_cache K/V: dispatch (kv_numel) + dst offset scale with live S/pos. + UpdateCacheParams uc = + make_update_cache_params(kv_numel, kv_off, k_cache_numel); + wgpuQueueWriteBuffer(gr.queue(), uc_k_buf, 0, &uc, sizeof(uc)); + wgpuQueueWriteBuffer(gr.queue(), uc_v_buf, 0, &uc, sizeof(uc)); + const uint32_t uc_wgc = utils::compute_1d_workgroup_count( + gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"); + gr.dispatch_at(uc_k_idx).workgroup_count_x = uc_wgc; + gr.dispatch_at(uc_v_idx).workgroup_count_x = uc_wgc; + + // QK: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(ctx/TN). + AttnWeightsParams qp = + make_attn_weights_params(s, Hq, Hkv, D, ctx, pos, g, scale); + wgpuQueueWriteBuffer(gr.queue(), qk_buf, 0, &qp, sizeof(qp)); + const int64_t qk_tiles = + Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(ctx, kSdpaTileN); + gr.dispatch_at(qk_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + + // softmax: one workgroup per (h,s) row. + SoftmaxParams sp = make_softmax_params(Hq, s, ctx); + wgpuQueueWriteBuffer(gr.queue(), softmax_buf, 0, &sp, sizeof(sp)); + gr.dispatch_at(softmax_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + gr.device(), static_cast(Hq * s), 1, "softmax(resize)"); + + // AV: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(D/TN). + ComputeOutParams op = make_compute_out_params(s, Hq, Hkv, D, ctx, g); + wgpuQueueWriteBuffer(gr.queue(), av_buf, 0, &op, sizeof(op)); + const int64_t av_tiles = + Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(D, kSdpaTileN); + gr.dispatch_at(av_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + + // Output attn has the same shape as q: [.., S, Hq, D]. + gr.set_cur_dims(out_id, gr.cur_dims(q_id)); + }; + // q and input_pos share one idempotent recompute; a double-fire is harmless. + graph.add_tensor_resize_hook(q_id, sdpa_resize); if (dynamic_pos) { - graph.add_resize_hook( - input_pos_id, - [input_pos_id, - S, - Hq, - Hkv, - D, - Cmax, - g, - scale, - qk_idx, - qk_wg, - uc_k_buf, - uc_v_buf, - qk_buf, - softmax_buf, - av_buf](WebGPUGraph& gr) { - const int32_t pos = gr.read_symint(input_pos_id); - if (pos < 0) { - throw std::runtime_error( - "WebGPU sdpa: input_pos must be non-negative"); - } - const int64_t ctx = S + pos; - if (ctx <= 0 || ctx > Cmax) { - throw std::runtime_error( - "WebGPU sdpa: context_len exceeds cache capacity"); - } - const uint32_t kv_off = static_cast( - static_cast(pos) * static_cast(Hkv) * - static_cast(D)); - const uint64_t aw_floats = static_cast(Hq) * - static_cast(S) * static_cast(ctx); - if (aw_floats > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa: Hq*S*context_len exceeds uint32 max"); - } - const uint64_t kv_numel = static_cast(S) * - static_cast(Hkv) * static_cast(D); - const uint64_t k_cache_numel = static_cast(Cmax) * - static_cast(Hkv) * static_cast(D); - - UpdateCacheParams uc = - make_update_cache_params(kv_numel, kv_off, k_cache_numel); - wgpuQueueWriteBuffer(gr.queue(), uc_k_buf, 0, &uc, sizeof(uc)); - wgpuQueueWriteBuffer(gr.queue(), uc_v_buf, 0, &uc, sizeof(uc)); - - AttnWeightsParams qp = - make_attn_weights_params(S, Hq, Hkv, D, ctx, pos, g, scale); - wgpuQueueWriteBuffer(gr.queue(), qk_buf, 0, &qp, sizeof(qp)); - const int64_t qk_tiles = Hq * utils::div_up(S, kSdpaTileM) * - utils::div_up(ctx, kSdpaTileN); - const uint32_t qk_wgc = utils::compute_1d_workgroup_count( - gr.device(), - static_cast(qk_tiles), - qk_wg, - "QK(resize)"); - gr.dispatch_at(qk_idx).workgroup_count_x = qk_wgc; - - SoftmaxParams sp = make_softmax_params(Hq, S, ctx); - wgpuQueueWriteBuffer(gr.queue(), softmax_buf, 0, &sp, sizeof(sp)); - - ComputeOutParams op = make_compute_out_params(S, Hq, Hkv, D, ctx, g); - wgpuQueueWriteBuffer(gr.queue(), av_buf, 0, &op, sizeof(op)); - }); + graph.add_resize_hook(input_pos_id, sdpa_resize); } } From cf9d401919d6a670e10cafcdbe0b87b376a30132 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:56 -0700 Subject: [PATCH 08/14] [ExecuTorch][WebGPU] Dynamic resize hook for slice_copy (dynamic start/end) Pull Request resolved: https://github.com/pytorch/executorch/pull/20581 **Make `slice_copy` support a dynamic gather range so the RoPE-freqs slice `[input_pos : input_pos + S]` works under one dynamic graph.** **Problem:** the static slice handler read `start` via a scalar reader that throws on a SymInt and ignored `end` (output length baked AOT). The RoPE-freqs slice uses a SymInt `input_pos` for start and a live S for the range, so the static op could neither build nor resize for it. **Solution:** read start/end as possibly-dynamic SymInts and add a resize hook that recomputes the gather offset and live output length each step. - Before: `start` is a static scalar (SymInt throws); `end` ignored; output length fixed at the serialized max. - After: `start`/`end` read via a SymInt-aware reader; a hook recomputes `out[dim] = (end - start + step - 1) / step`, rewrites `out_meta`/`in_meta`/`params` UBOs + the dispatch count, and sets the output `cur_dims`. **Implementation:** - Hook registered on the `start`/`end` value-ids when they are SymInts and on the input tensor always (inert until resized, so a static slice is byte-identical). - Output/input `TensorMeta` rebuilt from live dims; `dim`/`step` stay static. - Keep the uniforms alive via `own_uniform_buffer` so the hook can rewrite them. - Mirrors Vulkan `resize_slice_copy_node`. **Constraints:** fp32-only; `dim`/`step` static; numerics + layout unchanged; inert on a static graph. NOTE (stacking): this diff sits on top of the in-review `slice_copy` op (D108793168); rebase onto it once that op lands on master. Co-authored-with: Claude Code. ghstack-source-id: 399812835 @exported-using-ghexport Differential Revision: [D109906092](https://our.internmc.facebook.com/intern/diff/D109906092/) --- backends/webgpu/runtime/ops/slice/Slice.cpp | 104 +++++++++++++++++--- 1 file changed, 92 insertions(+), 12 deletions(-) diff --git a/backends/webgpu/runtime/ops/slice/Slice.cpp b/backends/webgpu/runtime/ops/slice/Slice.cpp index 1d4406bbd1a..261020e05d5 100644 --- a/backends/webgpu/runtime/ops/slice/Slice.cpp +++ b/backends/webgpu/runtime/ops/slice/Slice.cpp @@ -46,9 +46,40 @@ 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. +int64_t read_index(WebGPUGraph& graph, int id, int64_t dflt) { + switch (graph.get_value_type(id)) { + case WebGPUGraph::ValueType::SymInt: + return graph.read_symint(id); + case WebGPUGraph::ValueType::Int: { + const int64_t v = graph.get_int(id); + return v == INT64_MAX ? dflt : v; + } + case WebGPUGraph::ValueType::Null: + return dflt; + default: + throw std::runtime_error("slice: dynamic/unsupported start/end index"); + } +} + +bool is_symint(WebGPUGraph& graph, int id) { + return graph.get_value_type(id) == WebGPUGraph::ValueType::SymInt; +} + +// Clamp + normalize a (possibly negative) index into [0, size]. +int64_t norm_clamp(int64_t idx, int64_t size) { + if (idx < 0) { + idx += size; + } + return idx < 0 ? 0 : (idx > size ? size : idx); +} + void slice_impl(WebGPUGraph& graph, const std::vector& args) { - // args: [self, dim, start, end, step, out]; end unread (out shape is AOT). + // args: [self, dim, start, end, step, out]. start/end may be dynamic SymInts; + // a resize hook recomputes the live extent on `dim` (out[dim] / cur_dims). const int in_id = args.at(0); + const int start_id = args.at(2); + const int end_id = args.at(3); const int out_id = args.at(5); WGPUDevice device = graph.device(); @@ -63,17 +94,14 @@ void slice_impl(WebGPUGraph& graph, const std::vector& args) { if (dim < 0 || dim >= in_ndim) { throw std::runtime_error("slice: dim out of range"); } - const int64_t in_size = in_tensor.dims[dim]; - int64_t start = read_scalar(graph, args.at(2), 0, "start"); - if (start < 0) { - start += in_size; - } - // Clamp start to [0, in_size] (guards the gather offset; out size is AOT). - start = start < 0 ? 0 : (start > in_size ? in_size : start); const int64_t step = read_scalar(graph, args.at(4), 1, "step"); if (step < 1) { throw std::runtime_error("slice: step must be >= 1"); } + // start/end may be dynamic SymInts; seed from current (max) dims, the resize + // hook recomputes live. Clamp guards the gather offset. + const int64_t in_size = in_tensor.dims[dim]; + const int64_t start = norm_clamp(read_index(graph, start_id, 0), in_size); TensorMeta out_meta; TensorMeta in_meta; @@ -175,14 +203,66 @@ void slice_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); graph.add_dispatch({pipeline, bind_group, workgroup_count}); + const size_t dispatch_idx = graph.num_dispatches() - 1; + + // Dynamic shapes: live start/end -> out[dim] len + meta/params/dispatch. + auto recompute = [in_id, + out_id, + start_id, + end_id, + dim, + step, + wg_size, + out_meta_buf, + in_meta_buf, + params_buf, + dispatch_idx](WebGPUGraph& g) { + const auto& in_dims = g.cur_dims(in_id); + const int64_t live_in_size = in_dims[dim]; + const int64_t start = norm_clamp(read_index(g, start_id, 0), live_in_size); + const int64_t end = + norm_clamp(read_index(g, end_id, live_in_size), live_in_size); + const int64_t len = end > start ? (end - start + step - 1) / step : 0; + + // Out dims = live input dims (mirror Vulkan resize_slice_copy_node). + std::vector od = in_dims; + od[dim] = len; + g.set_cur_dims(out_id, od); + + WebGPUTensor t_out; + t_out.dims = od; + WebGPUTensor t_in; + t_in.dims = in_dims; + TensorMeta om; + TensorMeta im; + fill_tensor_meta(t_out, &om); + fill_tensor_meta(t_in, &im); + wgpuQueueWriteBuffer(g.queue(), out_meta_buf, 0, &om, sizeof(om)); + wgpuQueueWriteBuffer(g.queue(), in_meta_buf, 0, &im, sizeof(im)); + SliceParams p = {}; + p.dim = static_cast(dim); + p.start = static_cast(start); + p.step = static_cast(step); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), om.numel, wg_size, "slice(resize)"); + }; + if (is_symint(graph, start_id)) { + graph.add_resize_hook(start_id, recompute); + } + if (is_symint(graph, end_id) && end_id != start_id) { + graph.add_resize_hook(end_id, recompute); + } + graph.add_tensor_resize_hook(in_id, recompute); 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); - wgpuBufferRelease(params_buf); + // 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); + graph.own_uniform_buffer(params_buf); } } // namespace From cff45e770f025957c260388bd2504f325845c49a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:56 -0700 Subject: [PATCH 09/14] [ExecuTorch][WebGPU] Dynamic-shape integration test (allocate-at-max + per-op resize) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20582 **End-to-end validation that one graph built at the upper-bound seq-len serves every smaller live shape, matching the torch golden.** **Problem:** the dynamic-resize engine (allocate-at-max buffers + per-op resize hooks + output resize) had unit-level reasoning but no single oracle proving a graph built at S=MAX runs correctly at Sadd cascade), `rms(x)*x` (mul). - Cases I-L: dynamic `linear_q4gsw` (GEMM at several M), `sdpa_with_kv_cache` (GQA prefill at several S), `embedding_q4gsw` (int64 ids), `apply_rotary_emb` (two outputs). - Cases M-N: dynamic `sigmoid` (elementwise) and `select_copy(0, -1)` (negative index resolved against the live leading dim each call). - Graph-reuse variants: every dynamic op above (`rms_norm` incl. a grow-first smallest→largest order, the `rms(rms(x))` cascade, `linear_q4gsw`, `embedding_q4gsw`, `apply_rotary_emb`, `sigmoid`, `select_copy`) also runs ONE loaded graph across multiple live shapes — proving buffers never move so bind groups stay valid across every resize. **Implementation:** - `test/ops/dynamic_shape/test_dynamic_shape_export.py` exports each toy model through `VulkanPartitioner` with a dynamic dim and writes per-S torch goldens; reuses the existing op-test helpers for quant/sdpa/embedding/rope. - `test/native/test_dynamic_shape.cpp` loads each `.pte`, runs each live S, and compares at the per-op tolerance (rms 1e-3, quant 5e-3, sdpa 2e-3). Reuse tests split each per-op helper into load-once + run-at-shape so a single `Module` serves the whole shape sweep. - Multi-output ops select their output by full shape, never numel. **Constraints:** numerics computed with torch (no hand-rolled reference); toy models stay within the 65535 1D-dispatch cap; SDPA case is skipped gracefully if `sym_size.int`/`copy_` op coverage is incomplete (does not fail the suite). Co-authored-with: Claude Code. ghstack-source-id: 399812841 @exported-using-ghexport Differential Revision: [D109906090](https://our.internmc.facebook.com/intern/diff/D109906090/) --- backends/webgpu/CMakeLists.txt | 7 + .../webgpu/test/native/test_dynamic_shape.cpp | 540 ++++++++++++++++++ .../test_dynamic_shape_export.py | 459 +++++++++++++++ 3 files changed, 1006 insertions(+) create mode 100644 backends/webgpu/test/native/test_dynamic_shape.cpp create mode 100644 backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 7e3cb32de05..072bcf43b8b 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -194,6 +194,13 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) ) target_compile_options(webgpu_op_test_util_test PRIVATE -fexceptions) set_property(TARGET webgpu_op_test_util_test PROPERTY CXX_STANDARD 17) + + # Dynamic-shape integration test: a gtest binary with its own main() that + # brings up the device once (like webgpu_op_test). + add_webgpu_native_test( + webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp + ) + target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest) endif() add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp) endif() diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp new file mode 100644 index 00000000000..f41adbf09ec --- /dev/null +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -0,0 +1,540 @@ +/* + * 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. + */ + +// Native test for dynamic tensor shapes (Option 2). One graph is built at the +// upper-bound seq-len MAXS and run at several live S; the output must match the +// torch golden at each S (allocate-at-max + per-op resize hooks + output-EValue +// resize). Cases: +// A dyn_rms at S=MAXS -> golden (static-equivalent) +// B dyn_rms at S < MAXS (64, 8, 1) -> golden (resize shrinks dispatch) +// C ONE loaded graph reused across S -> all golden (buffers never moved => +// bind groups stayed valid) +// D static_rms (no dynamic dim) -> golden (static path unchanged) +// F dyn_rms_chain (rms(rms(x))) at 3 S -> golden (resize CASCADE, DD-4) +// G rms+residual H rms*x I dyn_linear J sdpa_dyn K emb_dyn L rope_dyn +// M dyn_sigmoid N dyn_select (select_copy(0,-1), dynamic S) +// .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. +// +// Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else +// /tmp/dynamic_shape. + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace executorch::backends::webgpu; +using namespace executorch::extension; +using namespace executorch::runtime; + +namespace { + +constexpr int kHidden = 64; + +// Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). +std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) + +std::vector read_bin(const std::string& path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) { + return {}; + } + const std::streamsize n = f.tellg(); + if (n < 0) { + return {}; + } + f.seekg(0); + std::vector v(static_cast(n) / sizeof(float)); + f.read(reinterpret_cast(v.data()), n); + return v; +} + +float max_err(const std::vector& a, const std::vector& b) { + if (a.size() != b.size() || a.empty()) { + return 1e30f; + } + float m = 0.0f; + for (size_t i = 0; i < a.size(); i++) { + m = std::fmax(m, std::fabs(a[i] - b[i])); + } + return m; +} + +// Run a [1,1,S,kHidden] input through `m` and compare to the golden. Shared by +// every single-output rms-shaped case (A-H, M). +void check_s(Module& m, const std::string& prefix, int s) { + const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(s); + auto input = read_bin(base + ".input.bin"); + ASSERT_FALSE(input.empty()) << "missing input: " << prefix << ".S" << s; + ASSERT_EQ(input.size(), static_cast(s) * kHidden) + << "wrong input size: " << prefix << ".S" << s; + auto t = make_tensor_ptr({1, 1, s, kHidden}, std::move(input)); + auto r = m.forward({EValue(t)}); + ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) + << prefix << " S=" << s + << " forward failed (err=" << (r.ok() ? 0 : (int)r.error()) << ")"; + const auto& out = r.get()[0].toTensor(); + const size_t numel = static_cast(s) * kHidden; + // Output EValue must have been resized to the live shape. + ASSERT_EQ(static_cast(out.numel()), numel) + << prefix << " S=" << s << " output numel mismatch"; + const float* d = out.const_data_ptr(); + std::vector got(d, d + numel); + auto golden = read_bin(base + ".golden.bin"); + const float e = max_err(got, golden); + EXPECT_LT(e, 1e-3f) << prefix << " S=" << s << " max_err=" << e + << " (got.size=" << got.size() + << " golden.size=" << golden.size() << ")"; +} + +// Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the +// register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. +constexpr int kLinK = 64; +constexpr int kLinN = 128; +constexpr int kLinNShmem = 2048; +// Run at [m_rows, kLinK] on an already-loaded module (so it can be +// reused across M without a fresh load), and compare to the golden. +void run_linear(Module& m, int m_rows, const char* prefix, int n) { + const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(m_rows); + auto input = read_bin(base + ".input.bin"); + auto golden = read_bin(base + ".golden.bin"); + ASSERT_FALSE(input.empty()) << "missing " << prefix << ".S" << m_rows; + auto t = make_tensor_ptr({m_rows, kLinK}, std::move(input)); + auto r = m.forward({EValue(t)}); + ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) + << prefix << " M=" << m_rows << " forward failed"; + const auto& out = r.get()[0].toTensor(); + const size_t numel = static_cast(m_rows) * n; + ASSERT_EQ(static_cast(out.numel()), numel) + << prefix << " M=" << m_rows << " output numel mismatch"; + std::vector got( + out.const_data_ptr(), out.const_data_ptr() + numel); + const float e = max_err(got, golden); + // 4-bit quant: looser tol (the kernel mirrors the dequant-matmul reference). + EXPECT_LT(e, 5e-3f) << prefix << " M=" << m_rows << " max_err=" << e; +} + +void check_linear(int m_rows) { + Module m(g_dir + "/dyn_linear.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear.pte"; + run_linear(m, m_rows, "dyn_linear", kLinN); +} + +void check_linear_shmem(int m_rows) { + Module m(g_dir + "/dyn_linear_shmem.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); +} + +// Dynamic SDPA (GQA prefill, input_pos=0): q[1,s,hq,d] k/v[1,s,hkv,d] +// caches[1,cmax,hkv,d]; attn output [1,s,hq,d] selected by shape (3 outputs). +constexpr int kSdHq = 8, kSdHkv = 2, kSdD = 16, kSdCmax = 64; +void check_sdpa(int s) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; + const std::string b = g_dir + "/sdpa_dyn.S" + std::to_string(s) + "."; + auto q = read_bin(b + "q.bin"); + auto k = read_bin(b + "k.bin"); + auto v = read_bin(b + "v.bin"); + auto kc = read_bin(b + "kc.bin"); + auto vc = read_bin(b + "vc.bin"); + auto golden = read_bin(b + "golden.bin"); + ASSERT_FALSE( + q.empty() || k.empty() || v.empty() || kc.empty() || vc.empty() || + golden.empty()) + << "missing sdpa_dyn.S" << s; + auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); + auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); + auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); + auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); + auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + auto r = + m.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tkc), EValue(tvc)}); + ASSERT_TRUE(r.ok()) << "sdpa S=" << s + << " forward failed (err=" << (int)r.error() << ")"; + // Select the attn output by full shape [1,s,hq,d] (never numel). + const float* attn = nullptr; + const size_t numel = static_cast(s) * kSdHq * kSdD; + for (size_t i = 0; i < r.get().size(); i++) { + if (!r.get()[i].isTensor()) { + continue; + } + const auto& t = r.get()[i].toTensor(); + if (t.dim() == 4 && t.size(1) == s && t.size(2) == kSdHq && + t.size(3) == kSdD) { + attn = t.const_data_ptr(); + break; + } + } + ASSERT_NE(attn, nullptr) << "sdpa S=" << s << ": no attn output of shape [1," + << s << "," << kSdHq << "," << kSdD << "]"; + std::vector got(attn, attn + numel); + const float e = max_err(got, golden); + EXPECT_LT(e, 2e-3f) << "sdpa_dyn S=" << s << " max_err=" << e; +} + +// Dynamic embedding: int64 token ids [N] -> [N, kEmbDim] fp32. The int64 host +// input exercises copy_inputs' int64->int32 narrow path under dynamic shapes. +constexpr int kEmbDim = 64; +// Run emb_dyn at N tokens on an already-loaded module (so it can be reused +// across N), and compare to the golden. +void run_embedding(Module& m, int n) { + const std::string b = g_dir + "/emb_dyn.S" + std::to_string(n) + "."; + std::ifstream f(b + "idx.bin", std::ios::binary | std::ios::ate); + ASSERT_TRUE(f.good()) << "missing emb_dyn.S" << n; + const std::streamsize nb = f.tellg(); + ASSERT_GE(nb, 0) << "missing emb_dyn.S" << n; + f.seekg(0); + std::vector idx(static_cast(nb) / sizeof(int64_t)); + f.read(reinterpret_cast(idx.data()), nb); + ASSERT_EQ(idx.size(), static_cast(n)) + << "wrong emb_dyn idx size S" << n; + auto golden = read_bin(b + "golden.bin"); + auto t = make_tensor_ptr({n}, std::move(idx)); // int64 (Long) host input + auto r = m.forward({EValue(t)}); + ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) + << "emb N=" << n + << " forward failed (err=" << (r.ok() ? 0 : (int)r.error()) << ")"; + const auto& out = r.get()[0].toTensor(); + const size_t numel = static_cast(n) * kEmbDim; + ASSERT_EQ(static_cast(out.numel()), numel) + << "emb N=" << n << " output numel mismatch"; + std::vector got( + out.const_data_ptr(), out.const_data_ptr() + numel); + const float e = max_err(got, golden); + EXPECT_LT(e, 5e-3f) << "emb_dyn N=" << n << " max_err=" << e; +} + +void check_embedding(int n) { + Module m(g_dir + "/emb_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn.pte"; + run_embedding(m, n); +} + +// Dynamic RoPE: xq[1,s,nh,hd] xk[1,s,nkv,hd] freqs[s,hd/2] -> xq_out/xk_out +// (2 outputs, selected by head count nh != nkv). +constexpr int kRopeNH = 8, kRopeNKV = 2, kRopeHD = 64; +// Run rope_dyn at seq-len s on an already-loaded module (so it can be reused +// across s), comparing xq_out/xk_out (selected by head count) to the goldens. +void run_rope(Module& m, int s) { + const std::string b = g_dir + "/rope_dyn.S" + std::to_string(s) + "."; + auto xq = read_bin(b + "xq.bin"); + auto xk = read_bin(b + "xk.bin"); + auto fc = read_bin(b + "fc.bin"); + auto fs = read_bin(b + "fs.bin"); + auto gq = read_bin(b + "gq.bin"); + auto gk = read_bin(b + "gk.bin"); + ASSERT_FALSE( + xq.empty() || xk.empty() || fc.empty() || fs.empty() || gq.empty() || + gk.empty()) + << "missing rope_dyn.S" << s; + auto txq = make_tensor_ptr({1, s, kRopeNH, kRopeHD}, std::move(xq)); + auto txk = make_tensor_ptr({1, s, kRopeNKV, kRopeHD}, std::move(xk)); + auto tfc = make_tensor_ptr({s, kRopeHD / 2}, std::move(fc)); + auto tfs = make_tensor_ptr({s, kRopeHD / 2}, std::move(fs)); + auto r = m.forward({EValue(txq), EValue(txk), EValue(tfc), EValue(tfs)}); + ASSERT_TRUE(r.ok()) << "rope S=" << s + << " forward failed (err=" << (int)r.error() << ")"; + // Select xq_out (nh heads) and xk_out (nkv heads) by shape. + const float *oq = nullptr, *okp = nullptr; + for (size_t i = 0; i < r.get().size(); i++) { + if (!r.get()[i].isTensor()) { + continue; + } + const auto& t = r.get()[i].toTensor(); + if (t.dim() == 4 && t.size(1) == s && t.size(3) == kRopeHD) { + if (t.size(2) == kRopeNH) { + oq = t.const_data_ptr(); + } else if (t.size(2) == kRopeNKV) { + okp = t.const_data_ptr(); + } + } + } + ASSERT_TRUE(oq != nullptr && okp != nullptr) + << "rope S=" << s << ": missing xq_out/xk_out by shape"; + std::vector gotq(oq, oq + static_cast(s) * kRopeNH * kRopeHD); + std::vector gotk( + okp, okp + static_cast(s) * kRopeNKV * kRopeHD); + const float e = std::fmax(max_err(gotq, gq), max_err(gotk, gk)); + EXPECT_LT(e, 1e-3f) << "rope_dyn S=" << s << " max_err=" << e; +} + +void check_rope(int s) { + Module m(g_dir + "/rope_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load rope_dyn.pte"; + run_rope(m, s); +} + +// Dynamic select_copy(0,-1): input [2,1,S,kHidden] -> output [1,S,kHidden]. The +// negative index resolves against the (static) leading dim live; the dynamic S +// flows to the output, so the resize hook recomputes its dispatch each S. +constexpr int kSelLead = 2; +// Run dyn_select at seq-len s on an already-loaded module (so it can be reused +// across s), and compare to the golden. +void run_select(Module& m, int s) { + const std::string base = g_dir + "/dyn_select.S" + std::to_string(s); + auto input = read_bin(base + ".input.bin"); + auto golden = read_bin(base + ".golden.bin"); + ASSERT_FALSE(input.empty() || golden.empty()) << "missing dyn_select.S" << s; + auto t = make_tensor_ptr({kSelLead, 1, s, kHidden}, std::move(input)); + auto r = m.forward({EValue(t)}); + ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) + << "select S=" << s + << " forward failed (err=" << (r.ok() ? 0 : (int)r.error()) << ")"; + const auto& out = r.get()[0].toTensor(); + const size_t numel = static_cast(s) * kHidden; + ASSERT_EQ(static_cast(out.numel()), numel) + << "select S=" << s << " output numel mismatch"; + std::vector got( + out.const_data_ptr(), out.const_data_ptr() + numel); + const float e = max_err(got, golden); + EXPECT_LT(e, 1e-3f) << "dyn_select S=" << s << " max_err=" << e; +} + +void check_select(int s) { + Module m(g_dir + "/dyn_select.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_select.pte"; + run_select(m, s); +} + +} // namespace + +// A + B: single dynamic rms_norm at S = MAXS .. 1 (fresh module load each S). +TEST(DynamicShape, RmsNormFreshLoad) { + for (int s : {128, 64, 8, 1}) { + Module m(g_dir + "/dyn_rms.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_rms.pte"; + check_s(m, "dyn_rms", s); + } +} + +// C: ONE loaded graph reused across S (buffers must not move => bind groups +// stay valid). +TEST(DynamicShape, RmsNormReusedGraph) { + Module m(g_dir + "/dyn_rms.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_rms.pte"; + for (int s : {128, 1, 64, 8, 128}) { + check_s(m, "dyn_rms", s); + } +} + +// C2: grow-only reuse — one loaded rms graph run smallest -> largest, so the +// FIRST resize grows the dispatch (every other reuse test starts at MAXS and +// only shrinks; this catches a hook with a shrink-only short-circuit). +TEST(DynamicShape, RmsNormGrowReused) { + Module m(g_dir + "/dyn_rms.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_rms.pte"; + for (int s : {1, 8, 64, 128}) { + check_s(m, "dyn_rms", s); + } +} + +// D: static rms_norm (no dynamic dim) — regression that the static path is +// unchanged. +TEST(DynamicShape, StaticRmsNorm) { + Module m(g_dir + "/static_rms.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load static_rms.pte"; + check_s(m, "static_rms", 8); +} + +// F: 2-op chain rms(rms(x)) — resize cascade. +TEST(DynamicShape, RmsChainCascade) { + for (int s : {128, 16, 1}) { + Module m(g_dir + "/dyn_rms_chain.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_rms_chain.pte"; + check_s(m, "dyn_rms_chain", s); + } +} + +// F2: cascade graph REUSED across S — one loaded rms(rms(x)) graph run across S +// (op1's output-resize feeds op2's input-resize on a persistent graph; the +// single-op reuse tests don't cover an inter-op resize cascade under reuse). +TEST(DynamicShape, RmsChainCascadeReusedGraph) { + Module m(g_dir + "/dyn_rms_chain.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_rms_chain.pte"; + for (int s : {128, 16, 1, 128}) { + check_s(m, "dyn_rms_chain", s); + } +} + +// G: rms(x)+x residual — cross-op (rms -> add) cascade. +TEST(DynamicShape, RmsResidualCascade) { + for (int s : {128, 32, 1}) { + Module m(g_dir + "/dyn_residual.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_residual.pte"; + check_s(m, "dyn_residual", s); + } +} + +// H: rms(x)*x — exercises the mul op resize. +TEST(DynamicShape, RmsMul) { + for (int s : {128, 32, 1}) { + Module m(g_dir + "/dyn_rmsmul.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_rmsmul.pte"; + check_s(m, "dyn_rmsmul", s); + } +} + +// I: dynamic 4-bit quantized linear (prefill GEMM) at several M. +TEST(DynamicShape, QuantizedLinear) { + for (int m_rows : {128, 32, 1}) { + check_linear(m_rows); + } +} + +// I2: dynamic linear reusing ONE loaded graph across M (buffers must not move +// => bind groups stay valid; the resize hook recomputes dispatch each M). +TEST(DynamicShape, QuantizedLinearReusedGraph) { + Module m(g_dir + "/dyn_linear.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear.pte"; + for (int m_rows : {128, 32, 1, 128}) { + run_linear(m, m_rows, "dyn_linear", kLinN); + } +} + +// I3: dynamic linear at N=2048 -> the shmem-GEMM route (K>=4096||N>=2048); the +// resize hook recomputes the shmem tile count for the live M on the fixed shmem +// pipeline (M=1 exercises a partial row-tile). +TEST(DynamicShape, QuantizedLinearShmem) { + for (int m_rows : {128, 32, 1}) { + check_linear_shmem(m_rows); + } +} + +// I4: shmem-routed linear reusing ONE loaded graph across M. +TEST(DynamicShape, QuantizedLinearShmemReusedGraph) { + Module m(g_dir + "/dyn_linear_shmem.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; + for (int m_rows : {128, 32, 1, 128}) { + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + } +} + +// J: dynamic SDPA (GQA prefill) at several seq-len S. The whole case skips +// while op coverage is pending (the dynamic-S build throws err 48 until +// registered). +TEST(DynamicShape, Sdpa) { + { + Module probe(g_dir + "/sdpa_dyn.pte"); + if (probe.load_forward() == Error::DelegateInvalidCompatibility) { + GTEST_SKIP() << "sdpa_dyn pending op coverage (err " + << (int)Error::DelegateInvalidCompatibility << ")"; + } + } + for (int s : {64, 16, 1}) { + check_sdpa(s); + } +} + +// K: dynamic embedding (int64 token ids) at several token counts. +TEST(DynamicShape, Embedding) { + for (int n : {16, 8, 1}) { + check_embedding(n); + } +} + +// K2: dynamic embedding reusing ONE loaded graph across N (buffers must not +// move +// => the multi-buffer bind group stays valid; resize recomputes blocks each N). +TEST(DynamicShape, EmbeddingReusedGraph) { + Module m(g_dir + "/emb_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn.pte"; + for (int n : {16, 8, 1, 16}) { + run_embedding(m, n); + } +} + +// L: dynamic RoPE (two outputs) at several seq-len S. +TEST(DynamicShape, Rope) { + for (int s : {16, 8, 1}) { + check_rope(s); + } +} + +// L2: dynamic RoPE reusing ONE loaded graph across S (buffers must not move => +// bind groups stay valid; the resize hook recomputes both outputs each S). +TEST(DynamicShape, RopeReusedGraph) { + Module m(g_dir + "/rope_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load rope_dyn.pte"; + for (int s : {16, 8, 1, 16}) { + run_rope(m, s); + } +} + +// M: dynamic sigmoid (elementwise) at several S. +TEST(DynamicShape, Sigmoid) { + for (int s : {128, 32, 1}) { + Module m(g_dir + "/dyn_sigmoid.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_sigmoid.pte"; + check_s(m, "dyn_sigmoid", s); + } +} + +// M2: dynamic sigmoid reusing ONE loaded graph across S. +TEST(DynamicShape, SigmoidReusedGraph) { + Module m(g_dir + "/dyn_sigmoid.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_sigmoid.pte"; + for (int s : {128, 32, 1, 128}) { + check_s(m, "dyn_sigmoid", s); + } +} + +// N: dynamic select_copy(0,-1) at several S. +TEST(DynamicShape, Select) { + for (int s : {128, 32, 1}) { + check_select(s); + } +} + +// N2: dynamic select_copy reusing ONE loaded graph across S (the resize hook +// re-resolves the negative index against the LIVE leading dim each call). +TEST(DynamicShape, SelectReusedGraph) { + Module m(g_dir + "/dyn_select.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_select.pte"; + for (int s : {128, 32, 1, 128}) { + run_select(m, s); + } +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + + // Artifacts dir: env wins, else first positional arg, else default (gtest + // flags were already stripped by InitGoogleTest above). + g_dir = "/tmp/dynamic_shape"; + if (argc > 1) { + g_dir = argv[1]; + } + if (const char* env = std::getenv("WEBGPU_DYNAMIC_SHAPE_DIR")) { + g_dir = env; + } + + WebGPUContext ctx; + try { + ctx = create_webgpu_context(); + } catch (const std::exception& e) { + std::printf("SKIP: no WebGPU device (%s)\n", e.what()); + return 0; + } + set_default_webgpu_context(&ctx); + + const int rc = RUN_ALL_TESTS(); + set_default_webgpu_context(nullptr); + destroy_webgpu_context(ctx); + return rc; +} diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py new file mode 100644 index 00000000000..969c5e06343 --- /dev/null +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -0,0 +1,459 @@ +# 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. + +"""Dynamic tensor-shape (Option 2) export tests via VulkanPartitioner. + +Exports ONE graph built at the upper-bound seq-len MAXS that the native runtime +(`test/native/test_dynamic_shape.cpp`) then runs at several live S, asserting the +output matches the torch golden and that the static path is unchanged. Numerics +are checked in the native test; this verifies the dynamic export side + writes +goldens. +""" + +import os +import unittest + +import torch +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +MAXS = 128 # upper bound for the dynamic seq-len dim (within the 1D dispatch cap) +HIDDEN = 64 + + +def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + x_f32 = x.to(torch.float32) + var = x_f32.pow(2).mean(dim=-1, keepdim=True) + return (x_f32 * torch.rsqrt(var + eps)) * weight + + +class RmsNormModule(torch.nn.Module): + def __init__(self, hidden: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = torch.nn.Parameter( + torch.linspace(0.5, 1.5, hidden, dtype=torch.float32) + ) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _rms(x, self.weight, self.eps) + + +class RmsChainModule(torch.nn.Module): + """rms(rms(x)) — two ops; exercises the resize-cascade (DD-4).""" + + def __init__(self, hidden: int, eps: float = 1e-6) -> None: + super().__init__() + self.w1 = torch.nn.Parameter( + torch.linspace(0.5, 1.5, hidden, dtype=torch.float32) + ) + self.w2 = torch.nn.Parameter( + torch.linspace(1.5, 0.5, hidden, dtype=torch.float32) + ) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _rms(_rms(x, self.w1, self.eps), self.w2, self.eps) + + +class RmsResidualModule(torch.nn.Module): + """rms(x) + x — rms op feeding an add op; proves the cross-op resize cascade.""" + + def __init__(self, hidden: int, eps: float = 1e-6) -> None: + super().__init__() + self.w = torch.nn.Parameter( + torch.linspace(0.5, 1.5, hidden, dtype=torch.float32) + ) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _rms(x, self.w, self.eps) + x + + +class RmsMulModule(torch.nn.Module): + """rms(x) * x — exercises the mul op (two same-shape dynamic operands).""" + + def __init__(self, hidden: int, eps: float = 1e-6) -> None: + super().__init__() + self.w = torch.nn.Parameter( + torch.linspace(0.5, 1.5, hidden, dtype=torch.float32) + ) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _rms(x, self.w, self.eps) * x + + +class SigmoidModule(torch.nn.Module): + """sigmoid(x) — elementwise; resize hook recomputes dispatch from live numel.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.sigmoid(x) + + +class SelectModule(torch.nn.Module): + """x.select(0, -1) — negative index resolved live + dynamic output dispatch.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.select(0, -1) + + +def _ramp(shape) -> torch.Tensor: + n = 1 + for d in shape: + n *= d + return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) + + +def _export(model, example_inputs, dynamic_shapes, path: str) -> None: + ep = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + assert found, f"Expected VulkanBackend delegate in {path}" + with open(path, "wb") as f: + f.write(et.buffer) + print(f"Exported {path}") + + +def _write_goldens(model, prefix: str, out_dir: str, s_values) -> None: + for s in s_values: + x = _ramp((1, 1, s, HIDDEN)) + with torch.no_grad(): + g = model(x) + x.detach().numpy().astype(" None: + """Write the dynamic + static .pte's and per-S goldens for the native test.""" + os.makedirs(out_dir, exist_ok=True) + s_dim = torch.export.Dim("s", min=1, max=MAXS) + + # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). + rms = RmsNormModule(HIDDEN) + _export( + rms, + (_ramp((1, 1, MAXS, HIDDEN)),), + {"x": {2: s_dim}}, + os.path.join(out_dir, "dyn_rms.pte"), + ) + _write_goldens(rms, "dyn_rms", out_dir, [MAXS, 64, 8, 1]) + + # 2) Two-op chain (cascade): rms(rms(x)). + chain = RmsChainModule(HIDDEN) + _export( + chain, + (_ramp((1, 1, MAXS, HIDDEN)),), + {"x": {2: s_dim}}, + os.path.join(out_dir, "dyn_rms_chain.pte"), + ) + _write_goldens(chain, "dyn_rms_chain", out_dir, [MAXS, 16, 1]) + + # 2b) rms(x)+x residual — cross-op (rms->add) cascade. + resid = RmsResidualModule(HIDDEN) + _export( + resid, + (_ramp((1, 1, MAXS, HIDDEN)),), + {"x": {2: s_dim}}, + os.path.join(out_dir, "dyn_residual.pte"), + ) + _write_goldens(resid, "dyn_residual", out_dir, [MAXS, 32, 1]) + + # 2c) rms(x)*x — exercises the mul op resize. + rmsmul = RmsMulModule(HIDDEN) + _export( + rmsmul, + (_ramp((1, 1, MAXS, HIDDEN)),), + {"x": {2: s_dim}}, + os.path.join(out_dir, "dyn_rmsmul.pte"), + ) + _write_goldens(rmsmul, "dyn_rmsmul", out_dir, [MAXS, 32, 1]) + + # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM + # (register-tiled N=128) + a shmem-GEMM-routed variant (N=2048). + _export_dynamic_linear(out_dir) + _export_dynamic_linear(out_dir, n=LIN_SHMEM_N, prefix="dyn_linear_shmem") + + # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). + _export_dynamic_sdpa(out_dir) + + # 2f) 4-bit embedding with a DYNAMIC token count (int64 indices). + _export_dynamic_embedding(out_dir) + + # 2g) Interleaved RoPE with a DYNAMIC seq-len S (two outputs xq/xk). + _export_dynamic_rope(out_dir) + + # 2h) Elementwise sigmoid with a DYNAMIC seq-len S. + sig = SigmoidModule() + _export( + sig, + (_ramp((1, 1, MAXS, HIDDEN)),), + {"x": {2: s_dim}}, + os.path.join(out_dir, "dyn_sigmoid.pte"), + ) + _write_goldens(sig, "dyn_sigmoid", out_dir, [MAXS, 32, 1]) + + # 2i) select_copy(0, -1) over a DYNAMIC seq-len S (negative live index). + _export_dynamic_select(out_dir) + + # 3) Static rms_norm (no dynamic dim) — regression: must stay byte-identical. + static = RmsNormModule(HIDDEN) + _export( + static, + (_ramp((1, 1, 8, HIDDEN)),), + None, + os.path.join(out_dir, "static_rms.pte"), + ) + _write_goldens(static, "static_rms", out_dir, [8]) + + +# Quantized linear: K x N weight, dynamic rows M; input [M, K], output [M, N]. +LIN_K = 64 +LIN_N = 128 +LIN_SHMEM_N = 2048 # N>=2048 routes linear_q4gsw to the shmem-GEMM path +LIN_GROUP = 32 +LIN_MAXM = 128 + + +def _export_dynamic_linear( + out_dir: str, n: int = LIN_N, prefix: str = "dyn_linear" +) -> None: + from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import ( + _fp64_golden, + _make_quantized_model, + ) + + model = _make_quantized_model(LIN_K, n, LIN_GROUP) + x = _ramp((LIN_MAXM, LIN_K)) + m_dim = torch.export.Dim("m", min=1, max=LIN_MAXM) + ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + assert any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ), "linear_q4gsw not delegated" + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + for m in [LIN_MAXM, 32, 1]: + xm = _ramp((m, LIN_K)) + g = _fp64_golden(model, xm).astype(" None: + from executorch.backends.webgpu.test.ops.sdpa.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + def cfg(s: int) -> "SdpaConfig": + return SdpaConfig("dyn", SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + + model = SdpaModule(0) + q, k, v, kc, vc = _det_inputs(cfg(SD_MAXS)) + s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) + ds = ({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None) + ep = torch.export.export(model, (q, k, v, kc, vc), dynamic_shapes=ds) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + assert any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ), "sdpa not delegated" + with open(os.path.join(out_dir, "sdpa_dyn.pte"), "wb") as f: + f.write(et.buffer) + print("Exported sdpa_dyn.pte") + for s in [SD_MAXS, 16, 1]: + c = cfg(s) + q, k, v, kc, vc = _det_inputs(c) + g = _golden(c, q, k, v, kc, vc) + for name, t in [ + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", g), + ]: + t.detach().cpu().numpy().astype(" [N, EMBED] fp32. +EMB_VOCAB = 64 +EMB_DIM = 64 +EMB_GROUP = 32 +EMB_MAXN = 16 + + +def _export_dynamic_embedding(out_dir: str) -> None: + from executorch.backends.webgpu.test.ops.embedding_q4gsw.test_embedding_q4gsw import ( + _make_quantized_model, + _quant_params, + Shape, + ) + + shape = Shape("dyn", EMB_VOCAB, EMB_DIM, EMB_GROUP, list(range(EMB_MAXN))) + qm = _make_quantized_model(shape) + idx_max = torch.arange(EMB_MAXN, dtype=torch.long) + n_dim = torch.export.Dim("n", min=1, max=EMB_MAXN) + ep = torch.export.export(qm, (idx_max,), dynamic_shapes=({0: n_dim},)) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + assert any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ), "embedding_q4gsw not delegated" + with open(os.path.join(out_dir, "emb_dyn.pte"), "wb") as f: + f.write(et.buffer) + print("Exported emb_dyn.pte") + weight, scales, group_size = _quant_params(qm) + for n in [EMB_MAXN, 8, 1]: + idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB + g = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, False + ) + idx.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.rope.test_rope import ( + _golden, + _inputs, + Shape, + ) + from executorch.examples.models.llama.rope import RotaryEmbedding + + xq, xk, fc, fs = _inputs(Shape("dyn", 1, ROPE_MAXS, ROPE_NH, ROPE_NKV, ROPE_HD)) + s_dim = torch.export.Dim("s", min=1, max=ROPE_MAXS) + ds = ({1: s_dim}, {1: s_dim}, {0: s_dim}, {0: s_dim}) + ep = torch.export.export( + RotaryEmbedding().eval(), (xq, xk, fc, fs), dynamic_shapes=ds + ) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + assert any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ), "apply_rotary_emb not delegated" + with open(os.path.join(out_dir, "rope_dyn.pte"), "wb") as f: + f.write(et.buffer) + print("Exported rope_dyn.pte") + for s in [ROPE_MAXS, 8, 1]: + xq, xk, fc, fs = _inputs(Shape("dyn", 1, s, ROPE_NH, ROPE_NKV, ROPE_HD)) + gq, gk = _golden(xq, xk, fc, fs) + for name, t in [ + ("xq", xq), + ("xk", xk), + ("fc", fc), + ("fs", fs), + ("gq", gq), + ("gk", gk), + ]: + t.detach().cpu().numpy().astype(" [1, S, HIDDEN]. +SEL_LEAD = 2 + + +def _export_dynamic_select(out_dir: str) -> None: + model = SelectModule() + s_dim = torch.export.Dim("s", min=1, max=MAXS) + ep = torch.export.export( + model, + (_ramp((SEL_LEAD, 1, MAXS, HIDDEN)),), + dynamic_shapes=({2: s_dim},), + ) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + assert any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ), "select_copy not delegated" + with open(os.path.join(out_dir, "dyn_select.pte"), "wb") as f: + f.write(et.buffer) + print("Exported dyn_select.pte") + for s in [MAXS, 32, 1]: + x = _ramp((SEL_LEAD, 1, s, HIDDEN)) + with torch.no_grad(): + g = model(x) + x.detach().numpy().astype(" None: + import tempfile + + with tempfile.TemporaryDirectory() as d: + export_dynamic_shape_cases(d) + self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.pte"))) + self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.S1.golden.bin"))) + + +if __name__ == "__main__": + unittest.main() From f8172f5891bcb13113b53b83b7314f81532e24f4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:57 -0700 Subject: [PATCH 10/14] =?UTF-8?q?[ExecuTorch][WebGPU]=202D=20compute=20dis?= =?UTF-8?q?patch=20=E2=80=94=20lift=20the=2065535=20per-dim=20cap=20(prefi?= =?UTF-8?q?ll=20path)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20583 **Lift the 65535 workgroup-per-dim dispatch cap so single-shot SDPA prefill runs at any sequence length.** **Problem**: The WebGPU backend is 1D-dispatch-only and throws when a kernel's workgroup count exceeds the device per-dim limit (`maxComputeWorkgroupsPerDimension`, spec floor 65535). SDPA prefill QK exceeds it around S~362 (softmax/AV at S=2048), blocking single-shot / long-context prefill. **Solution**: Fold a >limit 1D workgroup count into 2D; the shader reconstructs the linear index from `@builtin(num_workgroups)`. - **Before**: `compute_1d_workgroup_count` throws if `count > limit`; dispatch `(count, 1, 1)`. - **After**: `compute_2d_workgroup_count` returns `{count, 1}` (fast path) or a near-square `{x, y}` (`x = ceil(sqrt(count))` clamped to `limit`, `y = div_up(count, x)`); dispatch `(x, y, 1)`. A flat `{limit, div_up(count, limit)}` split would idle up to ~half the launched workgroups when `count` just exceeds `limit`; the near-square split holds the waste to `O(sqrt(count))` (e.g. 65536 -> `{256, 256}`, 0 inactive). **Implementation**: - `WgCount` + pure `fold_workgroup_count_2d` + `compute_2d_workgroup_count` in `WebGPUUtils.h` (device-free, unit-testable; `queried_max_workgroups` factored out of the 1D path) - `WebGPUDispatch.workgroup_count_y` (default 1, declared last so existing aggregate inits are unchanged); both `dispatchWorkgroups` calls + the profiling record pass `(x, y, 1)` - Per-kernel in-shader reconstruction: thread-form `idx = gid.x + gid.y*(num_workgroups.x*wg_size)` (QK/AV/add); row-form `row_idx = wid.x + wid.y*num_workgroups.x` (softmax — keeps a `valid` predicate, not an early return, so `workgroupBarrier()`s stay uniform) - `Sdpa.cpp`: QK/softmax/AV counts via the 2D helper; the dynamic-`input_pos` resize hook recomputes both x and y for QK - Reference: ET-Vulkan dispatches over natural N-D extents (never folds a flat count nor guards the per-dim limit) and MLX `get_2d_grid_dims` packs whole tensor dims; for our flattened scalar count the near-square split is the correct no-shape-info analog (a pack-to-limit split would reproduce the idle-half waste) **Constraints**: - `y=1` fast path keeps every non-folded dispatch byte-identical to the prior 1D path - Scope = prefill path only; `rms_norm`/`embedding`/`lm_head`/`update_cache` are row/token-indexed and never hit the cap, so they keep the 1D path - Throws if a 3rd dispatch dimension would be needed — unreachable for real prefill (the `uint32` element guard fires first at S~11585) Co-authored-with: Claude Code. ghstack-source-id: 399812920 @exported-using-ghexport Differential Revision: [D109517684](https://our.internmc.facebook.com/intern/diff/D109517684/) --- backends/webgpu/runtime/WebGPUGraph.cpp | 9 ++- backends/webgpu/runtime/WebGPUGraph.h | 1 + backends/webgpu/runtime/WebGPUUtils.h | 64 +++++++++++++++++-- backends/webgpu/runtime/ops/add/BinaryOp.cpp | 60 +++++++++-------- .../webgpu/runtime/ops/add/binary_add.wgsl | 6 +- .../webgpu/runtime/ops/add/binary_add_wgsl.h | 8 ++- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 45 ++++++++----- .../ops/sdpa/sdpa_compute_attn_weights.wgsl | 12 ++-- .../ops/sdpa/sdpa_compute_attn_weights_wgsl.h | 14 ++-- .../runtime/ops/sdpa/sdpa_compute_out.wgsl | 12 ++-- .../runtime/ops/sdpa/sdpa_compute_out_wgsl.h | 14 ++-- .../webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl | 7 +- .../runtime/ops/sdpa/sdpa_softmax_wgsl.h | 9 ++- 13 files changed, 176 insertions(+), 85 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index bc3883399a1..eb10fb44e28 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -814,7 +814,7 @@ void WebGPUGraph::execute() { wgpuComputePassEncoderSetBindGroup( pass, 0, dispatch.bind_group, 0, nullptr); wgpuComputePassEncoderDispatchWorkgroups( - pass, dispatch.workgroup_count_x, 1, 1); + pass, dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1); wgpuComputePassEncoderEnd(pass); wgpuComputePassEncoderRelease(pass); #ifdef WGPU_BACKEND_ENABLE_PROFILING @@ -822,7 +822,7 @@ void WebGPUGraph::execute() { qp->record( static_cast(i), dispatch.kernel_name, - {dispatch.workgroup_count_x, 1, 1}, + {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, {1, 1, 1}); } #endif // WGPU_BACKEND_ENABLE_PROFILING @@ -894,7 +894,10 @@ void WebGPUGraph::execute() { wgpuComputePassEncoderSetBindGroup( pass, 0, dispatches_[i].bind_group, 0, nullptr); wgpuComputePassEncoderDispatchWorkgroups( - pass, dispatches_[i].workgroup_count_x, 1, 1); + pass, + dispatches_[i].workgroup_count_x, + dispatches_[i].workgroup_count_y, + 1); wgpuComputePassEncoderEnd(pass); wgpuComputePassEncoderRelease(pass); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 5474ee4667a..0ebcd8071f9 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -48,6 +48,7 @@ struct WebGPUDispatch { WGPUBindGroup bind_group = nullptr; uint32_t workgroup_count_x = 1; std::string kernel_name; // bench label + uint32_t workgroup_count_y = 1; // 2D fold (>65535); 1 = unchanged 1D path // DMA copy command; default Compute keeps existing positional inits valid. enum class Kind { Compute, Copy }; Kind kind = Kind::Compute; diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index afa90d54aec..4326fde8a19 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -47,6 +48,49 @@ inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) { return desired; } +struct WgCount { + uint32_t x; + uint32_t y; +}; + +// Device's max workgroups per dispatch dimension; the WebGPU spec-default floor +// (65535) if the query fails — never under-reports a real device's capacity. +inline uint32_t queried_max_workgroups(WGPUDevice device) { + WGPULimits limits = {}; + return wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupsPerDimension > 0 + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; +} + +// Pure 2D fold of a 1D workgroup count (device-free, unit-testable): {count,1} +// when count <= max, else a near-square {x, y} with x ~ ceil(sqrt(count)) so +// the launched grid stays close to count. A flat {max, div_up(count, max)} +// split would leave up to ~half the workgroups inactive when count just exceeds +// max, and inactive workgroups still cost launch/scheduling; the near-square +// 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. +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}; +} + // 1D dispatch count (mirrors Vulkan div_up); throws if > device limit. inline uint32_t compute_1d_workgroup_count( WGPUDevice device, @@ -54,13 +98,7 @@ inline uint32_t compute_1d_workgroup_count( uint32_t workgroup_size, const char* op_name) { uint32_t count = div_up(num_threads, workgroup_size); - WGPULimits limits = {}; - uint32_t max_count = - wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && - limits.maxComputeWorkgroupsPerDimension > 0 - ? limits.maxComputeWorkgroupsPerDimension - : 65535u; // WebGPU spec-default floor - if (count > max_count) { + if (count > queried_max_workgroups(device)) { throw std::runtime_error( std::string("WebGPU ") + op_name + ": workgroup count exceeds the 1D dispatch limit"); @@ -68,6 +106,18 @@ inline uint32_t compute_1d_workgroup_count( return count; } +// 2D dispatch count: fold the 1D count across x/y when it exceeds the per-dim +// limit (lifts the cap, e.g. for SDPA prefill). Same fast path as compute_1d. +inline WgCount compute_2d_workgroup_count( + WGPUDevice device, + uint32_t num_threads, + uint32_t workgroup_size, + const char* op_name) { + uint32_t count = div_up(num_threads, workgroup_size); + return fold_workgroup_count_2d( + count, 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) { diff --git a/backends/webgpu/runtime/ops/add/BinaryOp.cpp b/backends/webgpu/runtime/ops/add/BinaryOp.cpp index 8c56ad6c15d..ca49e21f046 100644 --- a/backends/webgpu/runtime/ops/add/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/add/BinaryOp.cpp @@ -53,8 +53,8 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { uint32_t wg_size = utils::clamp_workgroup_size(device, kBinaryAddWorkgroupSizeX); - uint32_t workgroup_count = - utils::compute_1d_workgroup_count(device, num_elements, wg_size, "add"); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_elements, wg_size, "add"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -158,40 +158,38 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); const size_t dispatch_idx = graph.num_dispatches() - 1; // Dynamic shapes: recompute numel/dispatch; out follows the larger operand. WGPUBuffer params_buf = uniform_buffer; - auto add_resize = [in1_id, - in2_id, - out_id, - alpha, - wg_size, - dispatch_idx, - params_buf](WebGPUGraph& g) { - const auto& d1 = g.cur_dims(in1_id); - const auto& d2 = g.cur_dims(in2_id); - const uint64_t n1 = utils::numel_of(d1); - const uint64_t n2 = utils::numel_of(d2); - const uint64_t numel = n2 > n1 ? n2 : n1; - const uint64_t n_min = n2 > n1 ? n1 : n2; - // The flat add follows the larger operand and broadcasts the smaller; valid - // only when the smaller tiles evenly into it (rejects e.g. [4,1] vs [1,3], - // whose true [4,3] result this flat kernel cannot produce). - if (n_min == 0u || numel % n_min != 0u) { - throw std::runtime_error( - "add(resize): operands are not broadcast-compatible by numel"); - } - g.set_cur_dims(out_id, n2 > n1 ? d2 : d1); - AddParams p = {}; - p.num_elements = static_cast(numel); - p.alpha = alpha; - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( + auto add_resize = + [in1_id, in2_id, out_id, alpha, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d1 = g.cur_dims(in1_id); + const auto& d2 = g.cur_dims(in2_id); + const uint64_t n1 = utils::numel_of(d1); + const uint64_t n2 = utils::numel_of(d2); + const uint64_t numel = n2 > n1 ? n2 : n1; + const uint64_t n_min = n2 > n1 ? n1 : n2; + // The flat add follows the larger operand and broadcasts the smaller; + // valid only when the smaller tiles evenly into it (rejects e.g. [4,1] + // vs [1,3], whose true [4,3] result this flat kernel cannot produce). + if (n_min == 0u || numel % n_min != 0u) { + throw std::runtime_error( + "add(resize): operands are not broadcast-compatible by numel"); + } + g.set_cur_dims(out_id, n2 > n1 ? d2 : d1); + AddParams p = {}; + p.num_elements = static_cast(numel); + p.alpha = alpha; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( g.device(), static_cast(numel), wg_size, "add(resize)"); - }; + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; graph.add_tensor_resize_hook(in1_id, add_resize); graph.add_tensor_resize_hook(in2_id, add_resize); diff --git a/backends/webgpu/runtime/ops/add/binary_add.wgsl b/backends/webgpu/runtime/ops/add/binary_add.wgsl index ac88f184c6b..c1cb9d5ffd7 100644 --- a/backends/webgpu/runtime/ops/add/binary_add.wgsl +++ b/backends/webgpu/runtime/ops/add/binary_add.wgsl @@ -11,8 +11,10 @@ struct Params { override wg_size: u32 = 256; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +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/add/binary_add_wgsl.h b/backends/webgpu/runtime/ops/add/binary_add_wgsl.h index 1f2614d3467..829407bb383 100644 --- a/backends/webgpu/runtime/ops/add/binary_add_wgsl.h +++ b/backends/webgpu/runtime/ops/add/binary_add_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_add.wgsl - DO NOT EDIT. -// wgsl-sha256: c1ceec80c8d4d3d56986ad91ce0d7f9a57cd8467b8c3aa07a28da70e51d141d9 +// wgsl-sha256: e66bd67465c2a0296e09668df54f87605a4c91015a615f3734cdd0f140a74477 inline constexpr const char* kBinaryAddWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -28,8 +28,10 @@ struct Params { override wg_size: u32 = 256; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +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/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index c55f5cfa22f..17918863a6e 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -143,6 +143,7 @@ void build_dispatch( WGPUBuffer uniform_buffer, uint64_t uniform_size, uint32_t workgroup_count_x, + uint32_t workgroup_count_y, uint32_t wg_size, bool retain_uniform = false, const char* kernel_name = "") { @@ -216,7 +217,12 @@ void build_dispatch( bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count_x, kernel_name}); + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count_x, + kernel_name, + workgroup_count_y}); wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); @@ -257,6 +263,7 @@ static WGPUBuffer record_update_cache_dispatch( ubuf, sizeof(uc), wgc, + 1, uc_wg, retain_uniform, "update_cache"); @@ -478,7 +485,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } const int64_t qk_tiles = Hq * utils::div_up(S, kSdpaTileM) * utils::div_up(context_len, kSdpaTileN); - const uint32_t wgc = utils::compute_1d_workgroup_count( + const utils::WgCount wgc = utils::compute_2d_workgroup_count( device, static_cast(qk_tiles), qk_wg, "QK"); AttnWeightsParams p = make_attn_weights_params( S, Hq, Hkv, D, context_len, input_pos, g, scale); @@ -494,7 +501,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { 3, ubuf, sizeof(p), - wgc, + wgc.x, + wgc.y, qk_wg, true, "sdpa_compute_attn_weights"); @@ -505,7 +513,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { // Dispatch 4: softmax, one workgroup per (h,s) row of width context_len. { // One workgroup per (h,s) row; wg_size 1 keeps the device dispatch check. - const uint32_t wgc = utils::compute_1d_workgroup_count( + const utils::WgCount wgc = utils::compute_2d_workgroup_count( device, static_cast(Hq * S), 1, "softmax"); SoftmaxParams p = make_softmax_params(Hq, S, context_len); WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); @@ -518,7 +526,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { 2, ubuf, sizeof(p), - wgc, + wgc.x, + wgc.y, 0, true, "sdpa_softmax"); @@ -530,7 +539,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { { const int64_t av_tiles = Hq * utils::div_up(S, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const uint32_t wgc = utils::compute_1d_workgroup_count( + const utils::WgCount wgc = utils::compute_2d_workgroup_count( device, static_cast(av_tiles), av_wg, "AV"); ComputeOutParams p = make_compute_out_params(S, Hq, Hkv, D, context_len, g); WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); @@ -545,7 +554,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { 3, ubuf, sizeof(p), - wgc, + wgc.x, + wgc.y, av_wg, true, "sdpa_compute_out"); @@ -632,25 +642,28 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { wgpuQueueWriteBuffer(gr.queue(), qk_buf, 0, &qp, sizeof(qp)); const int64_t qk_tiles = Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(ctx, kSdpaTileN); - gr.dispatch_at(qk_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + const utils::WgCount qk_wgc = utils::compute_2d_workgroup_count( + gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + gr.dispatch_at(qk_idx).workgroup_count_x = qk_wgc.x; + gr.dispatch_at(qk_idx).workgroup_count_y = qk_wgc.y; // softmax: one workgroup per (h,s) row. SoftmaxParams sp = make_softmax_params(Hq, s, ctx); wgpuQueueWriteBuffer(gr.queue(), softmax_buf, 0, &sp, sizeof(sp)); - gr.dispatch_at(softmax_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - gr.device(), static_cast(Hq * s), 1, "softmax(resize)"); + const utils::WgCount sm_wgc = utils::compute_2d_workgroup_count( + gr.device(), static_cast(Hq * s), 1, "softmax(resize)"); + gr.dispatch_at(softmax_idx).workgroup_count_x = sm_wgc.x; + gr.dispatch_at(softmax_idx).workgroup_count_y = sm_wgc.y; // AV: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(D/TN). ComputeOutParams op = make_compute_out_params(s, Hq, Hkv, D, ctx, g); wgpuQueueWriteBuffer(gr.queue(), av_buf, 0, &op, sizeof(op)); const int64_t av_tiles = Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - gr.dispatch_at(av_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + const utils::WgCount av_wgc = utils::compute_2d_workgroup_count( + gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + gr.dispatch_at(av_idx).workgroup_count_x = av_wgc.x; + gr.dispatch_at(av_idx).workgroup_count_y = av_wgc.y; // Output attn has the same shape as q: [.., S, Hq, D]. gr.set_cur_dims(out_id, gr.cur_dims(q_id)); diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl index fd15767603d..014f0039048 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl @@ -53,17 +53,21 @@ fn store_qk(s: u32, c: u32, h: u32, raw: f32) { } @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { let nrt = (params.S + TM - 1u) / TM; let nct = (params.context_len + TN - 1u) / TN; let tiles = nrt * nct; let total = tiles * params.Hq; - if (gid.x >= total) { + // 2D dispatch fold: recover the linear tile index across x/y. + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= total) { return; } - let h = gid.x / tiles; - let rem = gid.x % tiles; + let h = idx / tiles; + let rem = idx % tiles; let row_tile = rem / nct; let col_tile = rem % nct; let kvh = h / params.g; diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_wgsl.h index ae250959e0e..d19de1ec214 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_compute_attn_weights.wgsl - DO NOT EDIT. -// wgsl-sha256: 4eef09b234fd926cdc0daf18d03e39cf4fd57dfa4bc67724b4878b7dc68d1254 +// wgsl-sha256: ac1616d310a67aeb915ee7b4f650c981cc423c5251de2e7acf8775bacdfd8e56 inline constexpr const char* kSdpaComputeAttnWeightsWGSL = R"( @group(0) @binding(0) var t_attn_weights: array; @group(0) @binding(1) var t_q: array>; @@ -70,17 +70,21 @@ fn store_qk(s: u32, c: u32, h: u32, raw: f32) { } @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { let nrt = (params.S + TM - 1u) / TM; let nct = (params.context_len + TN - 1u) / TN; let tiles = nrt * nct; let total = tiles * params.Hq; - if (gid.x >= total) { + // 2D dispatch fold: recover the linear tile index across x/y. + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= total) { return; } - let h = gid.x / tiles; - let rem = gid.x % tiles; + let h = idx / tiles; + let rem = idx % tiles; let row_tile = rem / nct; let col_tile = rem % nct; let kvh = h / params.g; diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl index 56242b0ddde..713345c0afa 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl @@ -64,17 +64,21 @@ fn store_out_vec4(s: u32, d0: u32, h: u32, val: vec4) { } @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { let nrt = (params.S + TM - 1u) / TM; let nct = (params.D + TN - 1u) / TN; let tiles = nrt * nct; let total = tiles * params.Hq; - if (gid.x >= total) { + // 2D dispatch fold: recover the linear tile index across x/y. + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= total) { return; } - let h = gid.x / tiles; - let rem = gid.x % tiles; + let h = idx / tiles; + let rem = idx % tiles; let row_tile = rem / nct; let col_tile = rem % nct; let kvh = h / params.g; diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_wgsl.h index 6bec079ac2b..80af51004f3 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_compute_out.wgsl - DO NOT EDIT. -// wgsl-sha256: 2ffa0eb520b1054e43a10fd13e6b287bd35777f1cfc29bd39e9d668772528191 +// wgsl-sha256: bbd7a511032fb901f9ee245357a7d681f062ba95bd36ff6ea61d70c71c023bee inline constexpr const char* kSdpaComputeOutWGSL = R"( @group(0) @binding(0) var t_out: array>; @group(0) @binding(1) var t_attn_weights_softmax: array; @@ -81,17 +81,21 @@ fn store_out_vec4(s: u32, d0: u32, h: u32, val: vec4) { } @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { let nrt = (params.S + TM - 1u) / TM; let nct = (params.D + TN - 1u) / TN; let tiles = nrt * nct; let total = tiles * params.Hq; - if (gid.x >= total) { + // 2D dispatch fold: recover the linear tile index across x/y. + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= total) { return; } - let h = gid.x / tiles; - let rem = gid.x % tiles; + let h = idx / tiles; + let rem = idx % tiles; let row_tile = rem / nct; let col_tile = rem % nct; let kvh = h / params.g; diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl index 6ef223c3a98..f0e349ceaff 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl @@ -20,9 +20,12 @@ var shared_sum: array; @compute @workgroup_size(WG_SIZE, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, - @builtin(local_invocation_id) lid: vec3) { + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { // One workgroup per (h, s) row of length context_len (= row_width). - let row_idx = wid.x; + // 2D dispatch fold: recover the linear row index. Keep the `valid` predicate + // below (NOT an early return) — the workgroupBarrier()s require uniform flow. + let row_idx = wid.x + wid.y * num_workgroups.x; let worker_id = lid.x; let base = row_idx * params.row_width; diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h index 94f0ab5790a..f1e7396fafe 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_softmax.wgsl - DO NOT EDIT. -// wgsl-sha256: e2714ec4c2400b37f6fd39c410075c519effc0273354a4f906fb924334809024 +// wgsl-sha256: bdd45cf344663533b243153200c507f41a90295751924e70452abcd5da4cdd5a inline constexpr const char* kSdpaSoftmaxWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_in: array; @@ -37,9 +37,12 @@ var shared_sum: array; @compute @workgroup_size(WG_SIZE, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, - @builtin(local_invocation_id) lid: vec3) { + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { // One workgroup per (h, s) row of length context_len (= row_width). - let row_idx = wid.x; + // 2D dispatch fold: recover the linear row index. Keep the `valid` predicate + // below (NOT an early return) — the workgroupBarrier()s require uniform flow. + let row_idx = wid.x + wid.y * num_workgroups.x; let worker_id = lid.x; let base = row_idx * params.row_width; From 88cbafc659c1f84da834c7f513ff357eed199124 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:57 -0700 Subject: [PATCH 11/14] =?UTF-8?q?[ExecuTorch][WebGPU]=202D=20compute=20dis?= =?UTF-8?q?patch=20tests=20=E2=80=94=20prefill=20golden=20+=20fold=20unit?= =?UTF-8?q?=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20584 **Test coverage for the 2D dispatch fold, stacked above the cap-lift op.** **Problem**: The 2D fold is load-bearing index math — a wrong `{x, y}` means out-of-bounds writes or dropped threads — and the prefill shapes that exercise it previously threw at the 1D cap, so they were untested. **Solution**: A device-free unit test for the fold arithmetic, plus two single-shot prefill SDPA golden configs that fold each kernel family. - **Before**: no coverage for >65535-workgroup dispatch; `llama1b_prefill_512`/`_2048` shapes threw at the cap - **After**: `fold_workgroup_count_2d` unit-tested at the cap boundaries, and the two prefill shapes run as goldens **Implementation**: - `test/native/test_dispatch_2d.cpp` — device-free unit test for `utils::fold_workgroup_count_2d`: the 1D fast path, the 2D fold, the real Llama-1B QK counts at S=512 (`{65535, 3}`) and S=2048 (`{65535, 33}`), and the needs-3rd-dimension throw; asserts each `{x, y}` covers `[0, count)` - `llama1b_prefill_512` + `llama1b_prefill_2048` configs appended to the byte-mirrored `CONFIGS` (`test_sdpa.py`) and `kSdpaConfigs` (`test_webgpu_native.cpp`) - Registers `webgpu_dispatch_2d_test` in CMake + the native CI script **Constraints**: - The Python/C++ config entries byte-mirror each other (kept in sync) - `add` shares the element-form path with QK, so it is covered structurally; a dedicated >16M-element `add` fold case is omitted as disproportionate Co-authored-with: Claude Code. ghstack-source-id: 399812923 @exported-using-ghexport Differential Revision: [D109517683](https://our.internmc.facebook.com/intern/diff/D109517683/) --- backends/webgpu/CMakeLists.txt | 8 +++ .../webgpu/scripts/test_webgpu_native_ci.sh | 4 +- .../webgpu/test/native/test_dispatch_2d.cpp | 60 +++++++++++++++++++ backends/webgpu/test/ops/test_sdpa.py | 3 + backends/webgpu/test/test_webgpu_native.cpp | 12 ++++ 5 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 backends/webgpu/test/native/test_dispatch_2d.cpp diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 072bcf43b8b..beb3f49e5b2 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -201,6 +201,14 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp ) target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest) + + # Device-free fold unit test (gtest_main provides main; no device needed). + add_webgpu_native_test( + webgpu_dispatch_2d_test test/native/test_dispatch_2d.cpp + ) + target_link_libraries( + webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main + ) endif() add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp) endif() diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 38195535732..0ed8c88e3b2 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -143,7 +143,7 @@ cmake \ "${EXECUTORCH_ROOT}" # ── Build + run every native test target that exists in this tree ──────────── -TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test) +TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dispatch_2d_test) BIN_DIR="${BUILD_DIR}/backends/webgpu" # Which targets are defined depends on which diffs are landed (native_test + @@ -212,6 +212,8 @@ if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" fi [[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test" +# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU). +[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test" echo "=== WebGPU native tests on Dawn: all run targets passed ===" diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp new file mode 100644 index 00000000000..582caa3c98e --- /dev/null +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -0,0 +1,60 @@ +/* + * 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 test for the pure 2D workgroup-count fold that lifts the +// 65535 per-dim dispatch cap. Exercises the fold arithmetic only — no GPU. + +#include + +#include + +#include +#include + +using executorch::backends::webgpu::utils::fold_workgroup_count_2d; +using executorch::backends::webgpu::utils::WgCount; + +namespace { + +constexpr uint32_t kMax = 65535u; + +// count <= max -> {count, 1}: the 1D fast path, byte-identical to the old path. +TEST(DispatchFold, FastPath1D) { + for (uint32_t count : {1u, kMax - 1u, kMax}) { + const WgCount got = fold_workgroup_count_2d(count, kMax, "test"); + EXPECT_EQ(got.x, count); + EXPECT_EQ(got.y, 1u); + } +} + +// count > max -> near-square {x, y}: fits the per-dim cap, covers every +// workgroup, and stays near-square so few invocations are inactive (launched - +// count is O(sqrt(count)); a flat {max, div_up} split would idle up to ~half). +TEST(DispatchFold, NearSquareFold) { + // Includes prefill-scale QK counts (Hq*ceil(S/4)*ceil(ctx/4)/wg) that fold: + // 131072 = S=2048 (32*512*512/64); 2097152 = large-S stress. + for (uint32_t count : + {kMax + 1u, 2u * kMax, 2u * kMax + 1u, 131072u, 2097152u}) { + const WgCount got = fold_workgroup_count_2d(count, kMax, "test"); + const uint64_t launched = static_cast(got.x) * got.y; + const uint32_t root = + static_cast(std::ceil(std::sqrt(static_cast(count)))); + EXPECT_LE(got.x, kMax) << "count=" << count; + EXPECT_LE(got.y, kMax) << "count=" << count; + EXPECT_GE(launched, count) << "count=" << count; + EXPECT_LT(launched - count, 2ull * root) + << "count=" << count << " launched=" << launched; + } +} + +// count > max^2 needs a 3rd dispatch dimension -> throws (out of scope). +TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) { + EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test")); +} + +} // namespace diff --git a/backends/webgpu/test/ops/test_sdpa.py b/backends/webgpu/test/ops/test_sdpa.py index 1f7a8242591..960c62e5203 100644 --- a/backends/webgpu/test/ops/test_sdpa.py +++ b/backends/webgpu/test/ops/test_sdpa.py @@ -61,6 +61,9 @@ class SdpaConfig: SdpaConfig("llama1b_decode", 32, 8, 64, 1, 512, 127), # D=6 is not a multiple of 4: the WebGPU head_dim%4 guard must reject it at load. SdpaConfig("reject_d6", 4, 4, 6, 4, 16, 0), + # 2D-dispatch cap (>65535 wg): S=512 folds QK; S=2048 folds QK+softmax+AV (cap+1). + SdpaConfig("llama1b_prefill_512", 32, 8, 64, 512, 512, 0), + SdpaConfig("llama1b_prefill_2048", 32, 8, 64, 2048, 2048, 0), ] diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index d0224231d8e..cb6d491f4ba 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -758,6 +758,18 @@ static const SdpaConfig kSdpaConfigs[] = { 16.0f, /*required=*/false, /*expect_reject=*/true}, + // 2D-dispatch cap (>65535 wg): S=512 folds QK; S=2048 folds QK+softmax+AV + // (cap+1). + {"llama1b_prefill_512", 32, 8, 64, 512, 512, 0, 16.0f, /*required=*/true}, + {"llama1b_prefill_2048", + 32, + 8, + 64, + 2048, + 2048, + 0, + 16.0f, + /*required=*/true}, }; // Ramp denominator; mirror of test_sdpa.py::_RAMP_DENOM (keep in sync). From 1b915a8ae2bfc0e0199494b0db8dccdf13477d35 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:58 -0700 Subject: [PATCH 12/14] [ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20651 **Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.** `mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line. Key changes: - `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`. - `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`). Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`). Co-authored-with: Claude Code. ghstack-source-id: 399812930 @exported-using-ghexport Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/) --- backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 17 +++++++++-------- backends/webgpu/runtime/ops/mul/binary_mul.wgsl | 7 +++++-- .../webgpu/runtime/ops/mul/binary_mul_wgsl.h | 9 ++++++--- backends/webgpu/runtime/ops/permute/Permute.cpp | 5 +++-- .../webgpu/runtime/ops/permute/permute.wgsl | 7 +++++-- .../webgpu/runtime/ops/permute/permute_wgsl.h | 9 ++++++--- 6 files changed, 34 insertions(+), 20 deletions(-) diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index 2ccb6c0e1bf..a14e583037a 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { const auto& in2_tensor = graph.get_tensor(in2_id); const auto& out_tensor = graph.get_tensor(out_id); - // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only). + // Rank guard (NCHW backend is <= 4 dims). if (out_tensor.dims.size() > kTensorMetaMaxNdim || in1_tensor.dims.size() > kTensorMetaMaxNdim || in2_tensor.dims.size() > kTensorMetaMaxNdim) { @@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { uint32_t wg_size = utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX); - uint32_t workgroup_count = - utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul"); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { 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}); + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "mul", 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; @@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om)); wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - g.device(), om.numel, wg_size, "mul(resize)"); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), om.numel, wg_size, "mul(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }; graph.add_tensor_resize_hook(in1_id, mul_resize); graph.add_tensor_resize_hook(in2_id, mul_resize); diff --git a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl index 1708cf08792..66722968e44 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl +++ b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl @@ -15,8 +15,11 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h b/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h index 4530d70e3dd..1c700615c7f 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h +++ b/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_mul.wgsl - DO NOT EDIT. -// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891 +// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e inline constexpr const char* kBinaryMulWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -32,8 +32,11 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/permute/Permute.cpp b/backends/webgpu/runtime/ops/permute/Permute.cpp index 5062c33cec1..c3b9b5c70c6 100644 --- a/backends/webgpu/runtime/ops/permute/Permute.cpp +++ b/backends/webgpu/runtime/ops/permute/Permute.cpp @@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector& args) { uint32_t wg_size = utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "permute"); WGPUConstantEntry wg_size_constant = {}; @@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y}); wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); diff --git a/backends/webgpu/runtime/ops/permute/permute.wgsl b/backends/webgpu/runtime/ops/permute/permute.wgsl index 521cfac1e66..e62fa5624d5 100644 --- a/backends/webgpu/runtime/ops/permute/permute.wgsl +++ b/backends/webgpu/runtime/ops/permute/permute.wgsl @@ -18,8 +18,11 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let out_bufi = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); if (out_bufi >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/permute/permute_wgsl.h b/backends/webgpu/runtime/ops/permute/permute_wgsl.h index 6ec41cc8446..b3d4684d54b 100644 --- a/backends/webgpu/runtime/ops/permute/permute_wgsl.h +++ b/backends/webgpu/runtime/ops/permute/permute_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from permute.wgsl - DO NOT EDIT. -// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197 +// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e inline constexpr const char* kPermuteWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,11 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let out_bufi = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); if (out_bufi >= out_meta.numel) { return; } From 3bd50ea4670ec204f965d8855e05b33f3b437b4b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:58 -0700 Subject: [PATCH 13/14] [ExecuTorch][WebGPU] Use requiredFeatures instance API on native + emscripten Dawn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20652 **Key the `timedWaitAny` instance setup to the actual Dawn API instead of `__EMSCRIPTEN__`, so native-rig Dawn and emscripten/emdawnwebgpu use the modern `requiredFeatures` path and only the vendored Dawn uses the legacy `capabilities.*` path.** The instance-descriptor setup was guarded by `#if defined(__EMSCRIPTEN__)`, which routed emscripten (emdawnwebgpu, emcc 4.0.19+) through the legacy `capabilities.*` API that no longer exists there. The guard now keys off the API actually present. Key changes: - `WebGPUDevice.cpp` — `#if defined(__EMSCRIPTEN__)` → `#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)`. The legacy `instance_desc.capabilities.*` path is taken only by the buck-vendored Dawn (which defines the macro); native cmake Dawn and emscripten leave it undefined and take the `requiredFeatures` / `WGPUInstanceFeatureName_TimedWaitAny` path. Co-authored-with: Claude Code. ghstack-source-id: 399812934 @exported-using-ghexport Differential Revision: [D110149678](https://our.internmc.facebook.com/intern/diff/D110149678/) --- backends/webgpu/runtime/WebGPUDevice.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backends/webgpu/runtime/WebGPUDevice.cpp b/backends/webgpu/runtime/WebGPUDevice.cpp index e69101851a2..9f48347c16b 100644 --- a/backends/webgpu/runtime/WebGPUDevice.cpp +++ b/backends/webgpu/runtime/WebGPUDevice.cpp @@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() { // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny. WGPUInstanceDescriptor instance_desc = {}; -#if defined(__EMSCRIPTEN__) + // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native + // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures. +#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES) instance_desc.capabilities.timedWaitAnyEnable = true; instance_desc.capabilities.timedWaitAnyMaxCount = 1; #else From 0c97a7f16f37fe3c9593c172f7253ef81b34b70c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 3 Jul 2026 14:36:59 -0700 Subject: [PATCH 14/14] [ExecuTorch][WebGPU] Convert remaining native tests to GTest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20706 Convert the remaining hand-rolled `int main()` + printf/`bool ok` native tests to GTest so the whole `backends/webgpu/test/` suite is uniform, filterable via `--gtest_filter`, and self-reporting (extends the GTest conversion already applied to `test_dynamic_shape`). The five converted files are a harness-only change — every test case, tensor shape, tolerance, artifact filename, and skip condition is preserved 1:1, only the pass/fail reporting mechanism changes — and this diff additionally wires the already-GTest `webgpu_dynamic_shape_test` into the CI runner so the dynamic-shape suite actually executes. Key changes: - `test/test_webgpu_native.cpp`, `test/native/test_dispatch_order.cpp`, `test/native/test_index.cpp`, `test/native/test_scratch_buffer.cpp`, `test/native/test_update_cache.cpp` — `main`+`printf`/`bool ok` accumulator → `TEST()` cases using `EXPECT_*`/`ASSERT_*`; each keeps a custom `main()` that brings up the WebGPU device once then `RUN_ALL_TESTS()` (device-absent still SKIPs by returning 0). `test_index`/`test_webgpu_native` use inclusive `EXPECT_LE(err, tol)` to match the original `err > tol` fail gate exactly. - `CMakeLists.txt` — move every native-test target into the `if(TARGET GTest::gtest)` block, linking `GTest::gtest`. - `scripts/test_webgpu_native_ci.sh` — add `-DEXECUTORCH_BUILD_TESTS=ON` to the native-test configure so the now-gtest-gated targets are defined, and wire `webgpu_dynamic_shape_test` into the runner: export its `.pte`s + goldens via `export_dynamic_shape_cases`, add it to the built/run target list behind the same `--target help` probe, and run it guarded (mirroring the `index` test). - `test/test_build_webgpu.sh` — add `-DEXECUTORCH_BUILD_TESTS=ON` so the local build script (which builds the now-gtest-gated targets unconditionally) still finds them. ghstack-source-id: 399812941 @exported-using-ghexport Differential Revision: [D110536636](https://our.internmc.facebook.com/intern/diff/D110536636/) --- backends/webgpu/CMakeLists.txt | 37 +- .../webgpu/scripts/test_webgpu_native_ci.sh | 13 +- .../test/native/test_dispatch_order.cpp | 154 +- backends/webgpu/test/native/test_index.cpp | 123 +- .../test/native/test_scratch_buffer.cpp | 109 +- .../webgpu/test/native/test_update_cache.cpp | 201 +-- backends/webgpu/test/test_build_webgpu.sh | 1 + backends/webgpu/test/test_webgpu_native.cpp | 1286 +++++++---------- 8 files changed, 825 insertions(+), 1099 deletions(-) diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index beb3f49e5b2..adbf4301413 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src) endfunction() if(EXECUTORCH_BUILD_WEBGPU_TEST) - add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp) - add_webgpu_native_test( - webgpu_dispatch_order_test test/native/test_dispatch_order.cpp - ) - add_webgpu_native_test( - webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp - ) - add_webgpu_native_test( - webgpu_update_cache_test test/native/test_update_cache.cpp - ) - - # Manifest-driven op-test framework: a generic gtest driver (webgpu_op_test) + - # its device-free util unit test. GTest needs -DEXECUTORCH_BUILD_TESTS=ON. + # All WebGPU native tests use GTest (device-dependent ones bring up the device + # in their own main(); the fold unit test is device-free via gtest_main). + # GTest needs -DEXECUTORCH_BUILD_TESTS=ON. if(NOT TARGET GTest::gtest) find_package(GTest QUIET) endif() @@ -195,12 +185,28 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) target_compile_options(webgpu_op_test_util_test PRIVATE -fexceptions) set_property(TARGET webgpu_op_test_util_test PROPERTY CXX_STANDARD 17) - # Dynamic-shape integration test: a gtest binary with its own main() that - # brings up the device once (like webgpu_op_test). + # Device-dependent native tests: each has its own main() that brings up the + # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main). + add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp) + target_link_libraries(webgpu_native_test PRIVATE GTest::gtest) + add_webgpu_native_test( + webgpu_dispatch_order_test test/native/test_dispatch_order.cpp + ) + target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest) + add_webgpu_native_test( + webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp + ) + target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest) + add_webgpu_native_test( + webgpu_update_cache_test test/native/test_update_cache.cpp + ) + target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest) add_webgpu_native_test( webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp ) target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest) + add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp) + target_link_libraries(webgpu_index_test PRIVATE GTest::gtest) # Device-free fold unit test (gtest_main provides main; no device needed). add_webgpu_native_test( @@ -210,5 +216,4 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main ) endif() - add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp) endif() diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 0ed8c88e3b2..810d8165303 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache" UPDATE_CACHE_OK=1 INDEX_DIR="/tmp/index" INDEX_OK=1 +DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape" +DYNAMIC_SHAPE_OK=1 EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte" EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin" EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin" @@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde export_all_index_models('${INDEX_DIR}') " || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; } +$PYTHON_EXECUTABLE -c " +from executorch.backends.webgpu.test.ops.dynamic_shape.test_dynamic_shape_export import export_dynamic_shape_cases +export_dynamic_shape_cases('${DYNAMIC_SHAPE_DIR}') +" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; } + # Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in # webgpu_native_test below (precise per-config error), so don't exit/mask here. $PYTHON_EXECUTABLE -c " @@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}" cmake \ -DEXECUTORCH_BUILD_WEBGPU=ON \ -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \ + -DEXECUTORCH_BUILD_TESTS=ON \ -DDawn_DIR="${Dawn_DIR}" \ -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ @@ -143,7 +151,7 @@ cmake \ "${EXECUTORCH_ROOT}" # ── Build + run every native test target that exists in this tree ──────────── -TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dispatch_2d_test) +TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test) BIN_DIR="${BUILD_DIR}/backends/webgpu" # Which targets are defined depends on which diffs are landed (native_test + @@ -211,6 +219,9 @@ fi if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" fi +if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then + "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" +fi [[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test" # Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU). [[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test" diff --git a/backends/webgpu/test/native/test_dispatch_order.cpp b/backends/webgpu/test/native/test_dispatch_order.cpp index 0f3eb5dea8e..d8aa627eff3 100644 --- a/backends/webgpu/test/native/test_dispatch_order.cpp +++ b/backends/webgpu/test/native/test_dispatch_order.cpp @@ -10,10 +10,13 @@ #include #include +#include + #include #include #include #include +#include #include #include #include @@ -24,23 +27,8 @@ using namespace executorch::runtime; namespace { -struct Case { - const char* name; - std::vector sizes; -}; - -// Mirrors _CASES in test_dispatch_order.py (add-chain or rms_norm+add chain). -const std::vector kCases = { - {"single", {16, 16}}, - {"chain3", {64, 64}}, - {"chain5_tiny", {1, 1}}, - {"chain5_wide", {7, 896}}, - {"chain8", {256, 256}}, - {"deep32", {128, 128}}, - {"large_chain", {1024, 1024}}, - {"het_small", {1, 1, 7, 896}}, - {"het_deep", {1, 1, 5, 256}}, -}; +// Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). +std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) std::vector read_f32_bin(const std::string& path) { std::ifstream f(path, std::ios::binary | std::ios::ate); @@ -59,53 +47,35 @@ std::vector read_f32_bin(const std::string& path) { return data; } -bool run_case(const std::string& dir, const Case& tc) { - printf("\n--- dispatch_order[%s] ---\n", tc.name); - const std::string base = dir + "/" + tc.name; +// Mirrors _CASES in test_dispatch_order.py (add-chain or rms_norm+add chain). +void run_case(const char* name, const std::vector& sizes) { + const std::string base = g_dir + "/" + name; std::vector input = read_f32_bin(base + ".input.bin"); std::vector golden = read_f32_bin(base + ".golden.bin"); - if (input.empty() || golden.empty()) { - printf("FAIL: could not read input/golden for %s\n", tc.name); - return false; - } + ASSERT_FALSE(input.empty() || golden.empty()) + << "could not read input/golden for " << name; Module module(base + ".pte"); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s.pte\n", tc.name); - return false; - } + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load " << name << ".pte"; size_t expected = 1; - for (int32_t d : tc.sizes) { + for (int32_t d : sizes) { expected *= static_cast(d); } - if (input.size() != expected) { - printf( - "FAIL: input numel %zu != expected %zu for %s\n", - input.size(), - expected, - tc.name); - return false; - } - auto x = make_tensor_ptr(tc.sizes, std::vector(input)); + ASSERT_EQ(input.size(), expected) + << "input numel " << input.size() << " != expected " << expected + << " for " << name; + auto x = make_tensor_ptr(sizes, std::vector(input)); auto result = module.forward({EValue(x)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor()) { - printf("FAIL: no tensor output\n"); - return false; - } + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (static_cast(out_tensor.numel()) != golden.size()) { - printf( - "FAIL: output numel %zu != golden %zu\n", - (size_t)out_tensor.numel(), - golden.size()); - return false; - } + ASSERT_EQ(static_cast(out_tensor.numel()), golden.size()) + << "output numel " << (size_t)out_tensor.numel() << " != golden " + << golden.size(); const float* out_data = out_tensor.const_data_ptr(); float max_abs_err = 0.0f; @@ -116,52 +86,76 @@ bool run_case(const std::string& dir, const Case& tc) { const float denom = std::max(std::abs(golden[i]), 1e-6f); max_rel_err = std::max(max_rel_err, abs_err / denom); } - printf( - "Max abs error: %e Max rel error: %e (%zu elements)\n", - max_abs_err, - max_rel_err, - golden.size()); // Lenient gate: pass iff abs<=tol OR rel<=tol (near-zero goldens). - if (max_abs_err > 1e-3f && max_rel_err > 1e-3f) { - printf("FAIL: dispatch_order[%s] exceeds tolerance 1e-3\n", tc.name); - return false; - } - printf("PASS: dispatch_order[%s]\n", tc.name); - return true; + EXPECT_FALSE(max_abs_err > 1e-3f && max_rel_err > 1e-3f) + << "dispatch_order[" << name + << "] exceeds tolerance 1e-3 (max_abs_err=" << max_abs_err + << " max_rel_err=" << max_rel_err << ", " << golden.size() + << " elements)"; } } // namespace +TEST(DispatchOrder, single) { + run_case("single", {16, 16}); +} + +TEST(DispatchOrder, chain3) { + run_case("chain3", {64, 64}); +} + +TEST(DispatchOrder, chain5_tiny) { + run_case("chain5_tiny", {1, 1}); +} + +TEST(DispatchOrder, chain5_wide) { + run_case("chain5_wide", {7, 896}); +} + +TEST(DispatchOrder, chain8) { + run_case("chain8", {256, 256}); +} + +TEST(DispatchOrder, deep32) { + run_case("deep32", {128, 128}); +} + +TEST(DispatchOrder, large_chain) { + run_case("large_chain", {1024, 1024}); +} + +TEST(DispatchOrder, het_small) { + run_case("het_small", {1, 1, 7, 896}); +} + +TEST(DispatchOrder, het_deep) { + run_case("het_deep", {1, 1, 5, 256}); +} + int main(int argc, char** argv) { - std::string dir = "/tmp/dispatch_order"; + ::testing::InitGoogleTest(&argc, argv); + + // Artifacts dir: env wins, else first positional arg, else default (gtest + // flags were already stripped by InitGoogleTest above). + g_dir = "/tmp/dispatch_order"; if (argc > 1) { - dir = argv[1]; + g_dir = argv[1]; } if (const char* env = std::getenv("WEBGPU_DISPATCH_ORDER_DIR")) { - dir = env; + g_dir = env; } WebGPUContext ctx; try { ctx = create_webgpu_context(); } catch (const std::exception& e) { - printf("SKIP: %s\n", e.what()); + std::printf("SKIP: %s\n", e.what()); return 0; } set_default_webgpu_context(&ctx); - printf("WebGPU device acquired (native); case dir: %s\n", dir.c_str()); - - bool ok = true; - for (const auto& tc : kCases) { - ok = run_case(dir, tc) && ok; - } + const int rc = RUN_ALL_TESTS(); set_default_webgpu_context(nullptr); destroy_webgpu_context(ctx); - - if (!ok) { - return 1; - } - printf("\nAll dispatch_order tests passed\n"); - return 0; + return rc; } diff --git a/backends/webgpu/test/native/test_index.cpp b/backends/webgpu/test/native/test_index.cpp index aed24c0a796..91f4ec9ea01 100644 --- a/backends/webgpu/test/native/test_index.cpp +++ b/backends/webgpu/test/native/test_index.cpp @@ -10,10 +10,14 @@ #include #include +#include + #include #include +#include #include #include +#include #include #include #include @@ -24,13 +28,8 @@ using namespace executorch::runtime; namespace { -// Names mirror test_index.py CONFIGS (self/idx/golden bins written per case). -constexpr const char* kIndexCases[] = { - "index_n16_m5", - "index_n8_rev", - "index_n32_m3", - "index_n4_rep", -}; +// Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). +std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) std::vector read_f32_bin(const std::string& path) { std::ifstream f(path, std::ios::binary | std::ios::ate); @@ -62,22 +61,19 @@ std::vector read_i32_bin(const std::string& path) { return data; } -bool run_case(const std::string& dir, const char* name) { - printf("\n--- Test: %s ---\n", name); - const std::string base = dir + "/" + name; +// index.Tensor: self [n] float, idx [m] int64 -> output [m]. Names mirror +// test_index.py CONFIGS (self/idx/golden bins written per case). +void run_case(const char* name) { + const std::string base = g_dir + "/" + name; std::vector self_data = read_f32_bin(base + ".self.bin"); std::vector idx32 = read_i32_bin(base + ".idx.bin"); std::vector golden = read_f32_bin(base + ".golden.bin"); - if (self_data.empty() || idx32.empty() || golden.empty()) { - printf("FAIL: could not read self/idx/golden for %s\n", name); - return false; - } + ASSERT_FALSE(self_data.empty() || idx32.empty() || golden.empty()) + << "could not read self/idx/golden for " << name; Module module(base + ".pte"); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s.pte\n", name); - return false; - } + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load " << name << ".pte"; const int32_t n = static_cast(self_data.size()); const int32_t m = static_cast(idx32.size()); @@ -87,33 +83,21 @@ bool run_case(const std::string& dir, const char* name) { auto idx = make_tensor_ptr({m}, std::vector(idx64)); auto result = module.forward({EValue(x), EValue(idx)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); // index.Tensor has exactly one output of shape [num_indices]; fail loud else. - if (outputs.size() != 1 || !outputs[0].isTensor()) { - printf("FAIL: expected exactly one tensor output\n"); - return false; - } + ASSERT_TRUE(outputs.size() == 1 && outputs[0].isTensor()) + << "expected exactly one tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (out_tensor.dim() != 1 || out_tensor.size(0) != m) { - printf( - "FAIL: output shape mismatch (dim %d size0 %d, expected [%d])\n", - (int)out_tensor.dim(), - (int)(out_tensor.dim() == 1 ? out_tensor.size(0) : -1), - m); - return false; - } - if (static_cast(out_tensor.numel()) != golden.size()) { - printf( - "FAIL: output numel %zu != golden %zu\n", - (size_t)out_tensor.numel(), - golden.size()); - return false; - } + ASSERT_TRUE(out_tensor.dim() == 1 && out_tensor.size(0) == m) + << "output shape mismatch (dim " << (int)out_tensor.dim() << " size0 " + << (int)(out_tensor.dim() == 1 ? out_tensor.size(0) : -1) + << ", expected [" << m << "])"; + ASSERT_EQ(static_cast(out_tensor.numel()), golden.size()) + << "output numel " << (size_t)out_tensor.numel() << " != golden " + << golden.size(); const float* out_data = out_tensor.const_data_ptr(); float max_abs_err = 0.0f; @@ -124,51 +108,54 @@ bool run_case(const std::string& dir, const char* name) { const float denom = std::max(std::abs(golden[i]), 1e-6f); max_rel_err = std::max(max_rel_err, abs_err / denom); } - printf( - "Max abs error: %e Max rel error: %e (%zu elements)\n", - max_abs_err, - max_rel_err, - golden.size()); - if (max_abs_err > 1e-3f || max_rel_err > 1e-3f) { - printf("FAIL: %s exceeds tolerance 1e-3\n", name); - return false; - } - printf("PASS: %s\n", name); - return true; + EXPECT_LE(max_abs_err, 1e-3f) << name << " max_abs_err=" << max_abs_err + << " (" << golden.size() << " elements)"; + EXPECT_LE(max_rel_err, 1e-3f) << name << " max_rel_err=" << max_rel_err + << " (" << golden.size() << " elements)"; } } // namespace +TEST(Index, N16M5) { + run_case("index_n16_m5"); +} + +TEST(Index, N8Rev) { + run_case("index_n8_rev"); +} + +TEST(Index, N32M3) { + run_case("index_n32_m3"); +} + +TEST(Index, N4Rep) { + run_case("index_n4_rep"); +} + int main(int argc, char** argv) { - std::string dir = "/tmp/index"; + ::testing::InitGoogleTest(&argc, argv); + + // Artifacts dir: env wins, else first positional arg, else default (gtest + // flags were already stripped by InitGoogleTest above). + g_dir = "/tmp/index"; if (argc > 1) { - dir = argv[1]; + g_dir = argv[1]; } if (const char* env = std::getenv("WEBGPU_INDEX_DIR")) { - dir = env; + g_dir = env; } WebGPUContext ctx; try { ctx = create_webgpu_context(); } catch (const std::exception& e) { - printf("SKIP: %s\n", e.what()); + std::printf("SKIP: %s\n", e.what()); return 0; } set_default_webgpu_context(&ctx); - printf("WebGPU device acquired (native); case dir: %s\n", dir.c_str()); - - bool ok = true; - for (const char* name : kIndexCases) { - ok = run_case(dir, name) && ok; - } + const int rc = RUN_ALL_TESTS(); set_default_webgpu_context(nullptr); destroy_webgpu_context(ctx); - - if (!ok) { - return 1; - } - printf("\nAll index tests passed\n"); - return 0; + return rc; } diff --git a/backends/webgpu/test/native/test_scratch_buffer.cpp b/backends/webgpu/test/native/test_scratch_buffer.cpp index 7a4df6e9d00..98cf3648c6b 100644 --- a/backends/webgpu/test/native/test_scratch_buffer.cpp +++ b/backends/webgpu/test/native/test_scratch_buffer.cpp @@ -14,17 +14,27 @@ #include +#include + #include #include #include -#include #include +#include #include using namespace executorch::backends::webgpu; namespace { +// WebGPU context; set from create_webgpu_context() in main() before +// RUN_ALL_TESTS(). +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +WGPUInstance g_instance = nullptr; +WGPUDevice g_device = nullptr; +WGPUQueue g_queue = nullptr; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + struct MapCb { std::atomic status{WGPUMapAsyncStatus_Error}; }; @@ -79,56 +89,46 @@ std::vector readback( return out; } +} // namespace + // Tier 1: allocation, zero-size guard, distinct non-null handles. -bool tier1_alloc(WGPUDevice device) { - printf("\n--- scratch[tier1: allocation] ---\n"); +TEST(ScratchBuffer, Tier1Alloc) { WebGPUGraph g; - g.set_device(device); + g.set_device(g_device); WGPUBuffer a = g.create_scratch_buffer(64 * sizeof(float)); WGPUBuffer z = g.create_scratch_buffer(0); // guarded to 4 bytes WGPUBuffer b = g.create_scratch_buffer(64 * sizeof(float)); - const bool ok = a && z && b && a != b && a != z && b != z; - printf(ok ? "PASS: allocation\n" : "FAIL: allocation\n"); - return ok; // graph dtor releases all three here + EXPECT_TRUE(a && z && b && a != b && a != z && b != z); + // graph dtor releases all three here } // Tier 2: host->scratch write, scratch->staging copy, read-back round-trip. -bool tier2_roundtrip( - WGPUInstance instance, - WGPUDevice device, - WGPUQueue queue) { - printf("\n--- scratch[tier2: copy round-trip] ---\n"); - bool ok = true; +TEST(ScratchBuffer, Tier2Roundtrip) { for (int n : {1, 7, 1024}) { WebGPUGraph g; - g.set_device(device); + g.set_device(g_device); WGPUBuffer s = g.create_scratch_buffer(n * sizeof(float)); std::vector in(n); for (int i = 0; i < n; i++) { in[i] = static_cast(i) * 0.5f + 1.0f; } - wgpuQueueWriteBuffer(queue, s, 0, in.data(), n * sizeof(float)); + wgpuQueueWriteBuffer(g_queue, s, 0, in.data(), n * sizeof(float)); std::vector back = - readback(instance, device, queue, s, n * sizeof(float)); + readback(g_instance, g_device, g_queue, s, n * sizeof(float)); float max_err = 0.0f; for (int i = 0; i < n; i++) { max_err = std::max(max_err, std::abs(back[i] - in[i])); } - printf(" n=%d max abs error %e\n", n, max_err); - if (max_err != 0.0f) { // pure copy: must be bit-exact - ok = false; - } + // pure copy: must be bit-exact + EXPECT_EQ(max_err, 0.0f) << "n=" << n << " max abs error " << max_err; } - printf(ok ? "PASS: copy round-trip\n" : "FAIL: copy round-trip\n"); - return ok; } // Tier 3a: bind scratch as a Storage buffer in a compute pass (its real use). -bool tier3_compute(WGPUInstance instance, WGPUDevice device, WGPUQueue queue) { - printf("\n--- scratch[tier3: compute Storage round-trip] ---\n"); +TEST(ScratchBuffer, Tier3Compute) { const int n = 256; WebGPUGraph g; - g.set_device(device); + g.set_device(g_device); WGPUBuffer s = g.create_scratch_buffer(n * sizeof(float)); const char* kWgsl = @@ -144,7 +144,7 @@ bool tier3_compute(WGPUInstance instance, WGPUDevice device, WGPUQueue queue) { wgsl.code = {kWgsl, WGPU_STRLEN}; WGPUShaderModuleDescriptor smd = {}; smd.nextInChain = &wgsl.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &smd); + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(g_device, &smd); WGPUBindGroupLayoutEntry ble = {}; ble.binding = 0; @@ -153,18 +153,18 @@ bool tier3_compute(WGPUInstance instance, WGPUDevice device, WGPUQueue queue) { WGPUBindGroupLayoutDescriptor bld = {}; bld.entryCount = 1; bld.entries = &ble; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bld); + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(g_device, &bld); WGPUPipelineLayoutDescriptor pld = {}; pld.bindGroupLayoutCount = 1; pld.bindGroupLayouts = &bgl; - WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(device, &pld); + WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(g_device, &pld); WGPUComputePipelineDescriptor cpd = {}; cpd.layout = pl; cpd.compute.module = shader; cpd.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipe = wgpuDeviceCreateComputePipeline(device, &cpd); + WGPUComputePipeline pipe = wgpuDeviceCreateComputePipeline(g_device, &cpd); WGPUBindGroupEntry bge = {}; bge.binding = 0; @@ -174,10 +174,10 @@ bool tier3_compute(WGPUInstance instance, WGPUDevice device, WGPUQueue queue) { bgd.layout = bgl; bgd.entryCount = 1; bgd.entries = &bge; - WGPUBindGroup bg = wgpuDeviceCreateBindGroup(device, &bgd); + WGPUBindGroup bg = wgpuDeviceCreateBindGroup(g_device, &bgd); WGPUCommandEncoderDescriptor ed = {}; - WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device, &ed); + WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(g_device, &ed); WGPUComputePassDescriptor pd = {}; WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(enc, &pd); wgpuComputePassEncoderSetPipeline(pass, pipe); @@ -187,18 +187,17 @@ bool tier3_compute(WGPUInstance instance, WGPUDevice device, WGPUQueue queue) { wgpuComputePassEncoderRelease(pass); WGPUCommandBufferDescriptor cd = {}; WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, &cd); - wgpuQueueSubmit(queue, 1, &cmd); + wgpuQueueSubmit(g_queue, 1, &cmd); wgpuCommandBufferRelease(cmd); wgpuCommandEncoderRelease(enc); std::vector back = - readback(instance, device, queue, s, n * sizeof(float)); + readback(g_instance, g_device, g_queue, s, n * sizeof(float)); float max_err = 0.0f; for (int i = 0; i < n; i++) { const float expected = static_cast(i) * 2.0f + 1.0f; max_err = std::max(max_err, std::abs(back[i] - expected)); } - printf(" max abs error %e (%d elements)\n", max_err, n); wgpuBindGroupRelease(bg); wgpuComputePipelineRelease(pipe); @@ -206,56 +205,40 @@ bool tier3_compute(WGPUInstance instance, WGPUDevice device, WGPUQueue queue) { wgpuBindGroupLayoutRelease(bgl); wgpuShaderModuleRelease(shader); - const bool ok = max_err == 0.0f; - printf( - ok ? "PASS: compute Storage round-trip\n" : "FAIL: compute round-trip\n"); - return ok; + EXPECT_EQ(max_err, 0.0f) << "max abs error " << max_err << " (" << n + << " elements)"; } // Tier 3b: many scratch buffers across repeated graphs; dtor must release all. -bool tier3_lifecycle(WGPUDevice device) { - printf("\n--- scratch[tier3: lifecycle/stress] ---\n"); - bool ok = true; +TEST(ScratchBuffer, Tier3Lifecycle) { for (int iter = 0; iter < 50; iter++) { WebGPUGraph g; - g.set_device(device); + g.set_device(g_device); for (int k = 0; k < 256; k++) { WGPUBuffer b = g.create_scratch_buffer(static_cast(k % 17) * sizeof(float)); - ok = ok && b != nullptr; + EXPECT_NE(b, nullptr); } } // each graph's dtor releases its 256 buffers here - printf( - ok ? "PASS: lifecycle/stress (50 graphs x 256 buffers)\n" - : "FAIL: lifecycle/stress (null buffer)\n"); - return ok; } -} // namespace +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); -int main() { WebGPUContext ctx; try { ctx = create_webgpu_context(); } catch (const std::exception& e) { - printf("SKIP: %s\n", e.what()); + std::printf("SKIP: %s\n", e.what()); return 0; } set_default_webgpu_context(&ctx); - printf("WebGPU device acquired (native)\n"); - - bool ok = true; - ok = tier1_alloc(ctx.device) && ok; - ok = tier2_roundtrip(ctx.instance, ctx.device, ctx.queue) && ok; - ok = tier3_compute(ctx.instance, ctx.device, ctx.queue) && ok; - ok = tier3_lifecycle(ctx.device) && ok; + g_instance = ctx.instance; + g_device = ctx.device; + g_queue = ctx.queue; + const int rc = RUN_ALL_TESTS(); set_default_webgpu_context(nullptr); destroy_webgpu_context(ctx); - - if (!ok) { - return 1; - } - printf("\nAll scratch_buffer tests passed\n"); - return 0; + return rc; } diff --git a/backends/webgpu/test/native/test_update_cache.cpp b/backends/webgpu/test/native/test_update_cache.cpp index 3f932ea7f03..dad859af669 100644 --- a/backends/webgpu/test/native/test_update_cache.cpp +++ b/backends/webgpu/test/native/test_update_cache.cpp @@ -10,10 +10,13 @@ #include #include +#include + #include #include #include #include +#include #include #include @@ -23,6 +26,9 @@ using namespace executorch::runtime; namespace { +// Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). +std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) + struct UpdateCacheCase { const char* name; int s; @@ -40,20 +46,10 @@ constexpr UpdateCacheCase kCases[] = { {"shape_b_offset", 3, 4, 8, 16, 10}, }; -bool run_case(const std::string& dir, const UpdateCacheCase& tc) { - printf( - "\n--- Test: update_cache[%s] (S=%d,H=%d,D=%d,Cmax=%d,pos=%d) ---\n", - tc.name, - tc.s, - tc.h, - tc.d, - tc.cmax, - tc.input_pos); - Module module(dir + "/" + tc.name + ".pte"); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s.pte\n", tc.name); - return false; - } +void run_case(const UpdateCacheCase& tc) { + Module module(g_dir + "/" + tc.name + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load " << tc.name << ".pte"; const int vnumel = tc.s * tc.h * tc.d; const int cnumel = tc.cmax * tc.h * tc.d; @@ -79,37 +75,24 @@ bool run_case(const std::string& dir, const UpdateCacheCase& tc) { auto v = make_tensor_ptr({1, tc.s, tc.h, tc.d}, std::vector(value)); auto c = make_tensor_ptr({1, tc.cmax, tc.h, tc.d}, std::vector(cache)); auto result = module.forward({EValue(v), EValue(c)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor()) { - printf("FAIL: no tensor output\n"); - return false; - } + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (static_cast(out_tensor.numel()) != cnumel) { - printf( - "FAIL: output numel %zu != expected %d\n", - (size_t)out_tensor.numel(), - cnumel); - return false; - } + ASSERT_EQ(static_cast(out_tensor.numel()), cnumel) + << "output numel " << (size_t)out_tensor.numel() << " != expected " + << cnumel; const float* out_data = out_tensor.const_data_ptr(); float max_abs_err = 0.0f; for (int i = 0; i < cnumel; i++) { max_abs_err = std::max(max_abs_err, std::abs(out_data[i] - ref[i])); } - printf("Max abs error: %e (checked %d elements)\n", max_abs_err, cnumel); // update_cache is a pure scatter copy: the output must be bit-exact. - if (max_abs_err > 0.0f) { - printf("FAIL: update_cache[%s] not bit-exact\n", tc.name); - return false; - } - printf("PASS: update_cache[%s]\n", tc.name); - return true; + EXPECT_EQ(max_abs_err, 0.0f) + << "update_cache[" << tc.name << "] not bit-exact (max abs error " + << max_abs_err << ", checked " << cnumel << " elements)"; } struct ReplayCase { @@ -120,18 +103,11 @@ struct ReplayCase { }; // Multi-step advancing-input_pos cache accumulation, mirroring VulkanSDPATest. -bool run_replay(const std::string& dir, const ReplayCase& rc) { +void run_replay(const ReplayCase& rc) { int cmax = 0; for (int s : rc.seq_lens) { cmax += s; } - printf( - "\n--- Replay: update_cache[%s] (H=%d,D=%d,Cmax=%d,%zu steps) ---\n", - rc.name, - rc.h, - rc.d, - cmax, - rc.seq_lens.size()); const int cnumel = cmax * rc.h * rc.d; std::vector cache(cnumel); @@ -141,7 +117,6 @@ bool run_replay(const std::string& dir, const ReplayCase& rc) { std::vector ref(cache); int input_pos = 0; - bool ok = true; for (size_t step = 0; step < rc.seq_lens.size(); step++) { const int s = rc.seq_lens[step]; const int vnumel = s * rc.h * rc.d; @@ -151,31 +126,22 @@ bool run_replay(const std::string& dir, const ReplayCase& rc) { value[i] = (base + static_cast(i)) * 0.25f; } - const std::string fname = dir + "/" + rc.name + "_step" + + const std::string fname = g_dir + "/" + rc.name + "_step" + std::to_string(step) + "_S" + std::to_string(s) + "_pos" + std::to_string(input_pos) + ".pte"; Module module(fname); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s\n", fname.c_str()); - return false; - } + ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << fname; auto v = make_tensor_ptr({1, s, rc.h, rc.d}, std::vector(value)); auto c = make_tensor_ptr({1, cmax, rc.h, rc.d}, std::vector(cache)); auto result = module.forward({EValue(v), EValue(c)}); - if (!result.ok()) { - printf( - "FAIL: forward failed step %zu (error %d)\n", - step, - (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed step " << step << " (error " + << (int)result.error() << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor() || - static_cast(outputs[0].toTensor().numel()) != cnumel) { - printf("FAIL: bad cache output at step %zu\n", step); - return false; - } + ASSERT_TRUE( + !outputs.empty() && outputs[0].isTensor() && + static_cast(outputs[0].toTensor().numel()) == cnumel) + << "bad cache output at step " << step; const float* out_data = outputs[0].toTensor().const_data_ptr(); const int dst_offset = input_pos * rc.h * rc.d; @@ -190,24 +156,12 @@ bool run_replay(const std::string& dir, const ReplayCase& rc) { max_abs_err = std::max(max_abs_err, std::abs(out_data[i] - ref[i])); cache[i] = out_data[i]; // thread the accumulated cache into the next step } - printf( - " step %zu (S=%d,pos=%d): max abs error %e\n", - step, - s, - input_pos, - max_abs_err); - if (max_abs_err > 0.0f) { // pure scatter copy: must be bit-exact - ok = false; - } + // pure scatter copy: must be bit-exact + EXPECT_EQ(max_abs_err, 0.0f) + << "step " << step << " (S=" << s << ",pos=" << input_pos + << "): max abs error " << max_abs_err; input_pos += s; } - - if (ok) { - printf("PASS: update_cache[%s] replay\n", rc.name); - } else { - printf("FAIL: update_cache[%s] replay\n", rc.name); - } - return ok; } struct NegativeCase { @@ -216,76 +170,75 @@ struct NegativeCase { }; // Single-op, single-guard-violation cases: rejection maps to the named guard. -bool run_negative_case(const std::string& dir, const NegativeCase& nc) { - printf( - "\n--- Negative: update_cache[%s] (expect rejection: %s) ---\n", - nc.name, - nc.guard); - Module module(dir + "/" + nc.name + ".pte"); +void run_negative_case(const NegativeCase& nc) { + Module module(g_dir + "/" + nc.name + ".pte"); const Error err = module.load_forward(); // init catches the guard throw -> this code; other errors = setup failure. - if (err != Error::DelegateInvalidCompatibility) { - printf( - "FAIL: %s.pte -> error %d; expected DelegateInvalidCompatibility " - "from the '%s' guard\n", - nc.name, - (int)err, - nc.guard); - return false; - } - printf("PASS: rejected with DelegateInvalidCompatibility (%s)\n", nc.guard); - return true; + EXPECT_EQ(err, Error::DelegateInvalidCompatibility) + << nc.name << ".pte -> error " << (int)err + << "; expected DelegateInvalidCompatibility from the '" << nc.guard + << "' guard"; } } // namespace -int main(int argc, char** argv) { - std::string dir = "/tmp/update_cache"; - if (argc > 1) { - dir = argv[1]; - } - if (const char* env = std::getenv("WEBGPU_UPDATE_CACHE_DIR")) { - dir = env; - } - - WebGPUContext ctx; - try { - ctx = create_webgpu_context(); - } catch (const std::exception& e) { - printf("SKIP: %s\n", e.what()); - return 0; - } - set_default_webgpu_context(&ctx); - printf("WebGPU device acquired (native); case dir: %s\n", dir.c_str()); - - bool ok = true; +// Single-step scatter cases (prefill / offset / shape variants): the op output +// must equal the inline integer-exact scatter reference. +TEST(UpdateCache, ScatterCases) { for (const auto& tc : kCases) { - ok = run_case(dir, tc) && ok; + run_case(tc); } +} +// Multi-step advancing-input_pos cache accumulation, mirroring VulkanSDPATest. +TEST(UpdateCache, Replay) { const std::vector kReplays = { {"seqA", 4, 4, {3, 1, 1, 5, 1, 1, 2}}, {"seqB", 2, 8, {3, 1, 1, 5, 1, 1}}, {"llama3", 8, 128, {111, 1, 1, 1, 57, 1, 1}}, }; for (const auto& rc : kReplays) { - ok = run_replay(dir, rc) && ok; + run_replay(rc); } +} +// Guard-violation cases: each must be rejected with +// DelegateInvalidCompatibility. +TEST(UpdateCache, Negative) { const NegativeCase kNegatives[] = { {"neg_batch", "batch must be 1"}, {"neg_fp16", "fp32-only"}, }; for (const auto& nc : kNegatives) { - ok = run_negative_case(dir, nc) && ok; + run_negative_case(nc); } +} - set_default_webgpu_context(nullptr); - destroy_webgpu_context(ctx); +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); - if (!ok) { - return 1; + // Artifacts dir: env wins, else first positional arg, else default (gtest + // flags were already stripped by InitGoogleTest above). + std::string dir = "/tmp/update_cache"; + if (argc > 1) { + dir = argv[1]; } - printf("\nAll update_cache tests passed\n"); - return 0; + if (const char* env = std::getenv("WEBGPU_UPDATE_CACHE_DIR")) { + dir = env; + } + g_dir = dir; + + WebGPUContext ctx; + try { + ctx = create_webgpu_context(); + } catch (const std::exception& e) { + std::printf("SKIP: %s\n", e.what()); + return 0; + } + set_default_webgpu_context(&ctx); + + const int rc = RUN_ALL_TESTS(); + set_default_webgpu_context(nullptr); + destroy_webgpu_context(ctx); + return rc; } diff --git a/backends/webgpu/test/test_build_webgpu.sh b/backends/webgpu/test/test_build_webgpu.sh index 2fd1dea1a52..1c79d17cf06 100755 --- a/backends/webgpu/test/test_build_webgpu.sh +++ b/backends/webgpu/test/test_build_webgpu.sh @@ -85,6 +85,7 @@ cmake \ -DEXECUTORCH_BUILD_WEBGPU=ON \ -DDawn_DIR="${Dawn_DIR}" \ -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \ + -DEXECUTORCH_BUILD_TESTS=ON \ -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index cb6d491f4ba..556eb0127b4 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -12,10 +12,14 @@ #include #include +#include + #include #include +#include #include #include +#include #include #include #include @@ -24,27 +28,31 @@ using namespace executorch::backends::webgpu; using namespace executorch::extension; using namespace executorch::runtime; +namespace { + +// Environment-derived config; captured in main() before RUN_ALL_TESTS(). +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +std::string g_update_cache_model_path; +std::string g_qlinear_dir; +std::string g_prepack_model_path, g_prepack_golden_path; +std::string g_prepack2_model_path, g_prepack2_golden_path; +std::string g_prepack_tied_model_path, g_prepack_tied_golden_path; +std::string g_sdpa_dir; +std::string g_symint_blob; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + #ifdef WGPU_BACKEND_ENABLE_PROFILING // Capacity-overrun must throw; runs without a device or TimestampQuery. -static bool test_query_pool_overrun_throws() { - printf("\n--- Test: WebGPUQueryPool capacity-overrun guard ---\n"); +void test_query_pool_overrun_throws() { WebGPUQueryPool qp; - try { - qp.reset(1); - } catch (const std::exception&) { - printf("PASS: reset beyond capacity throws\n"); - return true; - } - printf("FAIL: reset beyond capacity did not throw\n"); - return false; + EXPECT_THROW(qp.reset(1), std::exception) + << "reset beyond capacity did not throw"; } // WebGPUQueryPool roundtrip: time a probe pass; assert non-zero GPU duration. -static bool test_query_pool_roundtrip(const WebGPUContext& ctx) { - printf("\n--- Test: WebGPUQueryPool roundtrip ---\n"); +void test_query_pool_roundtrip(const WebGPUContext& ctx) { if (!ctx.timestamp_supported) { - printf("SKIP: adapter lacks TimestampQuery feature\n"); - return true; + GTEST_SKIP() << "adapter lacks TimestampQuery feature"; } WGPUDevice device = ctx.device; @@ -134,33 +142,21 @@ static bool test_query_pool_roundtrip(const WebGPUContext& ctx) { wgpuBindGroupRelease(bg); wgpuShaderModuleRelease(shader); - if (qp.results().size() != 1) { - printf("FAIL: expected 1 duration, got %zu\n", qp.results().size()); - return false; - } + ASSERT_EQ(qp.results().size(), 1u) + << "expected 1 duration, got " << qp.results().size(); const uint64_t dur = qp.results()[0].execution_duration_ns; printf(" probe duration: %llu ns\n", (unsigned long long)dur); - if (dur == 0) { - printf("FAIL: probe duration is zero (expected monotonic non-zero)\n"); - return false; - } - printf("PASS: WebGPUQueryPool roundtrip -- non-zero GPU kernel duration\n"); - return true; + EXPECT_NE(dur, 0u) << "probe duration is zero (expected monotonic non-zero)"; } #endif // WGPU_BACKEND_ENABLE_PROFILING -static bool test_update_cache(const std::string& model_path) { +void test_update_cache(const std::string& model_path) { // update_cache: value [1,2,2,4] scattered into cache [1,8,2,4] at // input_pos=0. - printf( - "\n--- Test: update_cache (value[1,2,2,4] -> cache[1,8,2,4], pos=0) ---\n"); - Module module(model_path); auto err = module.load_forward(); - if (err != Error::Ok) { - printf("FAIL: could not load forward method (error %d)\n", (int)err); - return false; - } + ASSERT_EQ(err, Error::Ok) + << "could not load forward method (error " << (int)err << ")"; printf("Model loaded: %s\n", model_path.c_str()); constexpr int S = 2, H = 2, D = 4, Cmax = 8; @@ -187,24 +183,15 @@ static bool test_update_cache(const std::string& model_path) { auto v = make_tensor_ptr({1, S, H, D}, std::vector(value)); auto c = make_tensor_ptr({1, Cmax, H, D}, std::vector(cache)); auto result = module.forward({EValue(v), EValue(c)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor()) { - printf("FAIL: no tensor output\n"); - return false; - } + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (out_tensor.numel() != cnumel) { - printf( - "FAIL: output numel %zu != expected %d\n", - (size_t)out_tensor.numel(), - cnumel); - return false; - } + ASSERT_EQ((int)out_tensor.numel(), cnumel) + << "output numel " << (size_t)out_tensor.numel() << " != expected " + << cnumel; const float* out_data = out_tensor.const_data_ptr(); float max_abs_err = 0.0f; @@ -212,15 +199,10 @@ static bool test_update_cache(const std::string& model_path) { max_abs_err = std::max(max_abs_err, std::abs(out_data[i] - ref[i])); } printf("Max abs error: %e (checked %d elements)\n", max_abs_err, cnumel); - if (max_abs_err > 1e-3f) { - printf("FAIL: max error exceeds tolerance 1e-3\n"); - return false; - } - printf("PASS: update_cache test\n"); - return true; + EXPECT_LE(max_abs_err, 1e-3f) << "max error exceeds tolerance 1e-3"; } -static std::vector load_golden(const std::string& path, size_t numel) { +std::vector load_golden(const std::string& path, size_t numel) { // Load a raw little-endian fp32 golden written by the export .py (the native // binary has no ATen/torch, so the reference is computed offline). std::vector g(numel); @@ -241,7 +223,7 @@ static std::vector load_golden(const std::string& path, size_t numel) { // value can't blow up the rel metric (the kernel's ~1e-8 abs error is the real // signal at llama3 scale). Sets the reported maxima; true iff all elements // pass. -static bool sdpa_within_tol( +bool sdpa_within_tol( const float* out, const float* golden, int n, @@ -277,7 +259,7 @@ struct Q4gswConfig { // Llama-3.2-1B linear shapes (q/o/k/v/gate/up/down + lm_head) + 4k/8k prefill. // tol scales with K (fp32 accum depth), not M; down_proj (K=8192) is looser. -static const Q4gswConfig kQ4gswConfigs[] = { +const Q4gswConfig kQ4gswConfigs[] = { // name M K N tol_abs tol_rel req heavy {"q_proj", 1, 2048, 2048, 1e-4f, 1e-3f, true, false}, {"kv_proj", 1, 2048, 512, 1e-4f, 1e-3f, true, false}, @@ -299,23 +281,36 @@ static const Q4gswConfig kQ4gswConfigs[] = { }; // /16 ramp over the flat index; mirrors test_quantized_linear.py _ramp_input. -static float q4gsw_ramp(int i) { +float q4gsw_ramp(int i) { return static_cast((i % 17) - 8) / 16.0f; } -// Fwd decl of the per-element abs-OR-rel tolerance helper (defined below). -static bool quant_within_tol( +// Per-element abs-OR-rel tolerance helper. +bool quant_within_tol( const float* out, const float* golden, int n, float atol, float rtol, float* ma, - float* mr); + float* mr) { + float max_abs = 0.0f, max_rel = 0.0f; + bool ok = true; + for (int i = 0; i < n; i++) { + const float ae = std::abs(out[i] - golden[i]); + const float re = ae / std::max(std::abs(golden[i]), 1e-6f); + max_abs = std::max(max_abs, ae); + max_rel = std::max(max_rel, re); + if (ae > atol && re > rtol) { + ok = false; + } + } + *ma = max_abs; + *mr = max_rel; + return ok; +} -static std::vector load_indices( - const std::string& path, - size_t numel) { +std::vector load_indices(const std::string& path, size_t numel) { // Load raw little-endian int32 indices written by the export .py. std::vector g(numel); FILE* f = std::fopen(path.c_str(), "rb"); @@ -330,7 +325,7 @@ static std::vector load_indices( return g; } -static bool test_embedding_q4gsw( +void test_embedding_q4gsw( const std::string& model_path, const std::string& indices_path, const std::string& golden_path, @@ -347,44 +342,29 @@ static bool test_embedding_q4gsw( Module module(model_path); auto err = module.load_forward(); - if (err != Error::Ok) { - printf("FAIL: could not load forward method (error %d)\n", (int)err); - return false; - } + ASSERT_EQ(err, Error::Ok) + << "could not load forward method (error " << (int)err << ")"; printf("Model loaded: %s\n", model_path.c_str()); std::vector idx32 = load_indices(indices_path, num_indices); std::vector golden = load_golden(golden_path, out_numel); - if (idx32.empty() || golden.empty()) { - printf( - "FAIL: could not load indices %s / golden %s\n", - indices_path.c_str(), - golden_path.c_str()); - return false; - } + ASSERT_FALSE(idx32.empty() || golden.empty()) + << "could not load indices " << indices_path << " / golden " + << golden_path; // int64 at the program boundary; copy_inputs narrows to the int32 buffer. std::vector idx64(idx32.begin(), idx32.end()); auto idx = make_tensor_ptr({num_indices}, std::move(idx64)); auto result = module.forward({EValue(idx)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor()) { - printf("FAIL: no tensor output\n"); - return false; - } + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (out_tensor.numel() != out_numel) { - printf( - "FAIL: output numel %zu != expected %d\n", - (size_t)out_tensor.numel(), - out_numel); - return false; - } + ASSERT_EQ((int)out_tensor.numel(), out_numel) + << "output numel " << (size_t)out_tensor.numel() << " != expected " + << out_numel; const float* out_data = out_tensor.const_data_ptr(); float max_abs_err = 0.0f, max_rel_err = 0.0f; @@ -401,39 +381,10 @@ static bool test_embedding_q4gsw( max_abs_err, max_rel_err, out_numel); - if (!pass) { - printf("FAIL: embedding_q4gsw exceeds tolerance 1e-3 (abs AND rel)\n"); - return false; - } - printf("PASS: embedding_q4gsw test\n"); - return true; -} - -static bool quant_within_tol( - const float* out, - const float* golden, - int n, - float atol, - float rtol, - float* ma, - float* mr) { - float max_abs = 0.0f, max_rel = 0.0f; - bool ok = true; - for (int i = 0; i < n; i++) { - const float ae = std::abs(out[i] - golden[i]); - const float re = ae / std::max(std::abs(golden[i]), 1e-6f); - max_abs = std::max(max_abs, ae); - max_rel = std::max(max_rel, re); - if (ae > atol && re > rtol) { - ok = false; - } - } - *ma = max_abs; - *mr = max_rel; - return ok; + EXPECT_TRUE(pass) << "embedding_q4gsw exceeds tolerance 1e-3 (abs AND rel)"; } -static bool test_rope( +void test_rope( const std::string& model_path, const std::string& xq_golden_path, const std::string& xk_golden_path, @@ -456,10 +407,8 @@ static bool test_rope( Module module(model_path); auto err = module.load_forward(); - if (err != Error::Ok) { - printf("FAIL: could not load forward method (error %d)\n", (int)err); - return false; - } + ASSERT_EQ(err, Error::Ok) + << "could not load forward method (error " << (int)err << ")"; printf("Model loaded: %s\n", model_path.c_str()); // ((i % mod) - off) / 16: exact in fp32, matches test_rope.py::_ramp. @@ -486,43 +435,30 @@ static bool test_rope( auto result = module.forward({EValue(xqt), EValue(xkt), EValue(fct), EValue(fst)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); // Outputs in graph order [0]=xq_out, [1]=xk_out (positional; the numel check // below guards a swap, since NH != NKV under GQA). - if (outputs.size() < 2 || !outputs[0].isTensor() || !outputs[1].isTensor()) { - printf("FAIL: expected 2 tensor outputs, got %zu\n", outputs.size()); - return false; - } + ASSERT_TRUE( + outputs.size() >= 2 && outputs[0].isTensor() && outputs[1].isTensor()) + << "expected 2 tensor outputs, got " << outputs.size(); const auto& xq_t = outputs[0].toTensor(); const auto& xk_t = outputs[1].toTensor(); - if (xq_t.numel() != xq_numel || xk_t.numel() != xk_numel) { - printf( - "FAIL: output shapes [%zu,%zu] != expected [%d,%d]\n", - (size_t)xq_t.numel(), - (size_t)xk_t.numel(), - xq_numel, - xk_numel); - return false; - } + ASSERT_TRUE(xq_t.numel() == xq_numel && xk_t.numel() == xk_numel) + << "output shapes [" << (size_t)xq_t.numel() << "," + << (size_t)xk_t.numel() << "] != expected [" << xq_numel << "," + << xk_numel << "]"; const float* xq_out = xq_t.const_data_ptr(); const float* xk_out = xk_t.const_data_ptr(); std::vector gq = load_golden(xq_golden_path, xq_numel); std::vector gk = load_golden(xk_golden_path, xk_numel); - if (gq.empty() || gk.empty()) { - printf( - "FAIL: could not load goldens %s / %s\n", - xq_golden_path.c_str(), - xk_golden_path.c_str()); - return false; - } + ASSERT_FALSE(gq.empty() || gk.empty()) + << "could not load goldens " << xq_golden_path << " / " << xk_golden_path; - // Per-element abs-OR-rel on xq and xk (shared helper, defined above). + // Per-element abs-OR-rel on xq and xk (shared helper). float maq = 0.0f, mrq = 0.0f, mak = 0.0f, mrk = 0.0f; const bool pass_q = quant_within_tol(xq_out, gq.data(), xq_numel, 1e-3f, 1e-3f, &maq, &mrq); @@ -536,15 +472,11 @@ static bool test_rope( max_abs_err, max_rel_err, xq_numel + xk_numel); - if (!(pass_q && pass_k)) { - printf("FAIL: apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)\n"); - return false; - } - printf("PASS: apply_rotary_emb test\n"); - return true; + EXPECT_TRUE(pass_q && pass_k) + << "apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)"; } -static bool test_prepack( +void test_prepack( const std::string& model_path, const std::string& golden_path, const std::string& label = "x + const w") { @@ -555,17 +487,12 @@ static bool test_prepack( Module module(model_path); auto err = module.load_forward(); - if (err != Error::Ok) { - printf("FAIL: could not load forward method (error %d)\n", (int)err); - return false; - } + ASSERT_EQ(err, Error::Ok) + << "could not load forward method (error " << (int)err << ")"; printf("Model loaded: %s\n", model_path.c_str()); std::vector golden = load_golden(golden_path, numel); - if (golden.empty()) { - printf("FAIL: could not load golden %s\n", golden_path.c_str()); - return false; - } + ASSERT_FALSE(golden.empty()) << "could not load golden " << golden_path; // ((i % 13) - 6) / 16: exact in fp32, matches test_prepack.py::_inputs. std::vector x_data(numel); @@ -575,23 +502,14 @@ static bool test_prepack( auto x = make_tensor_ptr({n, n}, std::vector(x_data)); auto result = module.forward({EValue(x)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor()) { - printf("FAIL: no tensor output\n"); - return false; - } + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (out_tensor.numel() != numel) { - printf( - "FAIL: output numel %zu != expected %d\n", - (size_t)out_tensor.numel(), - numel); - return false; - } + ASSERT_EQ((int)out_tensor.numel(), numel) + << "output numel " << (size_t)out_tensor.numel() << " != expected " + << numel; const float* out_data = out_tensor.const_data_ptr(); float max_abs_err = 0.0f, max_rel_err = 0.0f; @@ -604,16 +522,11 @@ static bool test_prepack( max_abs_err, max_rel_err, numel); - if (!within) { - printf("FAIL: prepack exceeds tolerance 1e-3\n"); - return false; - } - printf("PASS: prepack test\n"); - return true; + EXPECT_TRUE(within) << "prepack exceeds tolerance 1e-3"; } // Reconstruct _ramp_input bit-for-bit, run the op, compare to the fp64 golden. -static bool test_q4gsw_config( +void test_q4gsw_config( const Q4gswConfig& cfg, const std::string& pte, const std::string& golden_path) { @@ -625,10 +538,7 @@ static bool test_q4gsw_config( cfg.n); Module module(pte); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s\n", pte.c_str()); - return false; - } + ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << pte; const int in_numel = cfg.m * cfg.k; const int out_numel = cfg.m * cfg.n; @@ -639,30 +549,18 @@ static bool test_q4gsw_config( auto x = make_tensor_ptr({cfg.m, cfg.k}, std::vector(input)); auto result = module.forward({EValue(x)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); - if (outputs.empty() || !outputs[0].isTensor()) { - printf("FAIL: no tensor output\n"); - return false; - } + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - if (out_tensor.numel() != out_numel) { - printf( - "FAIL: output numel %zu != expected %d\n", - (size_t)out_tensor.numel(), - out_numel); - return false; - } + ASSERT_EQ((int)out_tensor.numel(), out_numel) + << "output numel " << (size_t)out_tensor.numel() << " != expected " + << out_numel; const float* out_data = out_tensor.const_data_ptr(); std::vector golden = load_golden(golden_path, out_numel); - if (golden.empty()) { - printf("FAIL: could not load golden %s\n", golden_path.c_str()); - return false; - } + ASSERT_FALSE(golden.empty()) << "could not load golden " << golden_path; float ma = 0.0f, mr = 0.0f; const bool pass = quant_within_tol( @@ -672,47 +570,8 @@ static bool test_q4gsw_config( ma, mr, out_numel); - if (!pass) { - printf( - "FAIL: linear_q4gsw %s exceeds tolerance (abs %g OR rel %g)\n", - cfg.name, - cfg.tol_abs, - cfg.tol_rel); - return false; - } - printf("PASS: linear_q4gsw %s\n", cfg.name); - return true; -} - -// q4gsw sweep: self-discover q4gsw_.pte; required=FAIL, heavy=gate, *ran. -static bool test_q4gsw_sweep(const std::string& dir, bool* ran) { - bool ok = true; - const bool heavy_run = std::getenv("WEBGPU_TEST_HEAVY") != nullptr; - for (const auto& cfg : kQ4gswConfigs) { - const std::string pte = dir + "q4gsw_" + cfg.name + ".pte"; - FILE* f = std::fopen(pte.c_str(), "rb"); - if (!f) { - if (cfg.required && !dir.empty()) { - printf( - "FAIL: required q4gsw config %s has no .pte in %s\n", - cfg.name, - dir.c_str()); - ok = false; - } - continue; - } - std::fclose(f); - if (cfg.heavy && !heavy_run) { - printf( - "SKIP: heavy q4gsw config %s (set WEBGPU_TEST_HEAVY=1 on a real GPU)\n", - cfg.name); - continue; - } - const std::string golden = dir + "q4gsw_" + cfg.name + ".golden.bin"; - *ran = true; - ok = test_q4gsw_config(cfg, pte, golden) && ok; - } - return ok; + EXPECT_TRUE(pass) << "linear_q4gsw " << cfg.name << " exceeds tolerance (abs " + << cfg.tol_abs << " OR rel " << cfg.tol_rel << ")"; } // Fused sdpa_with_kv_cache sweep config. Mirrors the Python CONFIGS table in @@ -730,7 +589,7 @@ struct SdpaConfig { bool expect_reject = false; // load MUST fail (e.g. D%4 guard), no golden }; -static const SdpaConfig kSdpaConfigs[] = { +const SdpaConfig kSdpaConfigs[] = { // name Hq Hkv D S Cmax pos denom {"gqa31_prefill", 6, 2, 8, 4, 16, 0, 16.0f}, // GQA 3:1 (original case) {"mha_ctxodd", 4, 4, 16, 3, 8, 0, 16.0f}, // MHA; context_len=3 (odd) @@ -777,14 +636,18 @@ constexpr float kSdpaRampDenom = 16.0f; // /denom ramp: ((i % mod) - off) / denom, exact in fp32 (power-of-two denom). // Mirrors test_sdpa.py::_ramp. -static float sdpa_ramp(int i, int mod, int off, float denom = kSdpaRampDenom) { +float sdpa_ramp(int i, int mod, int off, float denom = kSdpaRampDenom) { return static_cast((i % mod) - off) / denom; } // Step-indexed ramp; mirrors test_sdpa.py::_ramp_t bit-for-bit. denom defaults // to kSdpaRampDenom and must match the Python denom for bit-identity. -static float -sdpa_ramp_t(int i, int mod, int off, int t, float denom = kSdpaRampDenom) { +float sdpa_ramp_t( + int i, + int mod, + int off, + int t, + float denom = kSdpaRampDenom) { return static_cast(((i + 31 * t) % mod) - off) / denom; } @@ -800,13 +663,13 @@ struct SdpaSequence { std::vector seq_lens; }; -static const SdpaSequence kSdpaSequences[] = { +const SdpaSequence kSdpaSequences[] = { {"small", 8, 4, 4, 16, {3, 1, 1, 5, 1, 1, 2}}, {"small_d", 6, 2, 8, 16, {3, 1, 1, 5, 1, 1}}, {"llama3", 24, 8, 128, 256, {111, 1, 1, 1, 57, 1, 1}}, }; -static bool test_sdpa_config( +void test_sdpa_config( const SdpaConfig& cfg, const std::string& model_path, const std::string& golden_path) { @@ -825,17 +688,13 @@ static bool test_sdpa_config( auto err = module.load_forward(); if (cfg.expect_reject) { // D not a multiple of 4 must be rejected at load by the head_dim guard. - if (err != Error::Ok) { - printf("PASS: %s rejected at load (error %d)\n", cfg.name, (int)err); - return true; - } - printf("FAIL: %s loaded OK; head_dim%%4 guard did not fire\n", cfg.name); - return false; - } - if (err != Error::Ok) { - printf("FAIL: could not load forward method (error %d)\n", (int)err); - return false; + ASSERT_NE(err, Error::Ok) + << cfg.name << " loaded OK; head_dim%4 guard did not fire"; + printf("PASS: %s rejected at load (error %d)\n", cfg.name, (int)err); + return; } + ASSERT_EQ(err, Error::Ok) + << "could not load forward method (error " << (int)err << ")"; printf("Model loaded: %s\n", model_path.c_str()); const int qn = cfg.s * cfg.hq * cfg.d; @@ -869,10 +728,8 @@ static bool test_sdpa_config( auto result = module.forward( {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); - if (!result.ok()) { - printf("FAIL: forward failed (error %d)\n", (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; const auto& outputs = result.get(); // Select the attention output [1,S,Hq,D] by shape; the op returns @@ -891,32 +748,17 @@ static bool test_sdpa_config( attn_matches++; } } - if (attn_idx < 0) { - printf( - "FAIL: no attention output [1,%d,%d,%d] among %zu outputs\n", - cfg.s, - cfg.hq, - cfg.d, - outputs.size()); - return false; - } - if (attn_matches > 1) { - printf( - "FAIL: ambiguous attention output: %d tensors match shape [1,%d,%d,%d]\n", - attn_matches, - cfg.s, - cfg.hq, - cfg.d); - return false; - } + ASSERT_GE(attn_idx, 0) << "no attention output [1," << cfg.s << "," << cfg.hq + << "," << cfg.d << "] among " << outputs.size() + << " outputs"; + ASSERT_LE(attn_matches, 1) << "ambiguous attention output: " << attn_matches + << " tensors match shape [1," << cfg.s << "," + << cfg.hq << "," << cfg.d << "]"; const auto& out_tensor = outputs[attn_idx].toTensor(); const float* out_data = out_tensor.const_data_ptr(); std::vector golden = load_golden(golden_path, on); - if (golden.empty()) { - printf("FAIL: could not load golden %s\n", golden_path.c_str()); - return false; - } + ASSERT_FALSE(golden.empty()) << "could not load golden " << golden_path; float max_abs_err = 0.0f, max_rel_err = 0.0f; const bool pass = @@ -926,48 +768,14 @@ static bool test_sdpa_config( max_abs_err, max_rel_err, on); - if (!pass) { - printf( - "FAIL: %s exceeds tolerance (per-element abs 1e-4 OR rel 1e-3)\n", - cfg.name); - return false; - } - printf("PASS: sdpa test (%s)\n", cfg.name); - return true; -} - -// Run the full SDPA sweep. Each config self-discovers its embedded/on-disk -// sdpa_.pte; a config is skipped silently when its .pte is absent, so the -// same binary works whether one or all configs are embedded. Returns false only -// if a discovered config actually fails. Sets *ran true if any config ran. -static bool test_sdpa_sweep(const std::string& dir, bool* ran) { - bool ok = true; - for (const auto& cfg : kSdpaConfigs) { - const std::string pte = dir + "sdpa_" + cfg.name + ".pte"; - FILE* f = std::fopen(pte.c_str(), "rb"); - if (!f) { - // required config absent (dir set) = FAIL; otherwise skip silently. - if (cfg.required && !dir.empty()) { - printf( - "FAIL: required sdpa config %s has no .pte in %s\n", - cfg.name, - dir.c_str()); - ok = false; - } - continue; // not embedded in this binary - } - std::fclose(f); - const std::string golden = dir + "sdpa_" + cfg.name + ".golden.bin"; - *ran = true; - ok = test_sdpa_config(cfg, pte, golden) && ok; - } - return ok; + EXPECT_TRUE(pass) << cfg.name + << " exceeds tolerance (per-element abs 1e-4 OR rel 1e-3)"; } // Replay one sequence: thread the op's returned (mutated) KV cache across // steps, comparing each step's attention output to its accumulated-context // golden. -static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { +void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { printf( "\n--- Test: sdpa replay (%s: Hq=%d,Hkv=%d,D=%d,Cmax=%d, %zu steps) ---\n", seq.name, @@ -982,7 +790,6 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { int input_pos = 0; int k_idx = -1, v_idx = -1; // pinned at step 0 by content (caches share numel) - bool ok = true; for (size_t t = 0; t < seq.seq_lens.size(); t++) { const int s = seq.seq_lens[t]; @@ -990,10 +797,8 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { std::to_string(t) + "_S" + std::to_string(s) + "_pos" + std::to_string(input_pos); Module module(base + ".pte"); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s.pte\n", base.c_str()); - return false; - } + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load " << base << ".pte"; const int qn = s * seq.hq * seq.d; const int kvn = s * seq.hkv * seq.d; @@ -1016,13 +821,8 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { auto result = module.forward( {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); - if (!result.ok()) { - printf( - "FAIL: forward %s.pte (error %d)\n", - base.c_str(), - (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) + << "forward " << base << ".pte (error " << (int)result.error() << ")"; const auto& outs = result.get(); // The op returns [k_cache, v_cache, attn_output]: attn has a unique numel; @@ -1040,10 +840,8 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { cache_idxs.push_back(static_cast(i)); } } - if (attn_idx < 0 || cache_idxs.size() != 2) { - printf("FAIL: %s step%zu: expected 1 attn + 2 caches\n", seq.name, t); - return false; - } + ASSERT_TRUE(attn_idx >= 0 && cache_idxs.size() == 2) + << seq.name << " step" << t << ": expected 1 attn + 2 caches"; if (t == 0) { const float* c0 = outs[cache_idxs[0]].toTensor().const_data_ptr(); @@ -1063,18 +861,13 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { k_idx = cache_idxs[1]; v_idx = cache_idxs[0]; } else { - printf( - "FAIL: %s step0 cannot identify k/v cache by content\n", seq.name); - return false; + FAIL() << seq.name << " step0 cannot identify k/v cache by content"; } printf(" k/v cache outputs: k_idx=%d v_idx=%d\n", k_idx, v_idx); } std::vector golden = load_golden(base + ".golden.bin", qn); - if (golden.empty()) { - printf("FAIL: could not load %s.golden.bin\n", base.c_str()); - return false; - } + ASSERT_FALSE(golden.empty()) << "could not load " << base << ".golden.bin"; const float* ad = outs[attn_idx].toTensor().const_data_ptr(); float ma = 0.0f, mr = 0.0f; const bool step_ok = sdpa_within_tol(ad, golden.data(), qn, &ma, &mr); @@ -1086,13 +879,9 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { input_pos + s, ma, mr); - if (!step_ok) { - printf( - "FAIL: %s step%zu exceeds tolerance (per-element abs 1e-4 OR rel 1e-3)\n", - seq.name, - t); - ok = false; - } + EXPECT_TRUE(step_ok) + << seq.name << " step" << t + << " exceeds tolerance (per-element abs 1e-4 OR rel 1e-3)"; // Thread the device-written caches into the next step (K->K, V->V). const float* kd = outs[k_idx].toTensor().const_data_ptr(); @@ -1101,28 +890,6 @@ static bool test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { vc.assign(vd, vd + cn); input_pos += s; } - - if (ok) { - printf("PASS: sdpa replay (%s)\n", seq.name); - } - return ok; -} - -// Run all replay sequences whose step0 .pte is present (self-skip otherwise). -static bool test_sdpa_replay_sweep(const std::string& dir, bool* ran) { - bool ok = true; - for (const auto& seq : kSdpaSequences) { - const std::string step0 = dir + "sdpa_" + seq.name + "_step0_S" + - std::to_string(seq.seq_lens[0]) + "_pos0.pte"; - FILE* f = std::fopen(step0.c_str(), "rb"); - if (!f) { - continue; // sequence not embedded in this binary - } - std::fclose(f); - *ran = true; - ok = test_sdpa_replay(seq, dir) && ok; - } - return ok; } // Dynamic input_pos decode: ONE .pte (S=1, runtime SymInt input_pos) reused @@ -1133,7 +900,7 @@ static bool test_sdpa_replay_sweep(const std::string& dir, bool* ran) { // per-step input_pos actually being read + applied. negative=true pins // input_pos at 0 every step (stale context_len) and asserts the run DIVERGES, // proving the runtime input_pos + resize hook are load-bearing (no false-pass). -static bool test_sdpa_dynamic_decode( +void test_sdpa_dynamic_decode( const SdpaSequence& seq, const std::string& dir, bool negative) { @@ -1150,16 +917,12 @@ static bool test_sdpa_dynamic_decode( const std::string pte = dir + "sdpa_dyn_" + seq.name + ".pte"; Module module(pte); - if (module.load_forward() != Error::Ok) { - printf("FAIL: could not load %s\n", pte.c_str()); - return false; - } + ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << pte; const int cn = seq.cmax * seq.hkv * seq.d; std::vector kc(cn, 0.0f), vc(cn, 0.0f); int k_idx = -1, v_idx = -1; // pinned at step 0 by content (caches share numel) - bool ok = true; bool any_mismatch = false; for (int t = 0; t < kSteps; t++) { @@ -1190,10 +953,8 @@ static bool test_sdpa_dynamic_decode( EValue(kct), EValue(vct), EValue(ipt)}); - if (!result.ok()) { - printf("FAIL: forward step%d (error %d)\n", t, (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) + << "forward step" << t << " (error " << (int)result.error() << ")"; const auto& outs = result.get(); int attn_idx = -1; @@ -1209,10 +970,8 @@ static bool test_sdpa_dynamic_decode( cache_idxs.push_back(static_cast(i)); } } - if (attn_idx < 0 || cache_idxs.size() != 2) { - printf("FAIL: %s step%d: expected 1 attn + 2 caches\n", seq.name, t); - return false; - } + ASSERT_TRUE(attn_idx >= 0 && cache_idxs.size() == 2) + << seq.name << " step" << t << ": expected 1 attn + 2 caches"; if (t == 0) { const float* c0 = outs[cache_idxs[0]].toTensor().const_data_ptr(); const float* c1 = outs[cache_idxs[1]].toTensor().const_data_ptr(); @@ -1231,18 +990,14 @@ static bool test_sdpa_dynamic_decode( k_idx = cache_idxs[1]; v_idx = cache_idxs[0]; } else { - printf("FAIL: %s step0 cannot identify k/v cache\n", seq.name); - return false; + FAIL() << seq.name << " step0 cannot identify k/v cache"; } } const std::string gpath = dir + "sdpa_dyn_" + seq.name + "_step" + std::to_string(t) + ".golden.bin"; std::vector golden = load_golden(gpath, qn); - if (golden.empty()) { - printf("FAIL: could not load %s\n", gpath.c_str()); - return false; - } + ASSERT_FALSE(golden.empty()) << "could not load " << gpath; const float* ad = outs[attn_idx].toTensor().const_data_ptr(); float ma = 0.0f, mr = 0.0f; const bool step_ok = sdpa_within_tol(ad, golden.data(), qn, &ma, &mr); @@ -1265,46 +1020,24 @@ static bool test_sdpa_dynamic_decode( } if (negative) { + // The negative control must DIVERGE: a stale input_pos=0 every step cannot + // match the accumulating golden -- otherwise the oracle has no teeth. + EXPECT_TRUE(any_mismatch) + << seq.name + << " negative control matched the golden (oracle has no teeth)"; if (any_mismatch) { printf( "PASS: sdpa dynamic decode NEGATIVE (%s): stale input_pos diverges " "as expected\n", seq.name); - return true; } - printf( - "FAIL: %s negative control matched the golden (oracle has no teeth)\n", - seq.name); - return false; + return; } - if (any_mismatch) { - printf( - "FAIL: %s exceeds tolerance (per-element abs 1e-4 OR rel 1e-3)\n", - seq.name); - ok = false; - } - if (ok) { + EXPECT_FALSE(any_mismatch) + << seq.name << " exceeds tolerance (per-element abs 1e-4 OR rel 1e-3)"; + if (!any_mismatch) { printf("PASS: sdpa dynamic decode (%s)\n", seq.name); } - return ok; -} - -// Run dynamic decode (positive + negative control) for each param set whose -// sdpa_dyn_.pte is embedded (self-skip otherwise). -static bool test_sdpa_dynamic_decode_sweep(const std::string& dir, bool* ran) { - bool ok = true; - for (const auto& seq : kSdpaSequences) { - const std::string pte = dir + "sdpa_dyn_" + seq.name + ".pte"; - FILE* f = std::fopen(pte.c_str(), "rb"); - if (!f) { - continue; - } - std::fclose(f); - *ran = true; - ok = test_sdpa_dynamic_decode(seq, dir, /*negative=*/false) && ok; - ok = test_sdpa_dynamic_decode(seq, dir, /*negative=*/true) && ok; - } - return ok; } // In-graph mutable KV cache: ONE .pte whose k_cache/v_cache are mutable buffers @@ -1314,7 +1047,7 @@ static bool test_sdpa_dynamic_decode_sweep(const std::string& dir, bool* ran) { // Module each step re-seeds the cache to zeros, so it MUST diverge from the // accumulating golden at step>=1. Persistent-matches + fresh-diverges = proof // the pass comes from real accumulation, not a static artifact. -static bool test_sdpa_incache_decode( +void test_sdpa_incache_decode( const SdpaSequence& seq, const std::string& dir, bool fresh_per_step) { @@ -1333,10 +1066,8 @@ static bool test_sdpa_incache_decode( std::unique_ptr persistent; if (!fresh_per_step) { persistent = std::make_unique(pte); - if (persistent->load_forward() != Error::Ok) { - printf("FAIL: could not load %s\n", pte.c_str()); - return false; - } + ASSERT_EQ(persistent->load_forward(), Error::Ok) + << "could not load " << pte; } bool any_mismatch = false; @@ -1363,10 +1094,7 @@ static bool test_sdpa_incache_decode( Module* mod = persistent.get(); if (fresh_per_step) { fresh = std::make_unique(pte); - if (fresh->load_forward() != Error::Ok) { - printf("FAIL: could not load %s\n", pte.c_str()); - return false; - } + ASSERT_EQ(fresh->load_forward(), Error::Ok) << "could not load " << pte; mod = fresh.get(); } @@ -1374,10 +1102,8 @@ static bool test_sdpa_incache_decode( // buffers). auto result = mod->forward({EValue(qt), EValue(kt), EValue(vt), EValue(ipt)}); - if (!result.ok()) { - printf("FAIL: forward step%d (error %d)\n", t, (int)result.error()); - return false; - } + ASSERT_TRUE(result.ok()) + << "forward step" << t << " (error " << (int)result.error() << ")"; const auto& outs = result.get(); int attn_idx = -1; for (size_t i = 0; i < outs.size(); i++) { @@ -1387,18 +1113,13 @@ static bool test_sdpa_incache_decode( break; } } - if (attn_idx < 0) { - printf("FAIL: %s step%d: no attn output (numel %d)\n", seq.name, t, qn); - return false; - } + ASSERT_GE(attn_idx, 0) << seq.name << " step" << t + << ": no attn output (numel " << qn << ")"; const std::string gpath = dir + "sdpa_incache_" + seq.name + "_step" + std::to_string(t) + ".golden.bin"; std::vector golden = load_golden(gpath, qn); - if (golden.empty()) { - printf("FAIL: could not load %s\n", gpath.c_str()); - return false; - } + ASSERT_FALSE(golden.empty()) << "could not load " << gpath; const float* ad = outs[attn_idx].toTensor().const_data_ptr(); float ma = 0.0f, mr = 0.0f; const bool step_ok = sdpa_within_tol(ad, golden.data(), qn, &ma, &mr); @@ -1418,55 +1139,36 @@ static bool test_sdpa_incache_decode( if (fresh_per_step) { // The control must DIVERGE: a fresh Module per step has no accumulated // history, so it cannot match the accumulating golden at step>=1. + EXPECT_TRUE(any_mismatch) + << seq.name + << " static control matched the accumulating golden -- " + "accumulation was not actually exercised (false-pass risk)"; if (any_mismatch) { printf( "PASS: in-graph-cache STATIC CONTROL (%s) diverges as expected -- " "persistence is load-bearing; the positive pass is real accumulation\n", seq.name); - return true; } - printf( - "FAIL: %s static control matched the accumulating golden -- " - "accumulation was not actually exercised (false-pass risk)\n", - seq.name); - return false; + return; } + EXPECT_FALSE(any_mismatch) + << seq.name << " in-graph-cache decode exceeds tolerance"; if (!any_mismatch) { printf( "PASS: sdpa in-graph-cache decode (%s) -- cache accumulated in-graph " "with NO host threading\n", seq.name); - return true; - } - printf("FAIL: %s in-graph-cache decode exceeds tolerance\n", seq.name); - return false; -} - -static bool test_sdpa_incache_decode_sweep(const std::string& dir, bool* ran) { - bool ok = true; - for (const auto& seq : kSdpaSequences) { - const std::string pte = dir + "sdpa_incache_" + seq.name + ".pte"; - FILE* f = std::fopen(pte.c_str(), "rb"); - if (!f) { - continue; - } - std::fclose(f); - *ran = true; - ok = test_sdpa_incache_decode(seq, dir, /*fresh_per_step=*/false) && ok; - ok = test_sdpa_incache_decode(seq, dir, /*fresh_per_step=*/true) && ok; } - return ok; } // S1 SymInt round-trip: build a graph directly from a dynamic-input_pos SDPA // blob; confirm input_pos deserializes as a live SymInt and set/read // round-trips. -static bool test_symint_roundtrip(const std::string& blob_path) { +void test_symint_roundtrip(const std::string& blob_path) { printf("\n--- Test: symint round-trip (%s) ---\n", blob_path.c_str()); FILE* f = std::fopen(blob_path.c_str(), "rb"); if (!f) { - printf("SKIP: %s not present\n", blob_path.c_str()); - return true; + GTEST_SKIP() << blob_path << " not present"; } std::fseek(f, 0, SEEK_END); long n = std::ftell(f); @@ -1474,24 +1176,17 @@ static bool test_symint_roundtrip(const std::string& blob_path) { std::vector blob(static_cast(n)); size_t rd = std::fread(blob.data(), 1, blob.size(), f); std::fclose(f); - if (rd != blob.size()) { - printf("FAIL: short read of %s\n", blob_path.c_str()); - return false; - } + ASSERT_EQ(rd, blob.size()) << "short read of " << blob_path; auto header = WebGPUDelegateHeader::parse(blob.data()); - if (!header.ok()) { - printf("FAIL: delegate header parse\n"); - return false; - } + ASSERT_TRUE(header.ok()) << "delegate header parse"; const uint8_t* base = blob.data(); WebGPUGraph graph; try { graph.build( base + header->flatbuffer_offset, base + header->bytes_offset, nullptr); } catch (const std::exception& e) { - printf("FAIL: graph build: %s\n", e.what()); - return false; + FAIL() << "graph build: " << e.what(); } int sid = -1; @@ -1501,35 +1196,21 @@ static bool test_symint_roundtrip(const std::string& blob_path) { break; } } - if (sid < 0) { - printf( - "FAIL: no SymInt value deserialized (input_pos should be a SymInt)\n"); - return false; - } - if (graph.symint_buffer(sid) == nullptr) { - printf("FAIL: SymInt %d has no live uniform buffer\n", sid); - return false; - } - if (graph.read_symint(sid) != 0) { - printf( - "FAIL: SymInt %d placeholder != 0 (got %d)\n", - sid, - graph.read_symint(sid)); - return false; - } + ASSERT_GE(sid, 0) + << "no SymInt value deserialized (input_pos should be a SymInt)"; + ASSERT_NE(graph.symint_buffer(sid), nullptr) + << "SymInt " << sid << " has no live uniform buffer"; + ASSERT_EQ(graph.read_symint(sid), 0) + << "SymInt " << sid << " placeholder != 0 (got " << graph.read_symint(sid) + << ")"; graph.set_symint(sid, 7); - if (graph.read_symint(sid) != 7) { - printf("FAIL: set/read round-trip (got %d)\n", graph.read_symint(sid)); - return false; - } + ASSERT_EQ(graph.read_symint(sid), 7) + << "set/read round-trip (got " << graph.read_symint(sid) << ")"; // Execute-read: feed a fake input_pos=5 via the recorded select_as_symint // source and confirm update_symints_from_inputs populates the SymInt. const auto& srcs = graph.symint_sources(); - if (srcs.empty()) { - printf("FAIL: no select_as_symint source recorded\n"); - return false; - } + ASSERT_FALSE(srcs.empty()) << "no select_as_symint source recorded"; const auto& in_ids = graph.input_ids(); std::vector fake_inputs(in_ids.size()); int64_t fake_pos = 5; @@ -1539,29 +1220,24 @@ static bool test_symint_roundtrip(const std::string& blob_path) { } } graph.update_symints_from_inputs(fake_inputs); - if (graph.read_symint(srcs[0].symint_id) != 5) { - printf( - "FAIL: execute-read (got %d, want 5)\n", - graph.read_symint(srcs[0].symint_id)); - return false; - } + ASSERT_EQ(graph.read_symint(srcs[0].symint_id), 5) + << "execute-read (got " << graph.read_symint(srcs[0].symint_id) + << ", want 5)"; printf( "PASS: symint round-trip (SymInt %d: deserialize, live buffer, " "set 0->7, execute-read input_pos->5)\n", sid); - return true; } // Group 1: the resize-hook dirty-gating mechanism (no SDPA dependency). // A hook keyed to a SymInt must run via propagate_resize() iff that SymInt // changed since the last propagate_resize, and exactly once per change. -static bool test_resize_hook(const std::string& blob_path) { +void test_resize_hook(const std::string& blob_path) { printf("\n--- Test: resize-hook dirty-gating (%s) ---\n", blob_path.c_str()); FILE* f = std::fopen(blob_path.c_str(), "rb"); if (!f) { - printf("SKIP: %s not present\n", blob_path.c_str()); - return true; + GTEST_SKIP() << blob_path << " not present"; } std::fseek(f, 0, SEEK_END); long n = std::ftell(f); @@ -1569,23 +1245,16 @@ static bool test_resize_hook(const std::string& blob_path) { std::vector blob(static_cast(n)); size_t rd = std::fread(blob.data(), 1, blob.size(), f); std::fclose(f); - if (rd != blob.size()) { - printf("FAIL: short read of %s\n", blob_path.c_str()); - return false; - } + ASSERT_EQ(rd, blob.size()) << "short read of " << blob_path; auto header = WebGPUDelegateHeader::parse(blob.data()); - if (!header.ok()) { - printf("FAIL: delegate header parse\n"); - return false; - } + ASSERT_TRUE(header.ok()) << "delegate header parse"; const uint8_t* base = blob.data(); WebGPUGraph graph; try { graph.build( base + header->flatbuffer_offset, base + header->bytes_offset, nullptr); } catch (const std::exception& e) { - printf("FAIL: graph build: %s\n", e.what()); - return false; + FAIL() << "graph build: " << e.what(); } int sid = -1; @@ -1595,10 +1264,7 @@ static bool test_resize_hook(const std::string& blob_path) { break; } } - if (sid < 0) { - printf("FAIL: no SymInt value deserialized\n"); - return false; - } + ASSERT_GE(sid, 0) << "no SymInt value deserialized"; int run_count = 0; int last_seen = -1; @@ -1610,279 +1276,405 @@ static bool test_resize_hook(const std::string& blob_path) { // 1: change 0->3 then propagate -> hook runs once, sees 3. graph.set_symint(sid, 3); graph.propagate_resize(); - if (run_count != 1 || last_seen != 3) { - printf( - "FAIL: after set(3)+propagate run_count=%d last_seen=%d (want 1,3)\n", - run_count, - last_seen); - return false; - } + ASSERT_TRUE(run_count == 1 && last_seen == 3) + << "after set(3)+propagate run_count=" << run_count + << " last_seen=" << last_seen << " (want 1,3)"; // 2: propagate again with no change -> hook does NOT run. graph.propagate_resize(); - if (run_count != 1) { - printf( - "FAIL: propagate with clean dirty-set ran the hook (run_count=%d)\n", - run_count); - return false; - } + ASSERT_EQ(run_count, 1) + << "propagate with clean dirty-set ran the hook (run_count=" << run_count + << ")"; // 3: set to the SAME value -> not dirty -> hook does NOT run. graph.set_symint(sid, 3); graph.propagate_resize(); - if (run_count != 1) { - printf( - "FAIL: set(same)+propagate ran the hook (run_count=%d)\n", run_count); - return false; - } + ASSERT_EQ(run_count, 1) << "set(same)+propagate ran the hook (run_count=" + << run_count << ")"; // 4: change 3->8 then propagate -> hook runs again, sees 8. graph.set_symint(sid, 8); graph.propagate_resize(); - if (run_count != 2 || last_seen != 8) { - printf( - "FAIL: after set(8)+propagate run_count=%d last_seen=%d (want 2,8)\n", - run_count, - last_seen); - return false; - } + ASSERT_TRUE(run_count == 2 && last_seen == 8) + << "after set(8)+propagate run_count=" << run_count + << " last_seen=" << last_seen << " (want 2,8)"; printf( "PASS: resize-hook dirty-gating (SymInt %d: runs only on change, " "once per change; saw 3 then 8)\n", sid); - return true; } -int main(int argc, char** argv) { - std::string update_cache_model_path; - if (const char* env = std::getenv("WEBGPU_TEST_UPDATE_CACHE_MODEL")) { - update_cache_model_path = env; - } +// q4gsw embedding_q4gsw on-GPU configs: small + llama1b (env-gated, +// run-if-present). +struct EmbConfig { + const char* name; + const char* model_env; + const char* indices_env; + const char* golden_env; + int num_indices; + int embed; +}; +const EmbConfig kEmbConfigs[] = { + {"small", + "WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL", + "WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES", + "WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN", + 4, + 64}, + {"llama1b", + "WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL", + "WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES", + "WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN", + 4, + 2048}, +}; - // Quantized-linear sweep dir (mirrors WEBGPU_TEST_SDPA_DIR). - std::string qlinear_dir; - if (const char* env = std::getenv("WEBGPU_TEST_QUANTIZED_LINEAR_DIR")) { - qlinear_dir = env; - if (!qlinear_dir.empty() && qlinear_dir.back() != '/') { - qlinear_dir += '/'; - } - } +// apply_rotary_emb on-GPU configs: multi + decode (env-gated, run-if-present). +struct RopeConfig { + const char* name; + const char* model_env; + const char* xq_env; + const char* xk_env; + int S; + int NH; + int NKV; + int HD; +}; +const RopeConfig kRopeConfigs[] = { + {"multi", + "WEBGPU_TEST_ROPE_MODEL", + "WEBGPU_TEST_ROPE_XQ_GOLDEN", + "WEBGPU_TEST_ROPE_XK_GOLDEN", + 5, + 8, + 2, + 64}, + {"decode", + "WEBGPU_TEST_ROPE_DECODE_MODEL", + "WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN", + "WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN", + 1, + 32, + 8, + 64}, +}; - // embedding_q4gsw on-GPU configs: small + llama1b (env-gated, - // run-if-present). - struct EmbConfig { - const char* name; - const char* model_env; - const char* indices_env; - const char* golden_env; - int num_indices; - int embed; - }; - const EmbConfig emb_configs[] = { - {"small", - "WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL", - "WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES", - "WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN", - 4, - 64}, - {"llama1b", - "WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL", - "WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES", - "WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN", - 4, - 2048}, - }; +} // namespace - // apply_rotary_emb on-GPU configs: multi + decode (env-gated, - // run-if-present). - struct RopeConfig { - const char* name; - const char* model_env; - const char* xq_env; - const char* xk_env; - int S; - int NH; - int NKV; - int HD; - }; - const RopeConfig rope_configs[] = { - {"multi", - "WEBGPU_TEST_ROPE_MODEL", - "WEBGPU_TEST_ROPE_XQ_GOLDEN", - "WEBGPU_TEST_ROPE_XK_GOLDEN", - 5, - 8, - 2, - 64}, - {"decode", - "WEBGPU_TEST_ROPE_DECODE_MODEL", - "WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN", - "WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN", - 1, - 32, - 8, - 64}, - }; +#ifdef WGPU_BACKEND_ENABLE_PROFILING +TEST(WebGPUNative, QueryPoolOverrunThrows) { + test_query_pool_overrun_throws(); +} - std::string prepack_model_path, prepack_golden_path; - if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_MODEL")) { - prepack_model_path = env; - } - if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_GOLDEN")) { - prepack_golden_path = env; +TEST(WebGPUNative, QueryPoolRoundtrip) { + test_query_pool_roundtrip(*get_default_webgpu_context()); +} +#endif // WGPU_BACKEND_ENABLE_PROFILING + +TEST(WebGPUNative, UpdateCache) { + if (g_update_cache_model_path.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_UPDATE_CACHE_MODEL not set"; } + test_update_cache(g_update_cache_model_path); +} - std::string prepack2_model_path, prepack2_golden_path; - if (const char* env = std::getenv("WEBGPU_TEST_PREPACK2_MODEL")) { - prepack2_model_path = env; +// Guard python<->C++ ramp bit-identity: q4gsw_ramp(0) = -0.5 exactly. +TEST(WebGPUNative, Q4gswRampBitIdentity) { + EXPECT_LT(std::abs(q4gsw_ramp(0) - (-0.5f)), 1e-12f) + << "q4gsw_ramp bit-identity check"; +} + +// q4gsw sweep: self-discover q4gsw_.pte; required=FAIL, heavy=gate. +TEST(WebGPUNative, QuantizedLinearSweep) { + const std::string& dir = g_qlinear_dir; + const bool heavy_run = std::getenv("WEBGPU_TEST_HEAVY") != nullptr; + bool ran = false; + for (const auto& cfg : kQ4gswConfigs) { + const std::string pte = dir + "q4gsw_" + cfg.name + ".pte"; + FILE* f = std::fopen(pte.c_str(), "rb"); + if (!f) { + if (cfg.required && !dir.empty()) { + ADD_FAILURE() << "required q4gsw config " << cfg.name + << " has no .pte in " << dir; + } + continue; + } + std::fclose(f); + if (cfg.heavy && !heavy_run) { + printf( + "SKIP: heavy q4gsw config %s (set WEBGPU_TEST_HEAVY=1 on a real GPU)\n", + cfg.name); + continue; + } + const std::string golden = dir + "q4gsw_" + cfg.name + ".golden.bin"; + ran = true; + test_q4gsw_config(cfg, pte, golden); } - if (const char* env = std::getenv("WEBGPU_TEST_PREPACK2_GOLDEN")) { - prepack2_golden_path = env; + if (!dir.empty() && !ran) { + ADD_FAILURE() + << "WEBGPU_TEST_QUANTIZED_LINEAR_DIR set but no q4gsw config ran"; } +} - std::string prepack_tied_model_path, prepack_tied_golden_path; - if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_TIED_MODEL")) { - prepack_tied_model_path = env; +TEST(WebGPUNative, EmbeddingQ4gsw) { + bool any = false; + for (const auto& c : kEmbConfigs) { + const char* m = std::getenv(c.model_env); + const char* ip = std::getenv(c.indices_env); + const char* g = std::getenv(c.golden_env); + if (m && ip && g && *m && *ip && *g) { + any = true; + test_embedding_q4gsw(m, ip, g, c.num_indices, c.embed, c.name); + } } - if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_TIED_GOLDEN")) { - prepack_tied_golden_path = env; + if (!any) { + GTEST_SKIP() << "no embedding_q4gsw config env set"; } +} - // SDPA sweep: configs self-discover their sdpa_.pte/.golden.bin under - // this directory (default "" = the embedded-file root / cwd). Set - // WEBGPU_TEST_SDPA_DIR to point at the exported .pte directory (e.g. /tmp/). - std::string sdpa_dir; - if (const char* env = std::getenv("WEBGPU_TEST_SDPA_DIR")) { - sdpa_dir = env; - if (!sdpa_dir.empty() && sdpa_dir.back() != '/') { - sdpa_dir += '/'; +TEST(WebGPUNative, Rope) { + bool any = false; + for (const auto& c : kRopeConfigs) { + const char* m = std::getenv(c.model_env); + const char* xq = std::getenv(c.xq_env); + const char* xk = std::getenv(c.xk_env); + if (m && xq && xk && *m && *xq && *xk) { + any = true; + test_rope(m, xq, xk, c.S, c.NH, c.NKV, c.HD, c.name); } } - - WebGPUContext ctx; - try { - ctx = create_webgpu_context(); - } catch (const std::exception& e) { - printf("SKIP: %s\n", e.what()); - return 0; + if (!any) { + GTEST_SKIP() << "no apply_rotary_emb config env set"; } +} - set_default_webgpu_context(&ctx); - printf("WebGPU device acquired (native)\n"); +TEST(WebGPUNative, Prepack) { + if (g_prepack_model_path.empty() || g_prepack_golden_path.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_PREPACK_MODEL/GOLDEN not set"; + } + test_prepack(g_prepack_model_path, g_prepack_golden_path); +} - bool ok = true; -#ifdef WGPU_BACKEND_ENABLE_PROFILING - ok = test_query_pool_overrun_throws() && ok; - ok = test_query_pool_roundtrip(ctx) && ok; -#endif // WGPU_BACKEND_ENABLE_PROFILING - if (!update_cache_model_path.empty()) { - ok = test_update_cache(update_cache_model_path) && ok; +TEST(WebGPUNative, Prepack2) { + if (g_prepack2_model_path.empty() || g_prepack2_golden_path.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_PREPACK2_MODEL/GOLDEN not set"; } + test_prepack(g_prepack2_model_path, g_prepack2_golden_path, "x + w1 + w2"); +} - bool q4gsw_ran = false; - bool q4gsw_ok = test_q4gsw_sweep(qlinear_dir, &q4gsw_ran); - if (q4gsw_ran) { - ok = q4gsw_ok && ok; +TEST(WebGPUNative, PrepackTied) { + if (g_prepack_tied_model_path.empty() || g_prepack_tied_golden_path.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_PREPACK_TIED_MODEL/GOLDEN not set"; } - // Guard python<->C++ ramp bit-identity: q4gsw_ramp(0) = -0.5 exactly. - if (std::abs(q4gsw_ramp(0) - (-0.5f)) > 1e-12f) { - printf("FAIL: q4gsw_ramp bit-identity check\n"); - ok = false; + test_prepack( + g_prepack_tied_model_path, + g_prepack_tied_golden_path, + "x + w + w (tied weights, shared key)"); +} + +// SDPA sweep: configs self-discover sdpa_.pte; required=FAIL else skip. +TEST(WebGPUNative, SdpaSweep) { + const std::string& dir = g_sdpa_dir; + bool ran = false; + for (const auto& cfg : kSdpaConfigs) { + const std::string pte = dir + "sdpa_" + cfg.name + ".pte"; + FILE* f = std::fopen(pte.c_str(), "rb"); + if (!f) { + // required config absent (dir set) = FAIL; otherwise skip silently. + if (cfg.required && !dir.empty()) { + ADD_FAILURE() << "required sdpa config " << cfg.name + << " has no .pte in " << dir; + } + continue; // not embedded in this binary + } + std::fclose(f); + const std::string golden = dir + "sdpa_" + cfg.name + ".golden.bin"; + ran = true; + test_sdpa_config(cfg, pte, golden); } - if (!qlinear_dir.empty() && !q4gsw_ran) { - printf( - "FAIL: WEBGPU_TEST_QUANTIZED_LINEAR_DIR set but no q4gsw config ran\n"); - ok = false; + if (!dir.empty() && !ran) { + ADD_FAILURE() << "WEBGPU_TEST_SDPA_DIR set but no sdpa config found a .pte"; } +} - for (const auto& c : emb_configs) { - const char* m = std::getenv(c.model_env); - const char* ip = std::getenv(c.indices_env); - const char* g = std::getenv(c.golden_env); - if (m && ip && g && *m && *ip && *g) { - ok = test_embedding_q4gsw(m, ip, g, c.num_indices, c.embed, c.name) && ok; +// Guard python<->C++ ramp bit-identity (recorded: _ramp_t(0,17,8,2)=0.1875). +TEST(WebGPUNative, SdpaRampTBitIdentity) { + EXPECT_LT(std::abs(sdpa_ramp_t(0, 17, 8, 2) - 0.1875f), 1e-12f) + << "sdpa_ramp_t bit-identity check"; +} + +// Guard the adversarial denom path: sdpa_ramp(0,17,8,0.5)= -16.0 exactly. +TEST(WebGPUNative, SdpaRampDenomBitIdentity) { + EXPECT_LT(std::abs(sdpa_ramp(0, 17, 8, 0.5f) - (-16.0f)), 1e-12f) + << "sdpa_ramp denom bit-identity check"; +} + +// Replay sweep: run every sequence whose step0 .pte is present. +TEST(WebGPUNative, SdpaReplaySweep) { + const std::string& dir = g_sdpa_dir; + for (const auto& seq : kSdpaSequences) { + const std::string step0 = dir + "sdpa_" + seq.name + "_step0_S" + + std::to_string(seq.seq_lens[0]) + "_pos0.pte"; + FILE* f = std::fopen(step0.c_str(), "rb"); + if (!f) { + continue; // sequence not embedded in this binary } + std::fclose(f); + test_sdpa_replay(seq, dir); } +} - for (const auto& c : rope_configs) { - const char* m = std::getenv(c.model_env); - const char* xq = std::getenv(c.xq_env); - const char* xk = std::getenv(c.xk_env); - if (m && xq && xk && *m && *xq && *xk) { - ok = test_rope(m, xq, xk, c.S, c.NH, c.NKV, c.HD, c.name) && ok; +// Dynamic decode sweep: positive + negative control per embedded param set. +TEST(WebGPUNative, SdpaDynamicDecodeSweep) { + const std::string& dir = g_sdpa_dir; + for (const auto& seq : kSdpaSequences) { + const std::string pte = dir + "sdpa_dyn_" + seq.name + ".pte"; + FILE* f = std::fopen(pte.c_str(), "rb"); + if (!f) { + continue; } + std::fclose(f); + test_sdpa_dynamic_decode(seq, dir, /*negative=*/false); + test_sdpa_dynamic_decode(seq, dir, /*negative=*/true); } +} - if (!prepack_model_path.empty() && !prepack_golden_path.empty()) { - ok = test_prepack(prepack_model_path, prepack_golden_path) && ok; +// In-graph-cache decode sweep: persistent + fresh (static control) per set. +TEST(WebGPUNative, SdpaIncacheDecodeSweep) { + const std::string& dir = g_sdpa_dir; + for (const auto& seq : kSdpaSequences) { + const std::string pte = dir + "sdpa_incache_" + seq.name + ".pte"; + FILE* f = std::fopen(pte.c_str(), "rb"); + if (!f) { + continue; + } + std::fclose(f); + test_sdpa_incache_decode(seq, dir, /*fresh_per_step=*/false); + test_sdpa_incache_decode(seq, dir, /*fresh_per_step=*/true); } +} - if (!prepack2_model_path.empty() && !prepack2_golden_path.empty()) { - ok = test_prepack( - prepack2_model_path, prepack2_golden_path, "x + w1 + w2") && - ok; +// If an SDPA dir was given, the exports must have produced .ptes for every +// family; a self-skip there means a silent export failure, not a pass. +TEST(WebGPUNative, SdpaAllFamiliesRanWhenDirSet) { + const std::string& dir = g_sdpa_dir; + if (dir.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SDPA_DIR not set"; + } + auto has_glob = [&](const std::string& prefix, const std::string& suffix) { + for (const auto& seq : kSdpaSequences) { + const std::string p = dir + prefix + seq.name + suffix; + FILE* f = std::fopen(p.c_str(), "rb"); + if (f) { + std::fclose(f); + return true; + } + } + return false; + }; + bool sdpa_ran = false; + for (const auto& cfg : kSdpaConfigs) { + const std::string pte = dir + "sdpa_" + cfg.name + ".pte"; + FILE* f = std::fopen(pte.c_str(), "rb"); + if (f) { + std::fclose(f); + sdpa_ran = true; + break; + } } + const bool replay_ran = [&] { + for (const auto& seq : kSdpaSequences) { + const std::string step0 = dir + "sdpa_" + seq.name + "_step0_S" + + std::to_string(seq.seq_lens[0]) + "_pos0.pte"; + FILE* f = std::fopen(step0.c_str(), "rb"); + if (f) { + std::fclose(f); + return true; + } + } + return false; + }(); + const bool dyn_ran = has_glob("sdpa_dyn_", ".pte"); + const bool incache_ran = has_glob("sdpa_incache_", ".pte"); + EXPECT_TRUE(sdpa_ran && replay_ran && dyn_ran && incache_ran) + << "WEBGPU_TEST_SDPA_DIR set but an SDPA family found no .pte"; +} - if (!prepack_tied_model_path.empty() && !prepack_tied_golden_path.empty()) { - ok = test_prepack( - prepack_tied_model_path, - prepack_tied_golden_path, - "x + w + w (tied weights, shared key)") && - ok; +TEST(WebGPUNative, SymintRoundtrip) { + if (g_symint_blob.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; } + test_symint_roundtrip(g_symint_blob); +} - bool sdpa_ran = false; - bool sdpa_ok = test_sdpa_sweep(sdpa_dir, &sdpa_ran); - if (sdpa_ran) { - ok = sdpa_ok && ok; +TEST(WebGPUNative, ResizeHook) { + if (g_symint_blob.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; } + test_resize_hook(g_symint_blob); +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); - // Guard python<->C++ ramp bit-identity (recorded: _ramp_t(0,17,8,2)=0.1875). - if (std::abs(sdpa_ramp_t(0, 17, 8, 2) - 0.1875f) > 1e-12f) { - printf("FAIL: sdpa_ramp_t bit-identity check\n"); - ok = false; + if (const char* env = std::getenv("WEBGPU_TEST_UPDATE_CACHE_MODEL")) { + g_update_cache_model_path = env; } - // Guard the adversarial denom path: sdpa_ramp(0,17,8,0.5)= -16.0 exactly. - if (std::abs(sdpa_ramp(0, 17, 8, 0.5f) - (-16.0f)) > 1e-12f) { - printf("FAIL: sdpa_ramp denom bit-identity check\n"); - ok = false; + + // Quantized-linear sweep dir (mirrors WEBGPU_TEST_SDPA_DIR). + if (const char* env = std::getenv("WEBGPU_TEST_QUANTIZED_LINEAR_DIR")) { + g_qlinear_dir = env; + if (!g_qlinear_dir.empty() && g_qlinear_dir.back() != '/') { + g_qlinear_dir += '/'; + } } - bool replay_ran = false; - bool replay_ok = test_sdpa_replay_sweep(sdpa_dir, &replay_ran); - if (replay_ran) { - ok = replay_ok && ok; + if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_MODEL")) { + g_prepack_model_path = env; + } + if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_GOLDEN")) { + g_prepack_golden_path = env; } - bool dyn_ran = false; - bool dyn_ok = test_sdpa_dynamic_decode_sweep(sdpa_dir, &dyn_ran); - if (dyn_ran) { - ok = dyn_ok && ok; + if (const char* env = std::getenv("WEBGPU_TEST_PREPACK2_MODEL")) { + g_prepack2_model_path = env; + } + if (const char* env = std::getenv("WEBGPU_TEST_PREPACK2_GOLDEN")) { + g_prepack2_golden_path = env; } - bool incache_ran = false; - bool incache_ok = test_sdpa_incache_decode_sweep(sdpa_dir, &incache_ran); - if (incache_ran) { - ok = incache_ok && ok; + if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_TIED_MODEL")) { + g_prepack_tied_model_path = env; + } + if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_TIED_GOLDEN")) { + g_prepack_tied_golden_path = env; } - // If an SDPA dir was given, the exports must have produced .ptes for every - // family; a self-skip there means a silent export failure, not a pass. - if (!sdpa_dir.empty() && - !(sdpa_ran && replay_ran && dyn_ran && incache_ran)) { - printf("FAIL: WEBGPU_TEST_SDPA_DIR set but an SDPA family found no .pte\n"); - ok = false; + // SDPA sweep: configs self-discover their sdpa_.pte/.golden.bin under + // this directory (default "" = the embedded-file root / cwd). Set + // WEBGPU_TEST_SDPA_DIR to point at the exported .pte directory (e.g. /tmp/). + if (const char* env = std::getenv("WEBGPU_TEST_SDPA_DIR")) { + g_sdpa_dir = env; + if (!g_sdpa_dir.empty() && g_sdpa_dir.back() != '/') { + g_sdpa_dir += '/'; + } } if (const char* env = std::getenv("WEBGPU_TEST_SYMINT_BLOB")) { - ok = test_symint_roundtrip(env) && ok; - ok = test_resize_hook(env) && ok; + g_symint_blob = env; + } + + WebGPUContext ctx; + try { + ctx = create_webgpu_context(); + } catch (const std::exception& e) { + printf("SKIP: %s\n", e.what()); + return 0; } + set_default_webgpu_context(&ctx); + printf("WebGPU device acquired (native)\n"); + + const int rc = RUN_ALL_TESTS(); set_default_webgpu_context(nullptr); destroy_webgpu_context(ctx); - - if (!ok) { - return 1; - } - printf("\nAll tests passed\n"); - return 0; + return rc; }