Skip to content

[ExecuTorch][WebGPU] Add constant_pad_nd op#20857

Open
JCNTH wants to merge 3 commits into
gh/JCNTH/32/basefrom
gh/JCNTH/32/head
Open

[ExecuTorch][WebGPU] Add constant_pad_nd op#20857
JCNTH wants to merge 3 commits into
gh/JCNTH/32/basefrom
gh/JCNTH/32/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Adds aten.constant_pad_nd.default to the WebGPU backend, unblocking the DaViT window-padding path in vision models.

Problem — the backend had no constant_pad_nd handler, so any graph that pads a tensor (e.g. DaViT's window partitioning) could not fully delegate to WebGPU and threw at runtime.

Solution — a single gather/fill compute kernel: Before — no handler; aten.constant_pad_nd.default unsupported at runtime. After — one thread per output element gathers the source element when its coordinates land inside the input, otherwise writes the constant fill value.

Implementation:

  • The handler right-aligns the (rank 1..4) dims into fixed vec4<u32> params (out_dims, in_dims, left); leading slots get extent 1 / pad 0, so the WGSL is rank-agnostic and always iterates 4 dims.
  • The pad IntList is reversed-dim (innermost-first (left, right) pairs); the handler expands it to per-dim left/right, then validates out.dims[d] == in.dims[d] + left[d] + right[d] before any buffer allocation (loud-fail, no leak-on-throw).
  • The kernel decodes each output element's 4D coords (last dim fastest), subtracts each dim's left pad as an unsigned wrapping subtract (a negative coord wraps to a huge value and is rejected by the < in_dims bound check); if all four coords are in-bounds it copies inp[flat_in], else it writes value — a pure copy/fill, so bit-exact.
  • The fill value is read via utils::scalar_or (a Scalar may serialize as Int or Double), defaulting to 0.
  • 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 workgroup-count ceiling; the stride_x override lets the shader decode i = gid.y*stride_x + gid.x).
  • Mirrors Vulkan backends/vulkan/runtime/graph/ops/impl/Pad.cpp (same reversed-dim (before, after) pad convention and constant_pad_nd resize logic).

Constraints — fp32 only (nbytes == numel*4 guard); rank 1..4; pad must be even-length and no longer than the rank; the output element count must fit u32 (<= 2^32).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110836672

Differential Revision: D110836672

[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/20857

Note: Links to docs will display an error until the docs builds have been completed.

❌ 129 New Failures, 2 Unrelated Failures, 7 Unclassified Failures

As of commit 2183251 with merge base c2b273e (image):

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:

  • pull / unittest / windows / windows-job (gh) (matched win rule in flaky-rules.json)
    Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under 'C:\actions-runner\_work\executorch\executorch\test-infra\.github\actions\teardown-windows'. Did you forget to run actions/checkout before running your local action?
  • pull / unittest-editable / windows / windows-job (gh) (matched win rule in flaky-rules.json)
    Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under 'C:\actions-runner\_work\executorch\executorch\test-infra\.github\actions\teardown-windows'. Did you forget to run actions/checkout before running your local action?

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.

[ghstack-poisoned]
[ghstack-poisoned]
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. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants