Skip to content

Commit 1a20522

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] Add shape/copy ops (dim_order, expand_copy, fill) to the WebGPU backend
Pull Request resolved: #20929 Add `_clone_dim_order` (DMA copy), `aten.expand_copy` (broadcast), and `fill` (`aten.full`/`full_like`/`scalar_tensor`) handlers for the training tail. Key changes: - `runtime/ops/{dim_order,expand_copy,fill}/` — WGSL kernels + handlers (dim_order is a native buffer DMA copy) - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Reuses the shared Vulkan partitioner (`_clone_dim_order`/`expand_copy`/`full`* already registered); WebGPU kernels only. Co-authored-with: Claude Code. ghstack-source-id: 405026073 @exported-using-ghexport Differential Revision: [D111755131](https://our.internmc.facebook.com/intern/diff/D111755131/)
1 parent f1dea3b commit 1a20522

8 files changed

Lines changed: 496 additions & 0 deletions

File tree

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,9 @@ set(WEBGPU_SRCS
6262
runtime/ops/where/Where.cpp
6363
runtime/ops/boolean_op/BooleanOp.cpp
6464
runtime/ops/gather/Gather.cpp
65+
runtime/ops/expand_copy/ExpandCopy.cpp
66+
runtime/ops/fill/Fill.cpp
67+
runtime/ops/dim_order/DimOrder.cpp
6568
runtime/ops/linear/Linear.cpp
6669
runtime/ops/embedding/Embedding.cpp
6770
)

‎backends/webgpu/runtime/ops/dim_order/DimOrder.cpp‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
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/ops/OperatorRegistry.h>
11+
#include <executorch/backends/webgpu/runtime/ops/view_copy/view_copy.h>
12+
13+
#include <vector>
14+
15+
namespace executorch::backends::webgpu {
16+
17+
namespace {
18+
19+
// _clone_dim_order = numel-preserving flat copy (shared DMA helper).
20+
void clone_dim_order_impl(WebGPUGraph& graph, const std::vector<int>& args) {
21+
add_flat_copy(graph, args.at(0), args.at(args.size() - 1));
22+
}
23+
24+
} // namespace
25+
26+
WEBGPU_REGISTER_OPERATORS {
27+
WEBGPU_REGISTER_OP(
28+
dim_order_ops._clone_dim_order.default, clone_dim_order_impl);
29+
}
30+
31+
} // namespace executorch::backends::webgpu

‎backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
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/TensorMeta.h>
13+
#include <executorch/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h>
14+
15+
#include <webgpu/webgpu.h>
16+
17+
#include <stdexcept>
18+
19+
namespace executorch::backends::webgpu {
20+
21+
namespace {
22+
23+
// out coord -> in coord; size-1 in-dims broadcast via clamp (mirrors mul).
24+
void expand_copy_impl(WebGPUGraph& graph, const std::vector<int>& args) {
25+
const int in_id = args.at(0);
26+
const int out_id = args.at(args.size() - 1);
27+
28+
if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor ||
29+
graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) {
30+
throw std::runtime_error("expand_copy: in/out arg is not a tensor");
31+
}
32+
33+
WGPUDevice device = graph.device();
34+
const auto& in_tensor = graph.get_tensor(in_id);
35+
const auto& out_tensor = graph.get_tensor(out_id);
36+
37+
TensorMeta out_meta;
38+
TensorMeta in_meta;
39+
fill_tensor_meta(out_tensor, &out_meta);
40+
fill_tensor_meta_broadcast(in_tensor, out_meta.ndim, &in_meta);
41+
if (out_tensor.nbytes !=
42+
static_cast<size_t>(out_meta.numel) * sizeof(float) ||
43+
in_tensor.nbytes != static_cast<size_t>(in_meta.numel) * sizeof(float)) {
44+
throw std::runtime_error(
45+
"expand_copy: non-fp32 operand (nbytes != numel*4)");
46+
}
47+
48+
uint32_t wg_size =
49+
utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX);
50+
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
51+
device, out_meta.numel, wg_size, "expand_copy");
52+
53+
WGPUConstantEntry wg_size_constant = {};
54+
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
55+
wg_size_constant.value = static_cast<double>(wg_size);
56+
57+
WGPUBuffer out_meta_buf =
58+
utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
59+
WGPUBuffer in_meta_buf =
60+
utils::make_uniform(device, &in_meta, sizeof(TensorMeta));
61+
graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta));
62+
63+
WGPUShaderSourceWGSL wgsl_desc = {};
64+
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
65+
wgsl_desc.code = {kExpandCopyWGSL, WGPU_STRLEN};
66+
WGPUShaderModuleDescriptor shader_desc = {};
67+
shader_desc.nextInChain = &wgsl_desc.chain;
68+
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
69+
70+
WGPUBindGroupLayoutEntry entries[4] = {};
71+
entries[0].binding = 0;
72+
entries[0].visibility = WGPUShaderStage_Compute;
73+
entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
74+
entries[1].binding = 1;
75+
entries[1].visibility = WGPUShaderStage_Compute;
76+
entries[1].buffer.type = WGPUBufferBindingType_Storage;
77+
entries[2].binding = 2;
78+
entries[2].visibility = WGPUShaderStage_Compute;
79+
entries[2].buffer.type = WGPUBufferBindingType_Uniform;
80+
entries[3].binding = 3;
81+
entries[3].visibility = WGPUShaderStage_Compute;
82+
entries[3].buffer.type = WGPUBufferBindingType_Uniform;
83+
84+
WGPUBindGroupLayoutDescriptor bgl_desc = {};
85+
bgl_desc.entryCount = 4;
86+
bgl_desc.entries = entries;
87+
WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
88+
89+
WGPUPipelineLayoutDescriptor pl_desc = {};
90+
pl_desc.bindGroupLayoutCount = 1;
91+
pl_desc.bindGroupLayouts = &bgl;
92+
WGPUPipelineLayout pipeline_layout =
93+
wgpuDeviceCreatePipelineLayout(device, &pl_desc);
94+
95+
WGPUComputePipelineDescriptor pipeline_desc = {};
96+
pipeline_desc.layout = pipeline_layout;
97+
pipeline_desc.compute.module = shader;
98+
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
99+
pipeline_desc.compute.constantCount = 1;
100+
pipeline_desc.compute.constants = &wg_size_constant;
101+
WGPUComputePipeline pipeline =
102+
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
103+
104+
WGPUBindGroupEntry bg_entries[4] = {};
105+
bg_entries[0].binding = 0;
106+
bg_entries[0].buffer = in_tensor.buffer;
107+
bg_entries[0].size = in_tensor.nbytes;
108+
bg_entries[1].binding = 1;
109+
bg_entries[1].buffer = out_tensor.buffer;
110+
bg_entries[1].size = out_tensor.nbytes;
111+
bg_entries[2].binding = 2;
112+
bg_entries[2].buffer = out_meta_buf;
113+
bg_entries[2].size = sizeof(TensorMeta);
114+
bg_entries[3].binding = 3;
115+
bg_entries[3].buffer = in_meta_buf;
116+
bg_entries[3].size = sizeof(TensorMeta);
117+
118+
WGPUBindGroupDescriptor bg_desc = {};
119+
bg_desc.layout = bgl;
120+
bg_desc.entryCount = 4;
121+
bg_desc.entries = bg_entries;
122+
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
123+
124+
graph.add_dispatch({pipeline, bind_group, workgroup_count});
125+
126+
wgpuShaderModuleRelease(shader);
127+
wgpuBindGroupLayoutRelease(bgl);
128+
wgpuPipelineLayoutRelease(pipeline_layout);
129+
wgpuBufferRelease(out_meta_buf);
130+
wgpuBufferRelease(in_meta_buf);
131+
}
132+
133+
} // namespace
134+
135+
WEBGPU_REGISTER_OPERATORS {
136+
WEBGPU_REGISTER_OP(aten.expand_copy.default, expand_copy_impl);
137+
}
138+
139+
} // namespace executorch::backends::webgpu

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
@group(0) @binding(0) var<storage, read> input: array<f32>;
2+
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
3+
4+
struct TensorMeta {
5+
ndim: u32,
6+
numel: u32,
7+
sizes: vec4<u32>,
8+
strides: vec4<u32>,
9+
}
10+
@group(0) @binding(2) var<uniform> out_meta: TensorMeta;
11+
@group(0) @binding(3) var<uniform> in_meta: TensorMeta;
12+
13+
override wg_size: u32 = 64u;
14+
15+
@compute @workgroup_size(wg_size, 1, 1)
16+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
17+
let idx = gid.x;
18+
if (idx >= out_meta.numel) {
19+
return;
20+
}
21+
var rem = idx;
22+
var l: u32 = 0u;
23+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
24+
let coord = rem / out_meta.strides[d];
25+
rem = rem % out_meta.strides[d];
26+
l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d];
27+
}
28+
output[idx] = input[l];
29+
}

‎backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
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 expand_copy.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: ad996b7fd6eca5c5773715af3a9a117da83ef522042eb0ee918a623096417815
17+
inline constexpr const char* kExpandCopyWGSL = R"(
18+
@group(0) @binding(0) var<storage, read> input: array<f32>;
19+
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
20+
21+
struct TensorMeta {
22+
ndim: u32,
23+
numel: u32,
24+
sizes: vec4<u32>,
25+
strides: vec4<u32>,
26+
}
27+
@group(0) @binding(2) var<uniform> out_meta: TensorMeta;
28+
@group(0) @binding(3) var<uniform> in_meta: TensorMeta;
29+
30+
override wg_size: u32 = 64u;
31+
32+
@compute @workgroup_size(wg_size, 1, 1)
33+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
34+
let idx = gid.x;
35+
if (idx >= out_meta.numel) {
36+
return;
37+
}
38+
var rem = idx;
39+
var l: u32 = 0u;
40+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
41+
let coord = rem / out_meta.strides[d];
42+
rem = rem % out_meta.strides[d];
43+
l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d];
44+
}
45+
output[idx] = input[l];
46+
}
47+
)";
48+
49+
inline constexpr uint32_t kExpandCopyWorkgroupSizeX = 64;
50+
inline constexpr uint32_t kExpandCopyWorkgroupSizeY = 1;
51+
inline constexpr uint32_t kExpandCopyWorkgroupSizeZ = 1;
52+
53+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)