Skip to content

[ExecuTorch][WebGPU] Migrate 9 ops to make_compute_pipeline (layout:"auto") #20868

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 1 commit into
base: gh/JCNTH/43/base
Choose a base branch
from
Open

[ExecuTorch][WebGPU] Migrate 9 ops to make_compute_pipeline (layout:"auto") #20868

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
107 changes: 25 additions & 82 deletions backends/webgpu/runtime/ops/add/BinaryOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,90 +76,37 @@ void add_impl(WebGPUGraph& graph, const std::vector<int>& args) {

graph.add_uniform_buffer_bytes(sizeof(AddParams));

// Create shader module from built-in WGSL source
WGPUShaderSourceWGSL wgsl_desc = {};
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
wgsl_desc.code = {kBinaryAddWGSL, WGPU_STRLEN};

WGPUShaderModuleDescriptor shader_desc = {};
shader_desc.nextInChain = &wgsl_desc.chain;
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);

// Create bind group layout: 3 storage buffers + 1 uniform
WGPUBindGroupLayoutEntry entries[4] = {};

// input1 - storage buffer, read-only
entries[0].binding = 0;
entries[0].visibility = WGPUShaderStage_Compute;
entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;

// input2 - storage buffer, read-only
entries[1].binding = 1;
entries[1].visibility = WGPUShaderStage_Compute;
entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;

// output - storage buffer, read-write
entries[2].binding = 2;
entries[2].visibility = WGPUShaderStage_Compute;
entries[2].buffer.type = WGPUBufferBindingType_Storage;

// params - uniform buffer
entries[3].binding = 3;
entries[3].visibility = WGPUShaderStage_Compute;
entries[3].buffer.type = WGPUBufferBindingType_Uniform;

WGPUBindGroupLayoutDescriptor bgl_desc = {};
bgl_desc.entryCount = 4;
bgl_desc.entries = entries;
WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);

// Create pipeline layout
WGPUPipelineLayoutDescriptor pl_desc = {};
pl_desc.bindGroupLayoutCount = 1;
pl_desc.bindGroupLayouts = &bgl;
WGPUPipelineLayout pipeline_layout =
wgpuDeviceCreatePipelineLayout(device, &pl_desc);

// Create compute pipeline
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);

// Create bind group with actual buffers
const auto& in1_tensor = graph.get_tensor(in1_id);
const auto& in2_tensor = graph.get_tensor(in2_id);

WGPUBindGroupEntry bg_entries[4] = {};

bg_entries[0].binding = 0;
bg_entries[0].buffer = in1_tensor.buffer;
bg_entries[0].size = in1_tensor.nbytes;

bg_entries[1].binding = 1;
bg_entries[1].buffer = in2_tensor.buffer;
bg_entries[1].size = in2_tensor.nbytes;

bg_entries[2].binding = 2;
bg_entries[2].buffer = out_tensor.buffer;
bg_entries[2].size = out_tensor.nbytes;

bg_entries[3].binding = 3;
bg_entries[3].buffer = uniform_buffer;
bg_entries[3].size = sizeof(AddParams);

WGPUBindGroupDescriptor bg_desc = {};
bg_desc.layout = bgl;
bg_desc.entryCount = 4;
bg_desc.entries = bg_entries;
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
device,
kBinaryAddWGSL,
{
{0,
WGPUBufferBindingType_ReadOnlyStorage,
in1_tensor.buffer,
in1_tensor.nbytes},
{1,
WGPUBufferBindingType_ReadOnlyStorage,
in2_tensor.buffer,
in2_tensor.nbytes},
{2,
WGPUBufferBindingType_Storage,
out_tensor.buffer,
out_tensor.nbytes},
{3, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(AddParams)},
},
&wg_size_constant,
1);

graph.add_dispatch(
{pipeline, bind_group, workgroup_count.x, "", workgroup_count.y});
{bundle.pipeline,
bundle.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.
Expand Down Expand Up @@ -193,10 +140,6 @@ void add_impl(WebGPUGraph& graph, const std::vector<int>& args) {
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);
// Graph owns it so a resize hook can rewrite it; freed in the dtor.
graph.own_uniform_buffer(uniform_buffer);
}
Expand Down
93 changes: 26 additions & 67 deletions backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,76 +197,38 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
wgpuBufferUnmap(uniform_buffer);
graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams));

WGPUShaderSourceWGSL wgsl_desc = {};
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
wgsl_desc.code = {kEmbeddingQ4gswWGSL, WGPU_STRLEN};
WGPUShaderModuleDescriptor shader_desc = {};
shader_desc.nextInChain = &wgsl_desc.chain;
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);

// Bind group layout: out (rw) + indices/weight/scales (ro storage) + 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;
WGPUComputePipeline pipeline =
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);

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 = indices.buffer;
bg_entries[1].size = indices.nbytes;
bg_entries[2].binding = 2;
bg_entries[2].buffer = weight.buffer;
bg_entries[2].size = weight.nbytes;
bg_entries[3].binding = 3;
bg_entries[3].buffer = scales.buffer;
bg_entries[3].size = scales.nbytes;
bg_entries[4].binding = 4;
bg_entries[4].buffer = uniform_buffer;
bg_entries[4].size = sizeof(EmbeddingParams);

WGPUBindGroupDescriptor bg_desc = {};
bg_desc.layout = bgl;
bg_desc.entryCount = 5;
bg_desc.entries = bg_entries;
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
// Bindings: out (rw) + indices/weight/scales (ro storage) + uniform.
utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
device,
kEmbeddingQ4gswWGSL,
{
{0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
{1,
WGPUBufferBindingType_ReadOnlyStorage,
indices.buffer,
indices.nbytes},
{2,
WGPUBufferBindingType_ReadOnlyStorage,
weight.buffer,
weight.nbytes},
{3,
WGPUBufferBindingType_ReadOnlyStorage,
scales.buffer,
scales.nbytes},
{4,
WGPUBufferBindingType_Uniform,
uniform_buffer,
sizeof(EmbeddingParams)},
},
&wg_size_constant,
1);

const size_t dispatch_idx = graph.add_dispatch(
{pipeline, bind_group, workgroup_count, "embedding_q4gsw"});
{bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"});

// Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim].
const uint32_t gs_u = static_cast<uint32_t>(group_size);
Expand Down Expand Up @@ -297,9 +259,6 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
params_buf);
});

wgpuShaderModuleRelease(shader);
wgpuBindGroupLayoutRelease(bgl);
wgpuPipelineLayoutRelease(pipeline_layout);
// Graph owns it so the resize hook can rewrite it; freed in the dtor.
graph.own_uniform_buffer(uniform_buffer);
}
Expand Down
92 changes: 26 additions & 66 deletions backends/webgpu/runtime/ops/index/Index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,73 +109,33 @@ void index_impl(WebGPUGraph& graph, const std::vector<int>& args) {
utils::make_uniform(device, &params, sizeof(IndexParams));
graph.add_uniform_buffer_bytes(sizeof(IndexParams));

WGPUShaderSourceWGSL wgsl_desc = {};
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
wgsl_desc.code = {kIndexWGSL, WGPU_STRLEN};
WGPUShaderModuleDescriptor shader_desc = {};
shader_desc.nextInChain = &wgsl_desc.chain;
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);

// self (read), out (read_write), index (read i32), params (uniform).
WGPUBindGroupLayoutEntry entries[4] = {};
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_ReadOnlyStorage;
entries[3].binding = 3;
entries[3].visibility = WGPUShaderStage_Compute;
entries[3].buffer.type = WGPUBufferBindingType_Uniform;

WGPUBindGroupLayoutDescriptor bgl_desc = {};
bgl_desc.entryCount = 4;
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);

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[4] = {};
bg_entries[0].binding = 0;
bg_entries[0].buffer = self_tensor.buffer;
bg_entries[0].size = self_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 = index_tensor.buffer;
bg_entries[2].size = index_tensor.nbytes;
bg_entries[3].binding = 3;
bg_entries[3].buffer = uniform_buffer;
bg_entries[3].size = sizeof(IndexParams);

WGPUBindGroupDescriptor bg_desc = {};
bg_desc.layout = bgl;
bg_desc.entryCount = 4;
bg_desc.entries = bg_entries;
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

graph.add_dispatch({pipeline, bind_group, workgroup_count});

wgpuShaderModuleRelease(shader);
wgpuBindGroupLayoutRelease(bgl);
wgpuPipelineLayoutRelease(pipeline_layout);
utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
device,
kIndexWGSL,
{
{0,
WGPUBufferBindingType_ReadOnlyStorage,
self_tensor.buffer,
self_tensor.nbytes},
{1,
WGPUBufferBindingType_Storage,
out_tensor.buffer,
out_tensor.nbytes},
{2,
WGPUBufferBindingType_ReadOnlyStorage,
index_tensor.buffer,
index_tensor.nbytes},
{3,
WGPUBufferBindingType_Uniform,
uniform_buffer,
sizeof(IndexParams)},
},
&wg_size_constant,
1);

graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});

// The bind group keeps the uniform buffer alive until release.
wgpuBufferRelease(uniform_buffer);
}
Expand Down
Loading
Loading