Skip to content

[ExecuTorch][WebGPU] Add upsample_nearest2d op#20859

Open
JCNTH wants to merge 1 commit into
gh/JCNTH/34/basefrom
gh/JCNTH/34/head
Open

[ExecuTorch][WebGPU] Add upsample_nearest2d op#20859
JCNTH wants to merge 1 commit into
gh/JCNTH/34/basefrom
gh/JCNTH/34/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 10, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Adds aten.upsample_nearest2d.vec to 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.vec unsupported. After — one thread per output element (n, c, oh, ow) maps back to its nearest source pixel and copies it.

Implementation:

  • The kernel computes the source index with ATen's legacy nearest formula ih = floor(oh*IH/OH), iw = floor(ow*IW/OW) (integer division), then copies inp[((n*C+c)*IH+ih)*IW+iw] — a plain NCHW row-major gather, bit-exact.
  • OH/OW are taken from the output tensor's own dims (not re-derived from the size/scale args), and the handler validates that N/C match between input and output.
  • Adaptive 1D->2D dispatch via utils::compute_dispatch_grid (workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 ceiling; stride_x decode).
  • Divergence from Vulkan (intentional): this does NOT mirror Vulkan 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 ATen nearest_neighbor_compute_source_index (UpSample.h) floor(oh*IH/OH) form is the one that matches PyTorch's mode="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's non_2x_ratio case pins down.

Constraints — fp32 only; 4D NCHW input and output; nearest mode only; non-zero spatial dims; the output element count must fit u32.

Co-authored-with: Claude Code.

Differential Revision: D110836668

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 10, 2026

Copy link
Copy Markdown

🔗 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 9075bf2 with merge base aceeb40 (image):

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.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants