From aa943fc1c727cdfc4beb9f6315e8bb892585d511 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:26:28 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/runtime/ops/cat/Cat.cpp | 95 ++++----------- .../runtime/ops/rope/RotaryEmbedding.cpp | 114 +++++------------- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 81 +++---------- .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 79 ++++-------- 4 files changed, 95 insertions(+), 274 deletions(-) diff --git a/backends/webgpu/runtime/ops/cat/Cat.cpp b/backends/webgpu/runtime/ops/cat/Cat.cpp index 0cfb857745c..128d33b1ef3 100644 --- a/backends/webgpu/runtime/ops/cat/Cat.cpp +++ b/backends/webgpu/runtime/ops/cat/Cat.cpp @@ -113,42 +113,6 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - // Shared shader/layout; fresh pipeline+bind group per input (no double-free). - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kCatWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - uint32_t off_k = 0; for (size_t k = 0; k < ids.size(); k++) { const auto& in_tensor = graph.get_tensor(ids[k]); @@ -163,48 +127,35 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(CatParams)); graph.add_uniform_buffer_bytes(sizeof(TensorMeta) + sizeof(CatParams)); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(CatParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, wg_counts[k]}); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kCatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + out_meta_buf, + sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(CatParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, wg_counts[k]}); // Drop our refs; this input's bind group keeps its uniforms alive. wgpuBufferRelease(in_meta_buf); wgpuBufferRelease(params_buf); off_k += static_cast(in_tensor.dims[dim]); } - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Drop our ref to the shared out_meta; the bind groups keep it alive. wgpuBufferRelease(out_meta_buf); } diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 76eafb0e738..03de78233b0 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -41,12 +41,11 @@ struct RopeDispatch { size_t dispatch_index; }; -// Rotate one (x->out) with the shared shader; freqs shared between xq and xk. +// Rotate one (x->out) with its own pipeline; freqs shared between xq and xk. RopeDispatch add_rope_dispatch( WebGPUGraph& graph, WGPUDevice device, - WGPUComputePipeline pipeline, - WGPUBindGroupLayout bgl, + uint32_t wg_size, const WebGPUTensor& x, const WebGPUTensor& out, const WebGPUTensor& freqs_cos, @@ -78,31 +77,37 @@ RopeDispatch add_rope_dispatch( wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(RotaryParams)); - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out.buffer; - bg_entries[0].size = out.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = x.buffer; - bg_entries[1].size = x.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = freqs_cos.buffer; - bg_entries[2].size = freqs_cos.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = freqs_sin.buffer; - bg_entries[3].size = freqs_sin.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = uniform_buffer; - bg_entries[4].size = sizeof(RotaryParams); + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kRotaryEmbeddingWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_cos.buffer, + freqs_cos.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_sin.buffer, + freqs_sin.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(RotaryParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_index = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "apply_rotary_emb"}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count, + "apply_rotary_emb"}); // Graph owns it so a resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); @@ -255,59 +260,10 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { wg_size, "apply_rotary_emb"); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kRotaryEmbeddingWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group: out (rw) + in/freqs_cos/freqs_sin (ro) + uniform. - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - // One pipeline per dispatch; a shared handle would double-free. - WGPUComputePipeline pipeline_q = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - WGPUComputePipeline pipeline_k = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - RopeDispatch q_disp = add_rope_dispatch( graph, device, - pipeline_q, - bgl, + wg_size, xq, xq_out, freqs_cos, @@ -319,8 +275,7 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { RopeDispatch k_disp = add_rope_dispatch( graph, device, - pipeline_k, - bgl, + wg_size, xk, xk_out, freqs_cos, @@ -371,11 +326,6 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { }; graph.add_tensor_resize_hook(xq_id, rope_hook); graph.add_tensor_resize_hook(xk_id, rope_hook); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - // pipeline_q/pipeline_k owned by their dispatches; graph dtor frees. } } // namespace diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 50321ba4bdf..b431cf0b7f0 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -24,6 +24,7 @@ #include #include #include +#include namespace executorch::backends::webgpu { @@ -152,84 +153,40 @@ void build_dispatch( const char* kernel_name = "") { WGPUDevice device = graph.device(); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group layout: storage entries then the uniform. - constexpr uint32_t kMaxEntries = 8; - if (n_storage + 1 > kMaxEntries) { - throw std::runtime_error("WebGPU sdpa: n_storage exceeds kMaxEntries"); - } - WGPUBindGroupLayoutEntry bgl_entries[kMaxEntries] = {}; + std::vector bindings; + bindings.reserve(n_storage + 1); const uint32_t uniform_binding = n_storage; for (uint32_t i = 0; i < n_storage; i++) { - bgl_entries[i].binding = i; - bgl_entries[i].visibility = WGPUShaderStage_Compute; - bgl_entries[i].buffer.type = (i == 0) - ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage; + bindings.push_back( + {i, + (i == 0) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage, + storage_bindings[i].buffer, + storage_bindings[i].size}); } - bgl_entries[uniform_binding].binding = uniform_binding; - bgl_entries[uniform_binding].visibility = WGPUShaderStage_Compute; - bgl_entries[uniform_binding].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = n_storage + 1; - bgl_desc.entries = bgl_entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); + bindings.push_back( + {uniform_binding, + WGPUBufferBindingType_Uniform, + uniform_buffer, + uniform_size}); // QK/AV/update_cache have an `override wg_size`; softmax (0) keeps a const. WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); + const size_t constant_count = (wg_size != 0) ? 1 : 0; - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - if (wg_size != 0) { - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - } - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[kMaxEntries] = {}; - for (uint32_t i = 0; i < n_storage; i++) { - bg_entries[i].binding = i; - bg_entries[i].buffer = storage_bindings[i].buffer; - bg_entries[i].size = storage_bindings[i].size; - } - bg_entries[uniform_binding].binding = uniform_binding; - bg_entries[uniform_binding].buffer = uniform_buffer; - bg_entries[uniform_binding].size = uniform_size; - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = n_storage + 1; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, wgsl_source, bindings, &wg_size_constant, constant_count); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count_x, kernel_name, workgroup_count_y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); if (retain_uniform) { // Graph owns it so a resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index ffd3b24dce3..12a5260af86 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -20,6 +20,7 @@ #include #include #include +#include namespace executorch::backends::webgpu { @@ -72,70 +73,32 @@ void build_dispatch( const char* kernel_name) { WGPUDevice device = graph.device(); - 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); - - constexpr uint32_t kMaxEntries = 8; - if (n_storage + 1u > kMaxEntries) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: bind group entry count exceeds kMaxEntries"); - } - WGPUBindGroupLayoutEntry bgl_entries[kMaxEntries] = {}; + // Storage bindings then the uniform; n_rw leading bindings are read_write. + // Under auto layout the access class is derived from the shader, so `type` + // here is descriptive only. + std::vector bindings; + bindings.reserve(n_storage + 1); const uint32_t uniform_binding = n_storage; for (uint32_t i = 0; i < n_storage; i++) { - bgl_entries[i].binding = i; - bgl_entries[i].visibility = WGPUShaderStage_Compute; - bgl_entries[i].buffer.type = (i < n_rw) - ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage; - } - bgl_entries[uniform_binding].binding = uniform_binding; - bgl_entries[uniform_binding].visibility = WGPUShaderStage_Compute; - bgl_entries[uniform_binding].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = n_storage + 1; - bgl_desc.entries = bgl_entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[kMaxEntries] = {}; - for (uint32_t i = 0; i < n_storage; i++) { - bg_entries[i].binding = i; - bg_entries[i].buffer = storage_bindings[i].buffer; - bg_entries[i].size = storage_bindings[i].size; + bindings.push_back( + {i, + (i < n_rw) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage, + storage_bindings[i].buffer, + storage_bindings[i].size}); } - bg_entries[uniform_binding].binding = uniform_binding; - bg_entries[uniform_binding].buffer = uniform_buffer; - bg_entries[uniform_binding].size = uniform_size; + bindings.push_back( + {uniform_binding, + WGPUBufferBindingType_Uniform, + uniform_buffer, + uniform_size}); - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = n_storage + 1; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = + utils::make_compute_pipeline(device, wgsl_source, bindings); - graph.add_dispatch({pipeline, bind_group, workgroup_count_x, kernel_name}); + graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, workgroup_count_x, kernel_name}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); wgpuBufferRelease(uniform_buffer); }