Skip to content

Commit cd32734

Browse files
committed
[ExecuTorch][WebGPU] Dynamic resize hook for slice_copy (dynamic start/end)
Pull Request resolved: #20581 **Make `slice_copy` support a dynamic gather range so the RoPE-freqs slice `[input_pos : input_pos + S]` works under one dynamic graph.** **Problem:** the static slice handler read `start` via a scalar reader that throws on a SymInt and ignored `end` (output length baked AOT). The RoPE-freqs slice uses a SymInt `input_pos` for start and a live S for the range, so the static op could neither build nor resize for it. **Solution:** read start/end as possibly-dynamic SymInts and add a resize hook that recomputes the gather offset and live output length each step. - Before: `start` is a static scalar (SymInt throws); `end` ignored; output length fixed at the serialized max. - After: `start`/`end` read via a SymInt-aware reader; a hook recomputes `out[dim] = (end - start + step - 1) / step`, rewrites `out_meta`/`in_meta`/`params` UBOs + the dispatch count, and sets the output `cur_dims`. **Implementation:** - Hook registered on the `start`/`end` value-ids when they are SymInts and on the input tensor always (inert until resized, so a static slice is byte-identical). - Output/input `TensorMeta` rebuilt from live dims; `dim`/`step` stay static. - Keep the uniforms alive via `own_uniform_buffer` so the hook can rewrite them. - Mirrors Vulkan `resize_slice_copy_node`. **Constraints:** fp32-only; `dim`/`step` static; numerics + layout unchanged; inert on a static graph. NOTE (stacking): this diff sits on top of the in-review `slice_copy` op (D108793168); rebase onto it once that op lands on master. Co-authored-with: Claude Code. ghstack-source-id: 399812835 @exported-using-ghexport Differential Revision: [D109906092](https://our.internmc.facebook.com/intern/diff/D109906092/)
1 parent f2a3486 commit cd32734

1 file changed

Lines changed: 92 additions & 12 deletions

File tree

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

Lines changed: 92 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,40 @@ read_scalar(WebGPUGraph& graph, int id, int64_t dflt, const char* what) {
4646
}
4747
}
4848

49+
// Read a slice index (start/end) that MAY be a dynamic SymInt; else Int/Null.
50+
int64_t read_index(WebGPUGraph& graph, int id, int64_t dflt) {
51+
switch (graph.get_value_type(id)) {
52+
case WebGPUGraph::ValueType::SymInt:
53+
return graph.read_symint(id);
54+
case WebGPUGraph::ValueType::Int: {
55+
const int64_t v = graph.get_int(id);
56+
return v == INT64_MAX ? dflt : v;
57+
}
58+
case WebGPUGraph::ValueType::Null:
59+
return dflt;
60+
default:
61+
throw std::runtime_error("slice: dynamic/unsupported start/end index");
62+
}
63+
}
64+
65+
bool is_symint(WebGPUGraph& graph, int id) {
66+
return graph.get_value_type(id) == WebGPUGraph::ValueType::SymInt;
67+
}
68+
69+
// Clamp + normalize a (possibly negative) index into [0, size].
70+
int64_t norm_clamp(int64_t idx, int64_t size) {
71+
if (idx < 0) {
72+
idx += size;
73+
}
74+
return idx < 0 ? 0 : (idx > size ? size : idx);
75+
}
76+
4977
void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
50-
// args: [self, dim, start, end, step, out]; end unread (out shape is AOT).
78+
// args: [self, dim, start, end, step, out]. start/end may be dynamic SymInts;
79+
// a resize hook recomputes the live extent on `dim` (out[dim] / cur_dims).
5180
const int in_id = args.at(0);
81+
const int start_id = args.at(2);
82+
const int end_id = args.at(3);
5283
const int out_id = args.at(5);
5384

5485
WGPUDevice device = graph.device();
@@ -63,17 +94,14 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
6394
if (dim < 0 || dim >= in_ndim) {
6495
throw std::runtime_error("slice: dim out of range");
6596
}
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);
7397
const int64_t step = read_scalar(graph, args.at(4), 1, "step");
7498
if (step < 1) {
7599
throw std::runtime_error("slice: step must be >= 1");
76100
}
101+
// start/end may be dynamic SymInts; seed from current (max) dims, the resize
102+
// hook recomputes live. Clamp guards the gather offset.
103+
const int64_t in_size = in_tensor.dims[dim];
104+
const int64_t start = norm_clamp(read_index(graph, start_id, 0), in_size);
77105

78106
TensorMeta out_meta;
79107
TensorMeta in_meta;
@@ -175,14 +203,66 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
175203
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
176204

177205
graph.add_dispatch({pipeline, bind_group, workgroup_count});
206+
const size_t dispatch_idx = graph.num_dispatches() - 1;
207+
208+
// Dynamic shapes: live start/end -> out[dim] len + meta/params/dispatch.
209+
auto recompute = [in_id,
210+
out_id,
211+
start_id,
212+
end_id,
213+
dim,
214+
step,
215+
wg_size,
216+
out_meta_buf,
217+
in_meta_buf,
218+
params_buf,
219+
dispatch_idx](WebGPUGraph& g) {
220+
const auto& in_dims = g.cur_dims(in_id);
221+
const int64_t live_in_size = in_dims[dim];
222+
const int64_t start = norm_clamp(read_index(g, start_id, 0), live_in_size);
223+
const int64_t end =
224+
norm_clamp(read_index(g, end_id, live_in_size), live_in_size);
225+
const int64_t len = end > start ? (end - start + step - 1) / step : 0;
226+
227+
// Out dims = live input dims (mirror Vulkan resize_slice_copy_node).
228+
std::vector<int64_t> od = in_dims;
229+
od[dim] = len;
230+
g.set_cur_dims(out_id, od);
231+
232+
WebGPUTensor t_out;
233+
t_out.dims = od;
234+
WebGPUTensor t_in;
235+
t_in.dims = in_dims;
236+
TensorMeta om;
237+
TensorMeta im;
238+
fill_tensor_meta(t_out, &om);
239+
fill_tensor_meta(t_in, &im);
240+
wgpuQueueWriteBuffer(g.queue(), out_meta_buf, 0, &om, sizeof(om));
241+
wgpuQueueWriteBuffer(g.queue(), in_meta_buf, 0, &im, sizeof(im));
242+
SliceParams p = {};
243+
p.dim = static_cast<uint32_t>(dim);
244+
p.start = static_cast<uint32_t>(start);
245+
p.step = static_cast<uint32_t>(step);
246+
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
247+
g.dispatch_at(dispatch_idx).workgroup_count_x =
248+
utils::compute_1d_workgroup_count(
249+
g.device(), om.numel, wg_size, "slice(resize)");
250+
};
251+
if (is_symint(graph, start_id)) {
252+
graph.add_resize_hook(start_id, recompute);
253+
}
254+
if (is_symint(graph, end_id) && end_id != start_id) {
255+
graph.add_resize_hook(end_id, recompute);
256+
}
257+
graph.add_tensor_resize_hook(in_id, recompute);
178258

179259
wgpuShaderModuleRelease(shader);
180260
wgpuBindGroupLayoutRelease(bgl);
181261
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);
262+
// Graph owns the uniforms so the resize hook can rewrite them; freed in dtor.
263+
graph.own_uniform_buffer(out_meta_buf);
264+
graph.own_uniform_buffer(in_meta_buf);
265+
graph.own_uniform_buffer(params_buf);
186266
}
187267

188268
} // namespace

0 commit comments

Comments
 (0)