Skip to content

Commit 8cfa882

Browse files
JCNTHJCNTH
authored
[ExecuTorch][WebGPU] Add split_with_sizes_copy op (#20995)
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20868 * #20996 * __->__ #20995 * #20994 * #20993 * #20992 * #20991 * #20990 * #20989 * #20988 * #20987 * #20986 * #20876 * #20875 * #20874 * #20873 * #20872 * #20871 * #20866 * #20865 * #20864 * #20863 * #20862 * #20861 * #20860 * #20859 * #20858 * #20857 * #20856 * #20855 * #20854 * #20852 * #20851 * #20850 * #20849 * #20848 * #20846 * #20845 * #20844 * #20843 * #20842 **Add `aten.split_with_sizes_copy.default` to the WebGPU backend** — the split on YOLO's Detect head (separating the concatenated box / objectness / class predictions). **Problem**: The WebGPU delegate had no `split_with_sizes_copy`, so YOLO object-detection could not fully delegate to the GPU. **Solution**: Split `self` along `dim` into N contiguous chunks. Each chunk is a step-1 slice from the running offset, reusing the `slice.wgsl` gather kernel — one dispatch per output. Outputs arrive as a serialized ValueList. Each chunk writes its own distinct output buffer and reads only the shared input, so there is no cross-dispatch read-after-write hazard. **Implementation**: - `runtime/ops/split_with_sizes/SplitWithSizes.cpp` registering `aten.split_with_sizes_copy.default`; reuses `slice_wgsl.h` (no new shader). Uses `utils::make_compute_pipeline` (auto-derived bind-group layout) rather than hand-rolling the layout / pipeline / bind group. - Mirrors the Vulkan `split_with_sizes_copy` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only; `dim` normalized + range-checked; `outputs == sizes` count enforced (fail-loud). Reuses the `slice` op's `slice_wgsl.h`, so this diff stacks above `slice` and must land after it. No change to existing ops. @exported-using-ghexport Differential Revision: [D112417284](https://our.internmc.facebook.com/intern/diff/D112417284/) Differential Revision: [D112417284](https://our.internmc.facebook.com/intern/diff/D112417284)
1 parent e731b41 commit 8cfa882

1 file changed

Lines changed: 147 additions & 0 deletions

File tree

‎backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
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 <vector>
20+
21+
namespace executorch::backends::webgpu {
22+
23+
namespace {
24+
25+
struct SliceParams {
26+
uint32_t dim;
27+
uint32_t start;
28+
uint32_t step;
29+
uint32_t _pad;
30+
};
31+
32+
// aten.split_with_sizes_copy: N contiguous chunks via per-output slice gather.
33+
void split_with_sizes_impl(WebGPUGraph& graph, const std::vector<int>& args) {
34+
if (args.size() < 4) {
35+
throw std::runtime_error("WebGPU split_with_sizes: expected >=4 args");
36+
}
37+
const int in_id = args.at(0);
38+
const std::vector<int64_t>& sizes = graph.get_int_list(args.at(1));
39+
int64_t dim = graph.get_int(args.at(2));
40+
const std::vector<int>& outs = graph.get_value_list(args.at(args.size() - 1));
41+
42+
WGPUDevice device = graph.device();
43+
const auto& in_tensor = graph.get_tensor(in_id);
44+
const int in_ndim = static_cast<int>(in_tensor.dims.size());
45+
if (dim < 0) {
46+
dim += in_ndim;
47+
}
48+
if (dim < 0 || dim >= in_ndim) {
49+
throw std::runtime_error("split_with_sizes: dim out of range");
50+
}
51+
if (outs.size() != sizes.size()) {
52+
throw std::runtime_error("split_with_sizes: outputs != sizes count");
53+
}
54+
55+
// Validate the split contract up front (before any dispatch): each size is
56+
// non-negative and they sum to the split dim's extent -- a negative or
57+
// oversize value would otherwise wrap when cast to u32 into an out-of-bounds
58+
// gather offset.
59+
int64_t sizes_sum = 0;
60+
for (int64_t s : sizes) {
61+
if (s < 0) {
62+
throw std::runtime_error("split_with_sizes: negative split size");
63+
}
64+
sizes_sum += s;
65+
}
66+
if (sizes_sum != in_tensor.dims[dim]) {
67+
throw std::runtime_error(
68+
"split_with_sizes: sizes must sum to the split dim extent");
69+
}
70+
71+
TensorMeta in_meta;
72+
fill_tensor_meta(in_tensor, &in_meta);
73+
if (in_tensor.nbytes != static_cast<size_t>(in_meta.numel) * sizeof(float)) {
74+
throw std::runtime_error("split_with_sizes: non-fp32 input");
75+
}
76+
77+
uint32_t start = 0;
78+
for (size_t i = 0; i < outs.size(); i++) {
79+
const auto& out_tensor = graph.get_tensor(outs[i]);
80+
TensorMeta out_meta;
81+
fill_tensor_meta(out_tensor, &out_meta);
82+
if (out_tensor.nbytes !=
83+
static_cast<size_t>(out_meta.numel) * sizeof(float)) {
84+
throw std::runtime_error("split_with_sizes: non-fp32 output");
85+
}
86+
87+
SliceParams params = {};
88+
params.dim = static_cast<uint32_t>(dim);
89+
params.start = start;
90+
params.step = 1u;
91+
start += static_cast<uint32_t>(sizes[i]);
92+
93+
uint32_t wg_size =
94+
utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX);
95+
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
96+
device, out_meta.numel, wg_size, "split_with_sizes");
97+
98+
WGPUConstantEntry wg_size_constant = {};
99+
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
100+
wg_size_constant.value = static_cast<double>(wg_size);
101+
102+
WGPUBuffer out_meta_buf =
103+
utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
104+
WGPUBuffer in_meta_buf =
105+
utils::make_uniform(device, &in_meta, sizeof(TensorMeta));
106+
WGPUBuffer params_buf =
107+
utils::make_uniform(device, &params, sizeof(SliceParams));
108+
graph.add_uniform_buffer_bytes(
109+
2 * sizeof(TensorMeta) + sizeof(SliceParams));
110+
111+
utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
112+
device,
113+
kSliceWGSL,
114+
{
115+
{0,
116+
WGPUBufferBindingType_ReadOnlyStorage,
117+
in_tensor.buffer,
118+
in_tensor.nbytes},
119+
{1,
120+
WGPUBufferBindingType_Storage,
121+
out_tensor.buffer,
122+
out_tensor.nbytes},
123+
{2,
124+
WGPUBufferBindingType_Uniform,
125+
out_meta_buf,
126+
sizeof(TensorMeta)},
127+
{3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)},
128+
{4, WGPUBufferBindingType_Uniform, params_buf, sizeof(SliceParams)},
129+
},
130+
&wg_size_constant,
131+
1);
132+
133+
graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
134+
135+
graph.own_uniform_buffer(out_meta_buf);
136+
graph.own_uniform_buffer(in_meta_buf);
137+
graph.own_uniform_buffer(params_buf);
138+
}
139+
}
140+
141+
} // namespace
142+
143+
WEBGPU_REGISTER_OPERATORS {
144+
WEBGPU_REGISTER_OP(aten.split_with_sizes_copy.default, split_with_sizes_impl);
145+
}
146+
147+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)