diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index adbf4301413..b12763afd6a 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -125,6 +125,7 @@ function(add_webgpu_native_test test_name test_src) target_link_libraries( ${test_name} PRIVATE webgpu_backend + vulkan_schema ${WEBGPU_GPU_LIB} executorch_core extension_module_static @@ -169,10 +170,11 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) ) # Device-free util unit test: no backend/Dawn link (pure manifest/tolerance - # helpers), so it does NOT use the native-test helper. + # + dispatch-grid-math helpers), so it does NOT use the native-test helper. add_executable( - webgpu_op_test_util_test test/op_tests/test_driver_util.cpp - test/op_tests/driver_util.cpp + webgpu_op_test_util_test + test/op_tests/test_driver_util.cpp test/op_tests/driver_util.cpp + test/native/test_webgpu_utils.cpp ) target_include_directories( webgpu_op_test_util_test diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h new file mode 100644 index 00000000000..78536a83d63 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -0,0 +1,104 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Pure dispatch-grid math with zero WebGPU/Dawn dependency, so it is +// unit-testable without a WGPUDevice (split out of WebGPUUtils.h, which +// requires for its device-facing functions). + +#include +#include +#include +#include +#include + +namespace executorch::backends::webgpu::utils { + +// Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). +template +inline T div_up(T a, T b) { + return (a + b - 1) / b; +} + +// Product of a tensor's dims; the same accumulation was duplicated per-op. +inline uint64_t numel(const std::vector& dims) { + uint64_t n = 1; + for (int64_t d : dims) { + if (d < 0) { + throw std::runtime_error("numel: negative dimension"); + } + n *= static_cast(d); + } + return n; +} + +// Broadcasts a 1- or 2-element int list to (h, w); PyTorch's convention for +// kernel_size/stride/padding/dilation args. Was duplicated as a local `hw` +// lambda in conv2d/conv_transpose2d/max_pool2d. +inline void parse_hw( + const std::vector& v, + uint32_t& h, + uint32_t& w, + const char* op_name, + const char* arg_name) { + if (v.size() == 1) { + h = w = static_cast(v[0]); + } else if (v.size() == 2) { + h = static_cast(v[0]); + w = static_cast(v[1]); + } else { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + arg_name + + " must be 1 or 2 elements"); + } +} + +// Adaptive 1D->2D dispatch grid. `count_x`/`count_y` are the dispatch dims; +// `stride_x` (= count_x * wg_size) lets the shader decode a flat index as +// `gid.y * stride_x + gid.x`. Used by ops whose element count can exceed the +// 65535 per-dimension ceiling that compute_1d_workgroup_count throws on. +struct DispatchGrid { + uint32_t wg_size; + uint32_t count_x; + uint32_t count_y; + uint32_t stride_x; +}; + +// Given the workgroup count needed (1D) and the device's per-dimension +// dispatch-count ceiling, compute a near-square 2D grid rather than +// {max_dim, div_up(total, max_dim)} — maxing one dim pads the other with +// mostly-idle workgroups (up to ~2x the needed launch) when total isn't a +// clean multiple of max_dim. +inline DispatchGrid compute_dispatch_grid_from_limits( + uint32_t total, // workgroups needed (1D) + uint32_t wg_size, + uint32_t max_dim, + const char* op_name) { + DispatchGrid g; + g.wg_size = wg_size; + if (total <= max_dim) { + g.count_x = total; + g.count_y = 1; + } else { + uint32_t sq = + static_cast(std::ceil(std::sqrt(static_cast(total)))); + g.count_x = sq < max_dim ? sq : max_dim; + g.count_y = div_up(total, g.count_x); + if (g.count_y > + max_dim) { // > max_dim^2 * wg threads — astronomically large + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": dispatch exceeds 2D grid capacity"); + } + } + g.stride_x = g.count_x * g.wg_size; + return g; +} + +} // namespace executorch::backends::webgpu::utils diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index d9ab6467882..afd97b33968 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -401,6 +401,7 @@ void WebGPUGraph::build( value_lists_.resize(num_vals); doubles_.resize(num_vals, 0.0); bools_.resize(num_vals, false); + strings_.resize(num_vals); // Pre-scan the op chain: a constant may be DEFERRED (no eager GPU buffer; the // prepack node materializes it once) only if it is a prepack source AND never @@ -608,6 +609,14 @@ void WebGPUGraph::build( bools_[i] = val->value_as_Bool()->bool_val(); break; } + case vkgraph::GraphTypes::String: { + value_types_[i] = ValueType::String; + const auto* sv = val->value_as_String()->string_val(); + if (sv) { + strings_[i] = sv->str(); + } + break; + } case vkgraph::GraphTypes::SymInt: { // Live scalar: small Uniform buffer the CPU rewrites per execute. value_types_[i] = ValueType::SymInt; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 66f0e401de5..eb046660511 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -152,6 +152,10 @@ class WebGPUGraph { bool get_bool(int id) const { return bools_[id]; } + // String value (e.g. gelu's `approximate` kwarg). + const std::string& get_string(int id) const { + return strings_[id]; + } // Live-scalar (SymInt) API; mirrors the Vulkan SymInt/ParamsBuffer UBO. // set_symint writes the buffer + marks dirty only if the value changed. @@ -240,6 +244,20 @@ class WebGPUGraph { return dispatches_.size() - 1; } + // 2D sibling of add_dispatch (sets workgroup_count_y); returns the index. + size_t add_dispatch_2d( + WGPUComputePipeline pipeline, + WGPUBindGroup bind_group, + uint32_t count_x, + uint32_t count_y) { + WebGPUDispatch d; + d.pipeline = pipeline; + d.bind_group = bind_group; + d.workgroup_count_x = count_x; + d.workgroup_count_y = count_y; + return add_dispatch(d); + } + // In-graph buffer-to-buffer DMA (e.g. flat copy); returns the dispatch index. size_t add_buffer_copy(WGPUBuffer src, WGPUBuffer dst, size_t nbytes) { WebGPUDispatch d; @@ -376,6 +394,7 @@ class WebGPUGraph { std::vector> value_lists_; std::vector doubles_; std::vector bools_; + std::vector strings_; // SymInt (live scalar): id -> {live Uniform buffer, current value}, sparse. struct SymIntSlot { diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 90ead722ace..cc02ca39fb1 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -8,9 +8,13 @@ #pragma once +#include +#include + #include #include +#include #include #include #include @@ -20,22 +24,10 @@ namespace executorch::backends::webgpu::utils { -// Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). -template -inline T div_up(T a, T b) { - return (a + b - 1) / b; -} - -// Product of dims (live element count); used by dynamic-resize hooks. +// Product of dims (live element count); used by dynamic-resize hooks. Delegates +// to numel (single impl; keeps the negative-dim guard, no caller churn). inline uint64_t numel_of(const std::vector& dims) { - uint64_t n = 1; - for (int64_t v : dims) { - if (v < 0) { - throw std::runtime_error("numel_of: negative dimension"); - } - n *= static_cast(v); - } - return n; + return numel(dims); } // Clamp workgroup size to device limit (SwiftShader caps at 128). @@ -71,24 +63,15 @@ inline uint32_t queried_max_workgroups(WGPUDevice device) { // split keeps the waste to O(sqrt(count)). Throws if even a max*max grid is too // small (a 3rd dispatch dimension, out of scope). The shader reconstructs the // linear index from @builtin(num_workgroups), so any x/y factoring works. +// Now delegates to the DispatchMath fold (single fold impl); wg_size=1 makes +// the returned grid's stride_x collapse harmlessly. inline WgCount fold_workgroup_count_2d( uint32_t count, uint32_t max_count, const char* op_name) { - if (count <= max_count) { - return {count, 1u}; - } - uint32_t x = - static_cast(std::ceil(std::sqrt(static_cast(count)))); - x = std::min(x, max_count); - // ceil-div written overflow-safe (count >= 1 here) as count nears UINT32_MAX. - uint32_t y = 1u + (count - 1u) / x; - if (y > max_count) { - throw std::runtime_error( - std::string("WebGPU ") + op_name + - ": workgroup count needs a 3rd dispatch dimension (unsupported)"); - } - return {x, y}; + DispatchGrid g = compute_dispatch_grid_from_limits( + count, /*wg_size=*/1, max_count, op_name); + return {g.count_x, g.count_y}; } // 1D dispatch count (mirrors Vulkan div_up); throws if > device limit. @@ -118,6 +101,60 @@ inline WgCount compute_2d_workgroup_count( count, queried_max_workgroups(device), op_name); } +inline DispatchGrid compute_dispatch_grid( + WGPUDevice device, + uint32_t num_threads, + uint32_t desired_wg, + const char* op_name) { + // Single limits query shared by wg-size clamp + max-dim (avoid 2 queries). + WGPULimits limits = {}; + bool got_limits = wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success; + uint32_t wg = (got_limits && limits.maxComputeInvocationsPerWorkgroup > 0) + ? std::min(desired_wg, limits.maxComputeInvocationsPerWorkgroup) + : desired_wg; + if (wg == 0) { + wg = 1; + } + uint32_t max_dim = (got_limits && limits.maxComputeWorkgroupsPerDimension > 0) + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; // WebGPU spec-default floor + uint32_t total = div_up(num_threads, wg); // workgroups needed (1D) + return compute_dispatch_grid_from_limits(total, wg, max_dim, op_name); +} + +// 2D tile grid for tiled kernels: {div_up(n, tile), div_up(m, tile)} workgroups +// (one per tile); throws if either dim exceeds the device's per-dim limit. +inline WgCount compute_tile_grid_2d( + WGPUDevice device, + uint32_t n, + uint32_t m, + uint32_t tile, + const char* op_name) { + uint32_t max_wgs = queried_max_workgroups(device); + uint32_t x = div_up(n, tile); + uint32_t y = div_up(m, tile); + if (x > max_wgs || y > max_wgs) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": tile grid exceeds dispatch limit"); + } + return {x, y}; +} + +// For "one workgroup per row" kernels (native_layer_norm, sdpa_softmax) where +// num_rows IS the workgroup count already, unlike compute_dispatch_grid's +// element-count-needing-division use case. wg_size=1 makes the returned +// grid's stride_x collapse to exactly count_x, giving a correct flat +// *workgroup*-index decode (row_idx = workgroup_id.y * stride_x + +// workgroup_id.x). +inline DispatchGrid compute_row_dispatch_grid( + WGPUDevice device, + uint32_t num_rows, + const char* op_name) { + return compute_dispatch_grid_from_limits( + num_rows, /*wg_size=*/1, queried_max_workgroups(device), op_name); +} + // Create a uniform buffer mapped-at-creation, copy `size` bytes in, and unmap. inline WGPUBuffer make_uniform(WGPUDevice device, const void* data, size_t size) { @@ -143,6 +180,270 @@ make_uniform(WGPUDevice device, const void* data, size_t size) { // loop over any excess work (vs compute_1d_workgroup_count, which throws). inline uint32_t clamp_workgroup_count(WGPUDevice device, uint32_t desired) { return std::min(desired, queried_max_workgroups(device)); + return std::min(desired, queried_max_workgroups(device)); +} + +// Zero-filled storage buffer; used as a dummy binding for an optional tensor +// arg (bias/mask/affine) the shader gates off and never reads. +inline WGPUBuffer make_storage_zeros(WGPUDevice device, size_t size) { + WGPUBufferDescriptor desc = {}; + desc.size = size; + desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + desc.mappedAtCreation = true; + WGPUBuffer buf = wgpuDeviceCreateBuffer(device, &desc); + if (!buf) { + throw std::runtime_error("make_storage_zeros: buffer creation failed"); + } + void* ptr = wgpuBufferGetMappedRange(buf, 0, size); + if (!ptr) { + wgpuBufferRelease(buf); + throw std::runtime_error("make_storage_zeros: mapped range is null"); + } + std::memset(ptr, 0, size); + wgpuBufferUnmap(buf); + return buf; +} + +// Validates a dim is a multiple of 4 (vec4-alignment precondition for a +// vec4-typed buffer view); throws loud, mirroring this backend's no-silent- +// return convention. Was independently duplicated in sdpa/Sdpa.cpp, +// sdpa_fd_decode/SdpaFdDecode.cpp, and et_vk_sdpa/EtVkSdpa.cpp. +inline void +check_vec4_aligned(uint32_t dim, const char* op_name, const char* dim_name) { + if (dim % 4 != 0) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + dim_name + + " must be a multiple of 4"); + } +} + +// fp32 byte-size guard (no runtime dtype): the serialized bytes must equal the +// element count times sizeof(float), else the tensor is not fp32. Returns the +// element count; replaces the `numel(dims)*sizeof(float) != nbytes` check +// duplicated across the fp32-only op handlers. +inline uint64_t +check_fp32(const WebGPUTensor& t, const char* op_name, const char* label) { + uint64_t n = numel(t.dims); + if (t.nbytes != n * sizeof(float)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + label + " must be fp32"); + } + return n; +} + +// Elementwise unary I/O guard: both buffers bound, byte counts fp32-aligned, +// and input/output sizes equal. Mirrors the guard inlined in every elementwise +// op handler (gelu/sigmoid/...). +inline void check_elementwise_fp32_io( + const WebGPUTensor& in, + const WebGPUTensor& out, + const char* op_name) { + if (in.buffer == nullptr || out.buffer == nullptr) { + throw std::runtime_error(std::string(op_name) + ": null buffer binding"); + } + if (in.nbytes % sizeof(float) != 0 || out.nbytes % sizeof(float) != 0) { + throw std::runtime_error( + std::string(op_name) + ": operand not 4-byte aligned"); + } + if (in.nbytes != out.nbytes) { + throw std::runtime_error( + std::string(op_name) + ": input/output size mismatch"); + } +} + +// Narrow a u64 element/byte count to u32 for a dispatch/param field; throws if +// it would truncate (the count exceeds the u32 addressing range). +inline uint32_t checked_u32(uint64_t v, const char* op_name) { + if (v > 0xFFFFFFFFull) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": output too large"); + } + return static_cast(v); +} + +// Extract a Scalar arg as float: Double -> cast, Int -> cast, Null -> fallback. +// Any other type throws (this backend's no-silent-return convention: an +// unexpected scalar type is a validation failure, not a "use the default"). +// Scope: the 3 already-identical Double/Int/Null-permissive copies only +// (addmm's beta/alpha, constant_pad_nd's value, native_layer_norm's epsilon) +// -- NOT et_vk.sdpa/sdpa's scale extraction, which is a stricter +// Double-or-Null-only policy (rejects Int) that would silently change +// behavior if forced into this signature. +inline float scalar_or(WebGPUGraph& graph, int id, float fallback) { + const auto vt = graph.get_value_type(id); + if (vt == WebGPUGraph::ValueType::Double) { + return static_cast(graph.get_double(id)); + } + if (vt == WebGPUGraph::ValueType::Int) { + return static_cast(graph.get_int(id)); + } + if (vt == WebGPUGraph::ValueType::Null) { + return fallback; + } + throw std::runtime_error("scalar_or: expected Double, Int, or None"); +} + +// Resolves an optional tensor arg (bias/mask/affine) to the binding to use: +// the real tensor's buffer+size if present, else a zero-filled dummy. Caller +// releases `.owned_dummy` (if non-null) after the dispatch is recorded. +struct OptionalBinding { + WGPUBuffer buffer; + uint64_t nbytes; + WGPUBuffer owned_dummy; // non-null only when a dummy was allocated +}; + +inline OptionalBinding make_optional_binding( + WGPUDevice device, + bool present, + WGPUBuffer real_buffer, + uint64_t real_nbytes) { + if (present) { + return {real_buffer, real_nbytes, nullptr}; + } + // 16 bytes (not 4): WebGPU validates a binding's size against the shader's + // DECLARED type at dispatch time regardless of which branch runs, and a + // vec4-typed binding requires >=16 bytes even when never read. + WGPUBuffer dummy = make_storage_zeros(device, 16); + return {dummy, 16, dummy}; +} + +// One compute-shader binding: a bind-group-layout entry + its bind-group +// entry share every field except `type` (layout) vs `buffer`/`size` (bind). +struct BindingSpec { + uint32_t binding; + WGPUBufferBindingType type; + WGPUBuffer buffer; + uint64_t size; +}; + +// Owns the shader module, bind-group layout, and pipeline layout, releasing +// them on destruction. `pipeline` and `bind_group` are NOT released here — +// every op hands them to WebGPUGraph::add_dispatch, which keeps them alive +// for the life of the graph (mirrors the "kept by dispatch" comment repeated +// across every op handler). +struct ComputePipelineBundle { + WGPUShaderModule shader = nullptr; + WGPUBindGroupLayout bind_group_layout = nullptr; + WGPUPipelineLayout pipeline_layout = nullptr; + WGPUComputePipeline pipeline = nullptr; + WGPUBindGroup bind_group = nullptr; + + ComputePipelineBundle() = default; + ComputePipelineBundle(const ComputePipelineBundle&) = delete; + ComputePipelineBundle& operator=(const ComputePipelineBundle&) = delete; + ComputePipelineBundle& operator=(ComputePipelineBundle&&) = delete; + + ComputePipelineBundle(ComputePipelineBundle&& other) noexcept + : shader(other.shader), + bind_group_layout(other.bind_group_layout), + pipeline_layout(other.pipeline_layout), + pipeline(other.pipeline), + bind_group(other.bind_group) { + other.shader = nullptr; + other.bind_group_layout = nullptr; + other.pipeline_layout = nullptr; + other.pipeline = nullptr; + other.bind_group = nullptr; + } + + ~ComputePipelineBundle() { + if (shader != nullptr) { + wgpuShaderModuleRelease(shader); + } + if (bind_group_layout != nullptr) { + wgpuBindGroupLayoutRelease(bind_group_layout); + } + if (pipeline_layout != nullptr) { + wgpuPipelineLayoutRelease(pipeline_layout); + } + } +}; + +// Builds shader module -> bind-group layout -> pipeline layout -> compute +// pipeline -> bind group from one binding list, replacing the ~50-line +// sequence duplicated (with varying binding counts) across every op handler. +inline ComputePipelineBundle make_compute_pipeline( + WGPUDevice device, + const char* wgsl_source, + const std::vector& bindings, + const WGPUConstantEntry* constants = nullptr, + size_t constant_count = 0, + const char* entry_point = "main") { + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + std::vector layout_entries(bindings.size()); + std::vector bind_entries(bindings.size()); + for (size_t i = 0; i < bindings.size(); i++) { + layout_entries[i] = {}; + layout_entries[i].binding = bindings[i].binding; + layout_entries[i].visibility = WGPUShaderStage_Compute; + layout_entries[i].buffer.type = bindings[i].type; + + bind_entries[i] = {}; + bind_entries[i].binding = bindings[i].binding; + bind_entries[i].buffer = bindings[i].buffer; + bind_entries[i].size = bindings[i].size; + } + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = layout_entries.size(); + bgl_desc.entries = layout_entries.data(); + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN}; + pipeline_desc.compute.constantCount = constant_count; + pipeline_desc.compute.constants = constants; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = bind_entries.size(); + bg_desc.entries = bind_entries.data(); + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + ComputePipelineBundle bundle; + bundle.shader = shader; + bundle.bind_group_layout = bgl; + bundle.pipeline_layout = pipeline_layout; + bundle.pipeline = pipeline; + bundle.bind_group = bind_group; + return bundle; +} + +// The {wg_size, stride_x} override-constant pair every 2D-spill dispatch +// builds from its DispatchGrid; was hand-rolled identically at 7 call sites. +inline std::array make_grid_constants( + const DispatchGrid& grid) { + std::array constants = {}; + constants[0].key = {"wg_size", WGPU_STRLEN}; + constants[0].value = static_cast(grid.wg_size); + constants[1].key = {"stride_x", WGPU_STRLEN}; + constants[1].value = static_cast(grid.stride_x); + return constants; +} + +// The 1-element sibling of make_grid_constants: the single wg_size override- +// constant, hand-rolled at 5 call sites. +inline WGPUConstantEntry make_wg_size_constant(uint32_t wg_size) { + WGPUConstantEntry constant = {}; + constant.key = {"wg_size", WGPU_STRLEN}; + constant.value = static_cast(wg_size); + return constant; } } // namespace executorch::backends::webgpu::utils diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp new file mode 100644 index 00000000000..5a9ca95a5ca --- /dev/null +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -0,0 +1,64 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Device-free unit tests for the dispatch-grid math (WebGPUDispatchMath.h has +// zero WebGPU/Dawn dependency, unlike WebGPUUtils.h which needs a WGPUDevice +// for its other helpers). + +#include + +#include + +using namespace executorch::backends::webgpu; + +TEST(WebGPUUtils, DispatchGridStaysOneDimUnderCeiling) { + utils::DispatchGrid g = + utils::compute_dispatch_grid_from_limits(1000u, 256u, 65535u, "test"); + EXPECT_EQ(g.count_x, 1000u); + EXPECT_EQ(g.count_y, 1u); + EXPECT_EQ(g.stride_x, g.count_x * g.wg_size); +} + +TEST(WebGPUUtils, DispatchGridPastCeilingIsNearSquareNotMaxedOut) { + // total=65536, max_dim=65535: one workgroup past the 1D ceiling. + // The old {max_dim, div_up(total,max_dim)} fold gave (65535, 2) = 131070 + // launched workgroups for 65536 needed (~100% overhead). A near-square + // grid should launch close to the needed count instead. + const uint32_t total = 65536u; + const uint32_t max_dim = 65535u; + utils::DispatchGrid g = + utils::compute_dispatch_grid_from_limits(total, 256u, max_dim, "test"); + + EXPECT_LE(static_cast(g.count_x) * g.count_y, total + total / 10) + << "near-square grid should launch within ~10% of the needed count, " + "not ~2x it"; + // Grid must still cover every needed workgroup. + EXPECT_GE(static_cast(g.count_x) * g.count_y, total); + // Not the old maxed-out-count_x behavior. + EXPECT_NE(g.count_x, max_dim); + EXPECT_EQ(g.stride_x, g.count_x * g.wg_size); +} + +TEST(WebGPUUtils, DispatchGridExactSquareCase) { + // total=65536 factors exactly as 256*256 — the near-square grid should + // find this with zero waste. + utils::DispatchGrid g = utils::compute_dispatch_grid_from_limits( + 65536u, 1u, 65535u, "test"); + EXPECT_EQ(g.count_x, 256u); + EXPECT_EQ(g.count_y, 256u); + EXPECT_EQ(static_cast(g.count_x) * g.count_y, 65536u); +} + +TEST(WebGPUUtils, DispatchGridThrowsPastCapacity) { + // total > max_dim^2: even a near-square grid can't fit in the 2D ceiling. + const uint32_t max_dim = 4u; + EXPECT_THROW( + utils::compute_dispatch_grid_from_limits( + static_cast(max_dim) * max_dim + 1u, 1u, max_dim, "test"), + std::runtime_error); +} diff --git a/backends/webgpu/test/op_tests/driver_util.cpp b/backends/webgpu/test/op_tests/driver_util.cpp index aa6c337e0c8..24003a2c3fb 100644 --- a/backends/webgpu/test/op_tests/driver_util.cpp +++ b/backends/webgpu/test/op_tests/driver_util.cpp @@ -49,6 +49,7 @@ std::vector parse_manifest(const std::string& manifest_path) { InputRef ir; ir.path = join(base, ie.at("path").get()); ir.shape = ie.at("shape").get>(); + ir.dtype = ie.value("dtype", std::string("float32")); m.inputs.push_back(std::move(ir)); } const auto& g = e.at("golden"); @@ -64,6 +65,20 @@ std::vector parse_manifest(const std::string& manifest_path) { return out; } +std::vector load_int32_bin(const std::string& path, size_t numel) { + FILE* f = std::fopen(path.c_str(), "rb"); + if (!f) { + return {}; + } + std::vector g(numel); + const size_t n = std::fread(g.data(), sizeof(int32_t), numel, f); + std::fclose(f); + if (n != numel) { + return {}; + } + return g; +} + std::vector load_fp32_bin(const std::string& path, size_t numel) { FILE* f = std::fopen(path.c_str(), "rb"); if (!f) { diff --git a/backends/webgpu/test/op_tests/driver_util.h b/backends/webgpu/test/op_tests/driver_util.h index 3f5e7c47058..4c638c4f57d 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -17,6 +17,7 @@ namespace executorch::backends::webgpu { struct InputRef { std::string path; std::vector shape; + std::string dtype = "float32"; }; struct GoldenRef { @@ -42,6 +43,7 @@ std::vector parse_manifest(const std::string& manifest_path); /// Load raw little-endian fp32; empty on size/IO mismatch. std::vector load_fp32_bin(const std::string& path, size_t numel); +std::vector load_int32_bin(const std::string& path, size_t numel); /// Element OK if abs_err <= atol OR rel_err <= rtol (rel floored at /// |golden|=1e-6). Sets the reported maxima; true iff all elements pass. diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 66e7e55bafc..9001592fca3 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -38,7 +38,12 @@ def _materialize(spec) -> torch.Tensor: else: shape, gen = spec, "randn" if callable(gen): - return gen(shape).to(torch.float32) + _t = gen(shape) + return ( + _t.to(torch.int32) + if not _t.is_floating_point() + else _t.to(torch.float32) + ) if gen == "randn": return torch.randn(*shape) if gen == "ramp": @@ -97,7 +102,9 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: elif golden_dtype == "float64": # fp64 oracle (~1e-15); deepcopy keeps the original module fp32. double_module = copy.deepcopy(module).double() - golden = double_module(*[x.double() for x in inputs]) + golden = double_module( + *[x.double() if x.is_floating_point() else x for x in inputs] + ) else: golden = eager # gather/copy: fp64 is bit-identical, skip it golden_t = golden[out_index] if isinstance(golden, (tuple, list)) else golden @@ -119,8 +126,15 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: input_entries: list[dict] = [] for i, t in enumerate(inputs): rel = f"{case_id}.in{i}.bin" - _write_fp32(t, os.path.join(out_dir, rel)) - input_entries.append({"path": rel, "shape": list(t.shape), "dtype": "float32"}) + if t.dtype == torch.int32: + t.detach().cpu().numpy().astype(" tensors; for (const auto& in : e_.inputs) { const size_t n = numel(in.shape); - auto data = load_fp32_bin(in.path, n); - ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; std::vector sizes( in.shape.begin(), in.shape.end()); - tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); + if (in.dtype == "int32") { + auto data = load_int32_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); + } else { + auto data = load_fp32_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); + } } std::vector inputs; for (auto& t : tensors) {