Skip to content

[ExecuTorch][WebGPU] Migrate multi-dispatch ops to make_compute_pipeline #20869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JCNTH wants to merge 3 commits into
base: gh/JCNTH/44/base
Choose a base branch
from
+95 −274
Open

[ExecuTorch][WebGPU] Migrate multi-dispatch ops to make_compute_pipeline #20869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aa943fc
Update
JCNTH Jul 10, 2026
7064f0f
Update
JCNTH Jul 16, 2026
983a89f
Update
JCNTH Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 23 additions & 72 deletions backends/webgpu/runtime/ops/cat/Cat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,42 +113,6 @@ void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) {
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
wg_size_constant.value = static_cast<double>(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]);
Expand All @@ -163,48 +127,35 @@ void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) {
utils::make_uniform(device, &params, 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<uint32_t>(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);
}
Expand Down
114 changes: 32 additions & 82 deletions backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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<double>(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);
Expand Down Expand Up @@ -255,59 +260,10 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector<int>& 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<double>(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,
Expand All @@ -319,8 +275,7 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector<int>& args) {
RopeDispatch k_disp = add_rope_dispatch(
graph,
device,
pipeline_k,
bgl,
wg_size,
xk,
xk_out,
freqs_cos,
Expand Down Expand Up @@ -371,11 +326,6 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector<int>& 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
Expand Down
81 changes: 19 additions & 62 deletions backends/webgpu/runtime/ops/sdpa/Sdpa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>

namespace executorch::backends::webgpu {

Expand Down Expand Up @@ -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<utils::BindingSpec> 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<double>(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);
Expand Down
Loading
Loading