Skip to content

Commit 26d3458

Browse files
committed
[ExecuTorch][WebGPU] Dynamic resize hook for view_copy
Pull Request resolved: #20579 **Make `view_copy` track the live sequence length under dynamic shapes.** **Problem:** `view_copy` lowers to a flat DMA buffer copy (`add_buffer_copy`) sized at the build-time max shape. With one dynamic graph serving any seq-len S (prefill S=K, decode S=1), the copy moved the full max-S byte count and the output kept its max dims, so a downstream consumer read a live shape that was too large. **Solution:** register a tensor resize hook on the input so the copy follows the live input numel (a view preserves numel). - Before: `copy_nbytes` and the output dims are fixed at the serialized max. - After: the hook recomputes the live numel from `cur_dims(in)`, scales the single dynamic output dim to preserve numel, sets the output `cur_dims`, and rewrites the Copy dispatch's `copy_nbytes`. **Implementation:** - Keep the existing DMA path (`Kind::Copy`); the hook only rewrites `copy_nbytes` via `dispatch_at`, no new kernel. - Handle the aliased in/out fast path (no copy emitted) by still setting the output `cur_dims` so the resize cascade reaches consumers. - Mirrors Vulkan's `view_buffer` contiguous fast path; numel-preserving like the other dynamic-shape op hooks. **Constraints:** inert on a static graph (`cur_dims == dims`), so byte-identical to the prior behavior; fp32-only and numel-preserving invariants unchanged. Co-authored-with: Claude Code. ghstack-source-id: 399812833 @exported-using-ghexport Differential Revision: [D109906098](https://our.internmc.facebook.com/intern/diff/D109906098/)
1 parent 03fd3ca commit 26d3458

1 file changed

Lines changed: 43 additions & 4 deletions

File tree

backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88

99
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
10+
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
1011
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
1112
#include <executorch/backends/webgpu/runtime/ops/view_copy/view_copy.h>
1213

@@ -38,11 +39,49 @@ void add_flat_copy(WebGPUGraph& graph, int in_id, int out_id) {
3839
}
3940

4041
// Aliased in/out already in place; CopyBufferToBuffer rejects src == dst.
41-
if (in_tensor.buffer == out_tensor.buffer) {
42-
return;
43-
}
42+
const bool aliased = in_tensor.buffer == out_tensor.buffer;
43+
const size_t dispatch_idx = aliased
44+
? 0
45+
: graph.add_buffer_copy(
46+
in_tensor.buffer, out_tensor.buffer, out_tensor.nbytes);
4447

45-
graph.add_buffer_copy(in_tensor.buffer, out_tensor.buffer, out_tensor.nbytes);
48+
// Dynamic shapes: view preserves numel; copy_nbytes + out dims track live in.
49+
std::vector<int64_t> out_max = out_tensor.dims;
50+
graph.add_tensor_resize_hook(
51+
in_id, [in_id, out_id, out_max, dispatch_idx, aliased](WebGPUGraph& g) {
52+
const uint64_t target = utils::numel_of(g.cur_dims(in_id));
53+
std::vector<int64_t> od = out_max;
54+
const uint64_t maxnumel = utils::numel_of(out_max);
55+
if (maxnumel != target) {
56+
bool resolved = false;
57+
// Assumes one dynamic dim; picks the leftmost numel-divisible.
58+
for (size_t d = 0; d < od.size(); d++) {
59+
if (out_max[d] <= 0) {
60+
continue;
61+
}
62+
const uint64_t rest = maxnumel / static_cast<uint64_t>(out_max[d]);
63+
if (rest != 0 && target % rest == 0) {
64+
const uint64_t nd = target / rest;
65+
if (nd <= static_cast<uint64_t>(out_max[d])) {
66+
od[d] = static_cast<int64_t>(nd);
67+
resolved = true;
68+
break;
69+
}
70+
}
71+
}
72+
// Fail loud: a silent miss would leave od at max while copy_nbytes
73+
// shrinks to the live size, desyncing consumers from the real copy.
74+
if (!resolved) {
75+
throw std::runtime_error(
76+
"view_copy(resize): could not resolve live output shape");
77+
}
78+
}
79+
g.set_cur_dims(out_id, od);
80+
if (!aliased) {
81+
g.dispatch_at(dispatch_idx).copy_nbytes =
82+
static_cast<size_t>(target) * sizeof(float);
83+
}
84+
});
4685
}
4786

4887
namespace {

0 commit comments

Comments
 (0)