[ExecuTorch][WebGPU] Add upsample_nearest2d op#20859
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20859
Note: Links to docs will display an error until the docs builds have been completed. ❌ 129 New Failures, 3 Unrelated Failures, 7 Unclassified FailuresAs of commit 2a0ccc6 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):
Adds
aten.upsample_nearest2d.vecto the WebGPU backend, enabling the SAM2/SAM3 pixel-decoder / FPN nearest-neighbour upsample path.Problem — nearest-neighbour 2D upsample (
F.interpolate(..., mode="nearest")) had no WebGPU kernel, blocking the FPN 2x upsample chain in the SAM2/SAM3 pixel decoder.Solution — Before — no handler;
aten.upsample_nearest2d.vecunsupported. After — one thread per output element(n, c, oh, ow)maps back to its nearest source pixel and copies it.Implementation:
ih = floor(oh*IH/OH),iw = floor(ow*IW/OW)(integer division), then copiesinp[((n*C+c)*IH+ih)*IW+iw]— a plain NCHW row-major gather, bit-exact.OH/OWare taken from the output tensor's own dims (not re-derived from thesize/scaleargs), and the handler validates thatN/Cmatch between input and output.utils::compute_dispatch_grid(workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 ceiling;stride_xdecode).backends/vulkan/runtime/graph/ops/impl/Upsample.cpp, which computes the source index with the half-pixel-center reciprocal-scale formula. Half-pixel-center is correct for bilinear but wrong for non-exact nearest; the ATennearest_neighbor_compute_source_index(UpSample.h)floor(oh*IH/OH)form is the one that matches PyTorch'smode="nearest". The two agree on exact-integer ratios (e.g. 2x) but diverge on non-integer ratios (e.g. 5->8), which the op-test'snon_2x_ratiocase pins down.Constraints — fp32 only; 4D NCHW input and output;
nearestmode only; non-zero spatial dims; the output element count must fitu32.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D110836668
Differential Revision: D110836668