Skip to content

Commit 284d8a4

Browse files
JCNTHJCNTH
authored
[ExecuTorch][WebGPU] Add slice_copy op (aten.slice_copy.Tensor) (#20394)
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * __->__ #20394 * #20393 * #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Adds `aten.slice_copy.Tensor` to the WebGPU delegate as a gather: each output element is mapped back to its source input element along the sliced dim via `start + coord * step`. Composition (single compute dispatch): - `runtime/ops/slice/Slice.cpp` — reads `args = [self, dim, start, end, step, out]` via `read_scalar` (static `Int`/`Null`-sentinel default; throws on dynamic `SymInt`); normalizes negative `dim`/`start`, clamps `start` to `[0, in_size]`; builds two `TensorMeta` UBOs + a `SliceParams{dim, start, step}` uniform; guards fp32; dispatches over `compute_1d_workgroup_count(out.numel)` with `override wg_size`; releases all uniforms after the bind group. - `runtime/ops/slice/slice.wgsl` — delinearizes the output index over the contiguous output strides, maps the sliced-dim coordinate back to the input (`start + coord*step`), relinearizes over the input strides. @exported-using-ghexport Differential Revision: [D108793168](https://our.internmc.facebook.com/intern/diff/D108793168/) Differential Revision: [D108793168](https://our.internmc.facebook.com/intern/diff/D108793168)
1 parent 7d179a6 commit 284d8a4

4 files changed

Lines changed: 303 additions & 0 deletions

File tree

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ set(WEBGPU_SRCS
4747
runtime/ops/sigmoid/UnaryOp.cpp
4848
runtime/ops/squeeze/Squeeze.cpp
4949
runtime/ops/unsqueeze/Unsqueeze.cpp
50+
runtime/ops/slice/Slice.cpp
5051
)
5152

5253
add_library(webgpu_backend ${WEBGPU_SRCS})

‎backends/webgpu/runtime/ops/slice/Slice.cpp‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
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/slice/slice_wgsl.h>
14+
15+
#include <webgpu/webgpu.h>
16+
17+
#include <cstdint>
18+
#include <stdexcept>
19+
#include <string>
20+
#include <vector>
21+
22+
namespace executorch::backends::webgpu {
23+
24+
namespace {
25+
26+
struct SliceParams {
27+
uint32_t dim;
28+
uint32_t start;
29+
uint32_t step;
30+
uint32_t _pad;
31+
};
32+
33+
// Read scalar arg: Int->value (INT64_MAX->default), Null->default, else throw.
34+
int64_t
35+
read_scalar(WebGPUGraph& graph, int id, int64_t dflt, const char* what) {
36+
switch (graph.get_value_type(id)) {
37+
case WebGPUGraph::ValueType::Int: {
38+
const int64_t v = graph.get_int(id);
39+
return v == INT64_MAX ? dflt : v;
40+
}
41+
case WebGPUGraph::ValueType::Null:
42+
return dflt;
43+
default:
44+
throw std::runtime_error(
45+
std::string("slice: dynamic/unsupported ") + what);
46+
}
47+
}
48+
49+
void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
50+
// args: [self, dim, start, end, step, out]; end unread (out shape is AOT).
51+
const int in_id = args.at(0);
52+
const int out_id = args.at(5);
53+
54+
WGPUDevice device = graph.device();
55+
const auto& in_tensor = graph.get_tensor(in_id);
56+
const auto& out_tensor = graph.get_tensor(out_id);
57+
58+
const int in_ndim = static_cast<int>(in_tensor.dims.size());
59+
int64_t dim = read_scalar(graph, args.at(1), 0, "dim");
60+
if (dim < 0) {
61+
dim += in_ndim;
62+
}
63+
if (dim < 0 || dim >= in_ndim) {
64+
throw std::runtime_error("slice: dim out of range");
65+
}
66+
const int64_t in_size = in_tensor.dims[dim];
67+
int64_t start = read_scalar(graph, args.at(2), 0, "start");
68+
if (start < 0) {
69+
start += in_size;
70+
}
71+
// Clamp start to [0, in_size] (guards the gather offset; out size is AOT).
72+
start = start < 0 ? 0 : (start > in_size ? in_size : start);
73+
const int64_t step = read_scalar(graph, args.at(4), 1, "step");
74+
if (step < 1) {
75+
throw std::runtime_error("slice: step must be >= 1");
76+
}
77+
78+
TensorMeta out_meta;
79+
TensorMeta in_meta;
80+
fill_tensor_meta(out_tensor, &out_meta);
81+
fill_tensor_meta(in_tensor, &in_meta);
82+
if (out_tensor.nbytes !=
83+
static_cast<size_t>(out_meta.numel) * sizeof(float) ||
84+
in_tensor.nbytes != static_cast<size_t>(in_meta.numel) * sizeof(float)) {
85+
throw std::runtime_error("slice: non-fp32 operand (nbytes != numel * 4)");
86+
}
87+
88+
SliceParams params = {};
89+
params.dim = static_cast<uint32_t>(dim);
90+
params.start = static_cast<uint32_t>(start);
91+
params.step = static_cast<uint32_t>(step);
92+
93+
uint32_t wg_size = utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX);
94+
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
95+
device, out_meta.numel, wg_size, "slice");
96+
97+
WGPUConstantEntry wg_size_constant = {};
98+
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
99+
wg_size_constant.value = static_cast<double>(wg_size);
100+
101+
WGPUBuffer out_meta_buf =
102+
utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
103+
WGPUBuffer in_meta_buf =
104+
utils::make_uniform(device, &in_meta, sizeof(TensorMeta));
105+
WGPUBuffer params_buf =
106+
utils::make_uniform(device, &params, sizeof(SliceParams));
107+
graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(SliceParams));
108+
109+
WGPUShaderSourceWGSL wgsl_desc = {};
110+
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
111+
wgsl_desc.code = {kSliceWGSL, WGPU_STRLEN};
112+
WGPUShaderModuleDescriptor shader_desc = {};
113+
shader_desc.nextInChain = &wgsl_desc.chain;
114+
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
115+
116+
// Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms).
117+
WGPUBindGroupLayoutEntry entries[5] = {};
118+
entries[0].binding = 0;
119+
entries[0].visibility = WGPUShaderStage_Compute;
120+
entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
121+
entries[1].binding = 1;
122+
entries[1].visibility = WGPUShaderStage_Compute;
123+
entries[1].buffer.type = WGPUBufferBindingType_Storage;
124+
entries[2].binding = 2;
125+
entries[2].visibility = WGPUShaderStage_Compute;
126+
entries[2].buffer.type = WGPUBufferBindingType_Uniform;
127+
entries[3].binding = 3;
128+
entries[3].visibility = WGPUShaderStage_Compute;
129+
entries[3].buffer.type = WGPUBufferBindingType_Uniform;
130+
entries[4].binding = 4;
131+
entries[4].visibility = WGPUShaderStage_Compute;
132+
entries[4].buffer.type = WGPUBufferBindingType_Uniform;
133+
134+
WGPUBindGroupLayoutDescriptor bgl_desc = {};
135+
bgl_desc.entryCount = 5;
136+
bgl_desc.entries = entries;
137+
WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
138+
139+
WGPUPipelineLayoutDescriptor pl_desc = {};
140+
pl_desc.bindGroupLayoutCount = 1;
141+
pl_desc.bindGroupLayouts = &bgl;
142+
WGPUPipelineLayout pipeline_layout =
143+
wgpuDeviceCreatePipelineLayout(device, &pl_desc);
144+
145+
WGPUComputePipelineDescriptor pipeline_desc = {};
146+
pipeline_desc.layout = pipeline_layout;
147+
pipeline_desc.compute.module = shader;
148+
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
149+
pipeline_desc.compute.constantCount = 1;
150+
pipeline_desc.compute.constants = &wg_size_constant;
151+
WGPUComputePipeline pipeline =
152+
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
153+
154+
WGPUBindGroupEntry bg_entries[5] = {};
155+
bg_entries[0].binding = 0;
156+
bg_entries[0].buffer = in_tensor.buffer;
157+
bg_entries[0].size = in_tensor.nbytes;
158+
bg_entries[1].binding = 1;
159+
bg_entries[1].buffer = out_tensor.buffer;
160+
bg_entries[1].size = out_tensor.nbytes;
161+
bg_entries[2].binding = 2;
162+
bg_entries[2].buffer = out_meta_buf;
163+
bg_entries[2].size = sizeof(TensorMeta);
164+
bg_entries[3].binding = 3;
165+
bg_entries[3].buffer = in_meta_buf;
166+
bg_entries[3].size = sizeof(TensorMeta);
167+
bg_entries[4].binding = 4;
168+
bg_entries[4].buffer = params_buf;
169+
bg_entries[4].size = sizeof(SliceParams);
170+
171+
WGPUBindGroupDescriptor bg_desc = {};
172+
bg_desc.layout = bgl;
173+
bg_desc.entryCount = 5;
174+
bg_desc.entries = bg_entries;
175+
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
176+
177+
graph.add_dispatch({pipeline, bind_group, workgroup_count});
178+
179+
wgpuShaderModuleRelease(shader);
180+
wgpuBindGroupLayoutRelease(bgl);
181+
wgpuPipelineLayoutRelease(pipeline_layout);
182+
// Drop our refs; the bind group keeps the uniforms alive until release.
183+
wgpuBufferRelease(out_meta_buf);
184+
wgpuBufferRelease(in_meta_buf);
185+
wgpuBufferRelease(params_buf);
186+
}
187+
188+
} // namespace
189+
190+
WEBGPU_REGISTER_OPERATORS {
191+
WEBGPU_REGISTER_OP(aten.slice_copy.Tensor, slice_impl);
192+
}
193+
194+
} // namespace executorch::backends::webgpu

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

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
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+
struct Params {
14+
dim: u32,
15+
start: u32,
16+
step: u32,
17+
}
18+
@group(0) @binding(4) var<uniform> params: Params;
19+
20+
override wg_size: u32 = 64u;
21+
22+
@compute @workgroup_size(wg_size, 1, 1)
23+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
24+
let out_bufi = gid.x;
25+
if (out_bufi >= out_meta.numel) {
26+
return;
27+
}
28+
29+
// Gather: out_bufi -> in_bufi, sliced dim coord = start + coord*step.
30+
var rem = out_bufi;
31+
var in_bufi: u32 = 0u;
32+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
33+
let coord = rem / out_meta.strides[d];
34+
rem = rem % out_meta.strides[d];
35+
var in_coord = coord;
36+
if (d == params.dim) {
37+
in_coord = params.start + coord * params.step;
38+
}
39+
in_bufi = in_bufi + in_coord * in_meta.strides[d];
40+
}
41+
output[out_bufi] = input[in_bufi];
42+
}

‎backends/webgpu/runtime/ops/slice/slice_wgsl.h‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
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 slice.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 6639d985420d43a67de0847749918ab6216e0785399bdcae737d49b81c773526
17+
inline constexpr const char* kSliceWGSL = 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+
struct Params {
31+
dim: u32,
32+
start: u32,
33+
step: u32,
34+
}
35+
@group(0) @binding(4) var<uniform> params: Params;
36+
37+
override wg_size: u32 = 64u;
38+
39+
@compute @workgroup_size(wg_size, 1, 1)
40+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
41+
let out_bufi = gid.x;
42+
if (out_bufi >= out_meta.numel) {
43+
return;
44+
}
45+
46+
// Gather: out_bufi -> in_bufi, sliced dim coord = start + coord*step.
47+
var rem = out_bufi;
48+
var in_bufi: u32 = 0u;
49+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
50+
let coord = rem / out_meta.strides[d];
51+
rem = rem % out_meta.strides[d];
52+
var in_coord = coord;
53+
if (d == params.dim) {
54+
in_coord = params.start + coord * params.step;
55+
}
56+
in_bufi = in_bufi + in_coord * in_meta.strides[d];
57+
}
58+
output[out_bufi] = input[in_bufi];
59+
}
60+
)";
61+
62+
inline constexpr uint32_t kSliceWorkgroupSizeX = 64;
63+
inline constexpr uint32_t kSliceWorkgroupSizeY = 1;
64+
inline constexpr uint32_t kSliceWorkgroupSizeZ = 1;
65+
66+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)