From 31c44644872076f7906bcb51f00069188c835703 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:26:18 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUUtils.h | 32 ++++++++----------- .../runtime/ops/max_pool2d/MaxPool2d.cpp | 12 +++---- .../runtime/ops/max_pool2d/max_pool2d.wgsl | 5 ++- .../runtime/ops/max_pool2d/max_pool2d_wgsl.h | 7 ++-- 4 files changed, 22 insertions(+), 34 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 8702fd43664..9508604e6ba 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -375,33 +375,24 @@ inline ComputePipelineBundle make_compute_pipeline( 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); - + // layout = nullptr => WebGPU auto-derives the bind-group layout from the + // shader's statically-used @group/@binding declarations, replacing the + // hand-built bind-group layout + pipeline layout. Precondition: every + // declared binding must be statically referenced by the shader (optional + // bindings backed by a dummy buffer are, under a runtime guard) -- auto + // layout omits an unreferenced binding, whose bind-group entry would then + // mismatch. BindingSpec.type is unused under auto layout (kept so call sites + // need no change). WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; + pipeline_desc.layout = nullptr; pipeline_desc.compute.module = shader; pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN}; pipeline_desc.compute.constantCount = constant_count; @@ -409,6 +400,9 @@ inline ComputePipelineBundle make_compute_pipeline( WGPUComputePipeline pipeline = wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + // Owned reference (must be released) -> handed to the bundle dtor. + WGPUBindGroupLayout bgl = wgpuComputePipelineGetBindGroupLayout(pipeline, 0); + WGPUBindGroupDescriptor bg_desc = {}; bg_desc.layout = bgl; bg_desc.entryCount = bind_entries.size(); @@ -418,7 +412,7 @@ inline ComputePipelineBundle make_compute_pipeline( ComputePipelineBundle bundle; bundle.shader = shader; bundle.bind_group_layout = bgl; - bundle.pipeline_layout = pipeline_layout; + bundle.pipeline_layout = nullptr; // none created under auto layout bundle.pipeline = pipeline; bundle.bind_group = bind_group; return bundle; diff --git a/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp index 682eff61197..4ef0f3cb7fa 100644 --- a/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp +++ b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp @@ -24,7 +24,7 @@ namespace { struct PoolParams { uint32_t N, C, IH, IW, OH, OW; uint32_t kH, kW, sH, sW, pH, pW, dH, dW; - uint32_t _p0, _p1; + uint32_t write_indices, _p1; }; static_assert(sizeof(PoolParams) == 64, "PoolParams must be 64 bytes"); @@ -145,17 +145,13 @@ void max_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { params.pW = pW; params.dH = dH; params.dW = dW; + params.write_indices = has_indices ? 1u : 0u; WGPUBuffer uniform_buffer = utils::make_uniform(device, ¶ms, sizeof(PoolParams)); graph.add_uniform_buffer_bytes(sizeof(PoolParams)); auto grid_constants = utils::make_grid_constants(grid); - WGPUConstantEntry constants[3] = {}; - constants[0] = grid_constants[0]; - constants[1] = grid_constants[1]; - constants[2].key = {"write_indices", WGPU_STRLEN}; - constants[2].value = has_indices ? 1.0 : 0.0; // write_indices==0 -> the shader never stores to out_idx, so a tiny dummy // buffer is safe (mirrors NativeLayerNorm.cpp's dummy_affine pattern). @@ -177,8 +173,8 @@ void max_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { sizeof(PoolParams)}, {3, WGPUBufferBindingType_Storage, idx.buffer, idx.nbytes}, }, - constants, - 3); + grid_constants.data(), + grid_constants.size()); graph.add_dispatch_2d( bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); diff --git a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl index 3d7c67ddf9a..01f5153ce5b 100644 --- a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl +++ b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl @@ -13,7 +13,7 @@ struct Params { pW: u32, dH: u32, dW: u32, - _p0: u32, + write_indices: u32, // 1 = also write out_idx (runtime gate; keeps out_idx statically used) _p1: u32, } @@ -24,7 +24,6 @@ struct Params { override wg_size: u32 = 256; override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill -override write_indices: u32 = 0u; // 1 = also write out_idx (compile-time gate, mirrors Vulkan) // max_pool2d (values [+ optional indices]), NCHW row-major, fp32. One thread per // output element (n, c, oh, ow); gather the window, take the max. General @@ -73,7 +72,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } } out_vals[i] = best; - if (write_indices != 0u) { + if (params.write_indices != 0u) { out_idx[i] = best_idx; } } diff --git a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h index 1dbe0571e40..22881074c37 100644 --- a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h +++ b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from max_pool2d.wgsl - DO NOT EDIT. -// wgsl-sha256: 2215b7ba725face441baac5eca8e2133458b414377868354c3fda9046a88109b +// wgsl-sha256: 205fd7013aaacb5e5d4bda7499d436da6ca76e02f910962223c0dbd65f95e6ac inline constexpr const char* kMaxPool2dWGSL = R"( struct Params { N: u32, @@ -30,7 +30,7 @@ struct Params { pW: u32, dH: u32, dW: u32, - _p0: u32, + write_indices: u32, // 1 = also write out_idx (runtime gate; keeps out_idx statically used) _p1: u32, } @@ -41,7 +41,6 @@ struct Params { override wg_size: u32 = 256; override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill -override write_indices: u32 = 0u; // 1 = also write out_idx (compile-time gate, mirrors Vulkan) // max_pool2d (values [+ optional indices]), NCHW row-major, fp32. One thread per // output element (n, c, oh, ow); gather the window, take the max. General @@ -90,7 +89,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } } out_vals[i] = best; - if (write_indices != 0u) { + if (params.write_indices != 0u) { out_idx[i] = best_idx; } }