[ExecuTorch][WebGPU] Add split_with_sizes_copy op#20995
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20995
Note: Links to docs will display an error until the docs builds have been completed. ❌ 128 New Failures, 3 Unrelated Failures, 7 Unclassified FailuresAs of commit d53a8f6 with merge base c2b273e ( NEW FAILURES - The following jobs have failed:
UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
Stack from ghstack (oldest at bottom):
Add
aten.split_with_sizes_copy.defaultto 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
selfalongdiminto N contiguous chunks. Each chunk is a step-1 slice from the running offset, reusing theslice.wgslgather 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.cppregisteringaten.split_with_sizes_copy.default; reusesslice_wgsl.h(no new shader). Usesutils::make_compute_pipeline(auto-derived bind-group layout) rather than hand-rolling the layout / pipeline / bind group.split_with_sizes_copydelegate.WEBGPU_SRCSentry.Constraints: fp32-only;
dimnormalized + range-checked;outputs == sizescount enforced (fail-loud). Reuses thesliceop'sslice_wgsl.h, so this diff stacks abovesliceand must land after it. No change to existing ops.Differential Revision: D112417284