Skip to content

Commit 0513f92

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] Add gather and embedding ops to the WebGPU backend
Pull Request resolved: #20927 Add `aten.gather.default` (row-gather) + `aten.embedding.default` handlers for the training tail. Key changes: - `runtime/ops/{gather,embedding}/` — gather / row-gather WGSL kernels + handlers - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Reuses the shared Vulkan partitioner (`aten.gather`/`embedding` already registered); WebGPU kernels only. Co-authored-with: Claude Code. ghstack-source-id: 405026068 @exported-using-ghexport Differential Revision: [D111755125](https://our.internmc.facebook.com/intern/diff/D111755125/)
1 parent d1832bd commit 0513f92

7 files changed

Lines changed: 517 additions & 0 deletions

File tree

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,9 @@ set(WEBGPU_SRCS
6161
runtime/ops/sub/BinaryOp.cpp
6262
runtime/ops/where/Where.cpp
6363
runtime/ops/boolean_op/BooleanOp.cpp
64+
runtime/ops/gather/Gather.cpp
6465
runtime/ops/linear/Linear.cpp
66+
runtime/ops/embedding/Embedding.cpp
6567
)
6668

6769
add_library(webgpu_backend ${WEBGPU_SRCS})

‎backends/webgpu/runtime/ops/embedding/Embedding.cpp‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
10+
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
11+
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
12+
#include <executorch/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h>
13+
14+
#include <webgpu/webgpu.h>
15+
16+
#include <cstdint>
17+
#include <stdexcept>
18+
19+
namespace executorch::backends::webgpu {
20+
21+
namespace {
22+
23+
struct EmbeddingParams {
24+
uint32_t num_elements;
25+
uint32_t dim;
26+
uint32_t _pad[2];
27+
};
28+
29+
// aten.embedding: out[row, :] = weight[indices[row], :] (fp32 weight, i32 idx).
30+
void embedding_impl(WebGPUGraph& graph, const std::vector<int>& args) {
31+
const int weight_id = args.at(0);
32+
const int indices_id = args.at(1);
33+
const int out_id = args.at(args.size() - 1);
34+
35+
WGPUDevice device = graph.device();
36+
const auto& weight = graph.get_tensor(weight_id);
37+
const auto& indices = graph.get_tensor(indices_id);
38+
const auto& out = graph.get_tensor(out_id);
39+
40+
if (weight.buffer == nullptr || indices.buffer == nullptr ||
41+
out.buffer == nullptr) {
42+
throw std::runtime_error("embedding: null buffer binding");
43+
}
44+
if (weight.dims.size() != 2) {
45+
throw std::runtime_error("embedding: weight must be 2D [vocab, dim]");
46+
}
47+
const uint32_t dim = static_cast<uint32_t>(weight.dims[1]);
48+
if (dim == 0) {
49+
throw std::runtime_error("embedding: dim == 0");
50+
}
51+
const size_t out_numel = out.nbytes / sizeof(float);
52+
// Index is the int32 downcast of the int64 indices (mirror index op).
53+
const size_t index_numel = indices.nbytes / sizeof(int32_t);
54+
if (out.nbytes != out_numel * sizeof(float) ||
55+
weight.nbytes % sizeof(float) != 0 ||
56+
indices.nbytes != index_numel * sizeof(int32_t)) {
57+
throw std::runtime_error(
58+
"embedding: fp32 weight/out + i32 indices required");
59+
}
60+
if (out_numel != index_numel * dim) {
61+
throw std::runtime_error("embedding: out numel != num_indices * dim");
62+
}
63+
64+
uint32_t num_elements = static_cast<uint32_t>(out_numel);
65+
uint32_t wg_size =
66+
utils::clamp_workgroup_size(device, kEmbeddingWorkgroupSizeX);
67+
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
68+
device, num_elements, wg_size, "embedding");
69+
70+
WGPUConstantEntry wg_size_constant = {};
71+
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
72+
wg_size_constant.value = static_cast<double>(wg_size);
73+
74+
EmbeddingParams params = {};
75+
params.num_elements = num_elements;
76+
params.dim = dim;
77+
WGPUBuffer uniform_buffer =
78+
utils::make_uniform(device, &params, sizeof(EmbeddingParams));
79+
graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams));
80+
81+
WGPUShaderSourceWGSL wgsl_desc = {};
82+
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
83+
wgsl_desc.code = {kEmbeddingWGSL, WGPU_STRLEN};
84+
WGPUShaderModuleDescriptor shader_desc = {};
85+
shader_desc.nextInChain = &wgsl_desc.chain;
86+
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
87+
88+
WGPUBindGroupLayoutEntry entries[4] = {};
89+
entries[0].binding = 0;
90+
entries[0].visibility = WGPUShaderStage_Compute;
91+
entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
92+
entries[1].binding = 1;
93+
entries[1].visibility = WGPUShaderStage_Compute;
94+
entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
95+
entries[2].binding = 2;
96+
entries[2].visibility = WGPUShaderStage_Compute;
97+
entries[2].buffer.type = WGPUBufferBindingType_Storage;
98+
entries[3].binding = 3;
99+
entries[3].visibility = WGPUShaderStage_Compute;
100+
entries[3].buffer.type = WGPUBufferBindingType_Uniform;
101+
102+
WGPUBindGroupLayoutDescriptor bgl_desc = {};
103+
bgl_desc.entryCount = 4;
104+
bgl_desc.entries = entries;
105+
WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
106+
107+
WGPUPipelineLayoutDescriptor pl_desc = {};
108+
pl_desc.bindGroupLayoutCount = 1;
109+
pl_desc.bindGroupLayouts = &bgl;
110+
WGPUPipelineLayout pipeline_layout =
111+
wgpuDeviceCreatePipelineLayout(device, &pl_desc);
112+
113+
WGPUComputePipelineDescriptor pipeline_desc = {};
114+
pipeline_desc.layout = pipeline_layout;
115+
pipeline_desc.compute.module = shader;
116+
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
117+
pipeline_desc.compute.constantCount = 1;
118+
pipeline_desc.compute.constants = &wg_size_constant;
119+
WGPUComputePipeline pipeline =
120+
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
121+
122+
WGPUBindGroupEntry bg_entries[4] = {};
123+
bg_entries[0].binding = 0;
124+
bg_entries[0].buffer = weight.buffer;
125+
bg_entries[0].size = weight.nbytes;
126+
bg_entries[1].binding = 1;
127+
bg_entries[1].buffer = indices.buffer;
128+
bg_entries[1].size = indices.nbytes;
129+
bg_entries[2].binding = 2;
130+
bg_entries[2].buffer = out.buffer;
131+
bg_entries[2].size = out.nbytes;
132+
bg_entries[3].binding = 3;
133+
bg_entries[3].buffer = uniform_buffer;
134+
bg_entries[3].size = sizeof(EmbeddingParams);
135+
136+
WGPUBindGroupDescriptor bg_desc = {};
137+
bg_desc.layout = bgl;
138+
bg_desc.entryCount = 4;
139+
bg_desc.entries = bg_entries;
140+
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
141+
142+
graph.add_dispatch(
143+
{pipeline, bind_group, workgroup_count.x, "", workgroup_count.y});
144+
145+
wgpuShaderModuleRelease(shader);
146+
wgpuBindGroupLayoutRelease(bgl);
147+
wgpuPipelineLayoutRelease(pipeline_layout);
148+
wgpuBufferRelease(uniform_buffer);
149+
}
150+
151+
} // namespace
152+
153+
WEBGPU_REGISTER_OPERATORS {
154+
WEBGPU_REGISTER_OP(aten.embedding.default, embedding_impl);
155+
}
156+
157+
} // namespace executorch::backends::webgpu

‎backends/webgpu/runtime/ops/embedding/embedding.wgsl‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
@group(0) @binding(0) var<storage, read> weight: array<f32>;
2+
@group(0) @binding(1) var<storage, read> indices: array<i32>;
3+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
4+
5+
struct Params {
6+
num_elements: u32,
7+
dim: u32,
8+
}
9+
@group(0) @binding(3) var<uniform> params: Params;
10+
11+
override wg_size: u32 = 256;
12+
13+
@compute @workgroup_size(wg_size)
14+
fn main(
15+
@builtin(global_invocation_id) gid: vec3<u32>,
16+
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
17+
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
18+
if (idx >= params.num_elements) {
19+
return;
20+
}
21+
let row = idx / params.dim;
22+
let col = idx % params.dim;
23+
let row_id = u32(indices[row]);
24+
output[idx] = weight[row_id * params.dim + col];
25+
}

‎backends/webgpu/runtime/ops/embedding/embedding_wgsl.h‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
namespace executorch::backends::webgpu {
14+
15+
// @generated from embedding.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: e7c94da588a9dd40c0e1be7d3c6ef33310c71390c57f6f330ae23a18866d0ea7
17+
inline constexpr const char* kEmbeddingWGSL = R"(
18+
@group(0) @binding(0) var<storage, read> weight: array<f32>;
19+
@group(0) @binding(1) var<storage, read> indices: array<i32>;
20+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
21+
22+
struct Params {
23+
num_elements: u32,
24+
dim: u32,
25+
}
26+
@group(0) @binding(3) var<uniform> params: Params;
27+
28+
override wg_size: u32 = 256;
29+
30+
@compute @workgroup_size(wg_size)
31+
fn main(
32+
@builtin(global_invocation_id) gid: vec3<u32>,
33+
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
34+
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
35+
if (idx >= params.num_elements) {
36+
return;
37+
}
38+
let row = idx / params.dim;
39+
let col = idx % params.dim;
40+
let row_id = u32(indices[row]);
41+
output[idx] = weight[row_id * params.dim + col];
42+
}
43+
)";
44+
45+
inline constexpr uint32_t kEmbeddingWorkgroupSizeX = 256;
46+
inline constexpr uint32_t kEmbeddingWorkgroupSizeY = 1;
47+
inline constexpr uint32_t kEmbeddingWorkgroupSizeZ = 1;
48+
49+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)