[ExecuTorch][WebGPU] Add aten.sub.Tensor with TensorMeta broadcast#20985
[ExecuTorch][WebGPU] Add aten.sub.Tensor with TensorMeta broadcast#20985JCNTH wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20985
Note: Links to docs will display an error until the docs builds have been completed. ❌ 128 New Failures, 3 Unrelated Failures, 2 Unclassified FailuresAs of commit 92ff4d3 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.sub.Tensorwith TensorMeta broadcastProblem: The WebGPU backend had no
subop. A flat modulo-tiled broadcast (the old rigaddform) is wrong for a middle/spatial broadcast[N,C,1,1]— e.g. InstanceNorm'sx - mean, wheremeanis[N,C,1,1]broadcast acrossH,W. A flat tiling cannot express that middle/spatial broadcast, so the result is numerically wrong.Solution: Add
aten.sub.Tensormirroringmul's 3xTensorMetaUBO broadcast — the same-shape elementwise fast path plus an out-index relinearize on shape mismatch — with analpha_pipelineoverrideconstant so the kernel computesin1 - alpha_ * in2.Implementation:
runtime/ops/sub/binary_sub.wgsl+BinaryOp.cppregisteraten.sub.Tensor. The handler fills three broadcastTensorMetauniforms (out/in1/in2) and rebuilds them plus the dispatch on dynamic resize;alphais read from the scalar arg and baked into the pipeline as thealpha_override (fixed for the op's lifetime, so resize never rewrites it). Mirrors Vulkanaten.sub.Tensor/ themulop.Constraints: fp32-only (every operand
nbytes == numel * 4), rank <= 4 (NCHW).Co-authored-with: Claude Code.
Differential Revision: D112378933