Skip to content

[ExecuTorch][WebGPU] Add aten.index.Tensor (1D-self gather)#20464

Open
JulianCloudNTH wants to merge 2 commits into
gh/JulianCloudNTH/59/basefrom
gh/JulianCloudNTH/59/head
Open

[ExecuTorch][WebGPU] Add aten.index.Tensor (1D-self gather)#20464
JulianCloudNTH wants to merge 2 commits into
gh/JulianCloudNTH/59/basefrom
gh/JulianCloudNTH/59/head

Conversation

@JulianCloudNTH

@JulianCloudNTH JulianCloudNTH commented Jun 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Adds the WebGPU delegate handler for aten.index.Tensor, the 1D-self advanced-index
gather out[i] = self[index[i]] (output shape == index shape). This is the form the
VulkanPartitioner delegates -- it requires a 1D self and exactly one non-None index
(op_registry.py); 2D mask/freqs gathers stay on CPU. It mirrors the Vulkan delegate's
index_tensor op (IndexTensor.cpp + index_tensor_buffer.glsl) as a single compute
dispatch over the output elements, each reading the int32 index and gathering the
corresponding fp32 self element.

The op is composed as:

  • index.wgsl: one workgroup-strided pass, out[i] = self[u32(index[i])], guarded by a
    numel bound; buffer-only, fp32 self/out, int32 index, 1D dispatch via the shared
    WebGPUUtils helpers (clamp workgroup size + 1D count).
  • Index.cpp: validates the args (self/out tensors; indices ValueList with exactly one
    index tensor; fp32 self/out; int32 index; out numel == index numel), failing loud on
    any violation, then records the dispatch. row_width is dropped (always 1 for 1D self).

Differential Revision: D109478967

@JulianCloudNTH
Update
8ca983e 
[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jun 23, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20464

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

❌ 4 New Failures, 1 Pending, 1 Unrelated Failure

As of commit 83b260c with merge base e03f777 (image):

NEW FAILURES - The following jobs have failed:

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 was referenced Jun 23, 2026
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Open
Open
Open
Open
Open
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 23, 2026
@github-actions

github-actions Bot commented Jun 23, 2026

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.

@JulianCloudNTH
Update
83b260c 
[ghstack-poisoned]
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] Add select_copy op (aten.select_copy.int) (#20362)
01ffde5 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* #20394
* #20393
* #20392
* #20391
* #20390
* #20363
* __->__ #20362
* #20361
* #20360
* #20359



Adds `aten.select_copy.int` to the WebGPU delegate as a gather: picks a
fixed index along one dim, producing an output of rank (input rank - 1).

Composition (single dispatch):
- `select/Select.cpp` — reads `[self, dim, index, out]` (static `Int`
via `read_scalar`; throws on dynamic `SymInt`), normalizes +
bounds-checks dim/index, builds 2 `TensorMeta` UBOs + a
`SelectParams{dim,index}`, fp32 guard, 1D-dispatch over `numel`,
releases uniforms after the bind group.
- `select/select.wgsl` — seeds the input offset with `index *
in.strides[dim]`, delinearizes the output index, maps each out dim to
its in dim (shifted past the selected dim), relinearizes on input
strides.
@exported-using-ghexport

Differential Revision:
[D108793166](https://our.internmc.facebook.com/intern/diff/D108793166/)

Differential Revision:
[D108793166](https://our.internmc.facebook.com/intern/diff/D108793166)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] select_copy op test suite (cases.py op-test fram…
d1e327e 
…ework) (#20363)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* #20394
* #20393
* #20392
* #20391
* #20390
* __->__ #20363
* #20362
* #20361
* #20360
* #20359



Registers `aten.select_copy.int` in the `cases.py` op-test framework: a
`_select_suite` of 4 configs (leading/middle/last dim + negative index)
that `generate_op_tests` exports and compares to a torch golden on Dawn.
Also adds `test/ops/select/test_select.py` (`SelectModule` + `CONFIGS` +
an export-delegation/eager smoke test) and the `aten.select_copy.int`
partitioner-allowlist entry in `tester.py`.
@exported-using-ghexport

Differential Revision:
[D108793161](https://our.internmc.facebook.com/intern/diff/D108793161/)

Differential Revision:
[D108793161](https://our.internmc.facebook.com/intern/diff/D108793161)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] Add sigmoid op (aten.sigmoid.default) (#20390)
9bc3f97 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* #20394
* #20393
* #20392
* #20391
* __->__ #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Adds `aten.sigmoid.default` to the WebGPU delegate: element-wise
`1/(1+exp(-x))` over a flat fp32 buffer. On the Llama critical path
(`F.silu` -> `sigmoid` + `mul`).

Composition (single dispatch):
- `sigmoid/UnaryOp.cpp` — binds input (storage, read-only) + output
(storage) + a `Params{num_elements}` uniform, 1D-dispatches over
`num_elements` with `override wg_size` (clamped to the device limit);
mirrors the `add` op (uniform mapped-at-creation, released after the
bind group).
- `sigmoid/sigmoid.wgsl` — guards `idx >= num_elements` and writes the
logistic of each element.
@exported-using-ghexport

Differential Revision:
[D108793157](https://our.internmc.facebook.com/intern/diff/D108793157/)

Differential Revision:
[D108793157](https://our.internmc.facebook.com/intern/diff/D108793157)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] sigmoid op test suite (cases.py op-test framewor…
20b80f7 
…k) (#20391)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* #20394
* #20393
* #20392
* __->__ #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Registers `aten.sigmoid.default` in the `cases.py` op-test framework: a
`_sigmoid_suite` (hard-coded shapes + a saturation case over a
`linspace(-12, 12)` input) that `generate_op_tests` exports and compares
to an fp64 torch golden on Dawn. Also adds
`test/ops/sigmoid/test_sigmoid.py` (`SigmoidModule` + `N` + `_det_input`
+ an export-delegation/eager smoke test) and the `aten.sigmoid.default`
partitioner-allowlist entry in `tester.py`.
@exported-using-ghexport

Differential Revision:
[D108793159](https://our.internmc.facebook.com/intern/diff/D108793159/)

Differential Revision:
[D108793159](https://our.internmc.facebook.com/intern/diff/D108793159)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] Add squeeze_copy + unsqueeze_copy (flat copies) (#…
0ee3a5c 
…20392)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* #20394
* #20393
* __->__ #20392
* #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Adds `aten.squeeze_copy.dims` and `aten.unsqueeze_copy.default` to the
WebGPU delegate. Both are numel-preserving shape ops; on a dense
row-major buffer backend they are the same flat copy as `view_copy` —
only the shape metadata differs (mirrors the Vulkan delegate, which
routes both through `add_view_copy_node`).

Composition (no new kernel):
- `squeeze/Squeeze.cpp` — reads `args = [self, dims, out]`, ignores the
AOT-fixed `dims`, calls `add_flat_copy(graph, in, out)` from
`runtime/ops/view_copy/view_copy.h`.
- `unsqueeze/Unsqueeze.cpp` — reads `args = [self, dim, out]`, ignores
the AOT-fixed `dim`, calls `add_flat_copy(graph, in, out)`.
@exported-using-ghexport

Differential Revision:
[D108793153](https://our.internmc.facebook.com/intern/diff/D108793153/)

Differential Revision:
[D108793153](https://our.internmc.facebook.com/intern/diff/D108793153)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] squeeze_copy + unsqueeze_copy test suites (cases…
abfff1e 
….py op-test framework) (#20393)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* #20394
* __->__ #20393
* #20392
* #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Registers `aten.squeeze_copy.dims` and `aten.unsqueeze_copy.default` in
the `cases.py` op-test framework: a `_squeeze_suite` of 3 configs
(squeeze leading/middle/multiple size-1 dims) and a `_unsqueeze_suite`
of 3 configs (insert dim at front/middle/last) that `generate_op_tests`
exports via `VulkanPartitioner` and compares to a torch golden on Dawn.
Also adds `test/ops/squeeze/test_squeeze.py` (`SqueezeModule` +
`CONFIGS` + `_op_delegated` smoke test),
`test/ops/unsqueeze/test_unsqueeze.py` (`UnsqueezeModule` + `CONFIGS` +
`_op_delegated` smoke test), and the two partitioner-allowlist entries
in `tester.py`.
@exported-using-ghexport

Differential Revision:
[D108793152](https://our.internmc.facebook.com/intern/diff/D108793152/)

Differential Revision:
[D108793152](https://our.internmc.facebook.com/intern/diff/D108793152)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] Add slice_copy op (aten.slice_copy.Tensor) (#20394)
284d8a4 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* #20395
* __->__ #20394
* #20393
* #20392
* #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Adds `aten.slice_copy.Tensor` to the WebGPU delegate as a gather: each
output element is mapped back to its source input element along the
sliced dim via `start + coord * step`.

Composition (single compute dispatch):
- `runtime/ops/slice/Slice.cpp` — reads `args = [self, dim, start, end,
step, out]` via `read_scalar` (static `Int`/`Null`-sentinel default;
throws on dynamic `SymInt`); normalizes negative `dim`/`start`, clamps
`start` to `[0, in_size]`; builds two `TensorMeta` UBOs + a
`SliceParams{dim, start, step}` uniform; guards fp32; dispatches over
`compute_1d_workgroup_count(out.numel)` with `override wg_size`;
releases all uniforms after the bind group.
- `runtime/ops/slice/slice.wgsl` — delinearizes the output index over
the contiguous output strides, maps the sliced-dim coordinate back to
the input (`start + coord*step`), relinearizes over the input strides.
@exported-using-ghexport

Differential Revision:
[D108793168](https://our.internmc.facebook.com/intern/diff/D108793168/)

Differential Revision:
[D108793168](https://our.internmc.facebook.com/intern/diff/D108793168)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] slice_copy op test suite (cases.py op-test frame…
88eca88 
…work) (#20395)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* #20396
* __->__ #20395
* #20394
* #20393
* #20392
* #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Registers `aten.slice_copy.Tensor` in the `cases.py` op-test framework:
a `_slice_suite` of 4 configs (leading-dim slice `[:,1:5]`, last-dim
slice `[...,1:3]`, step-2 `[:,0:8:2]`, negative-end `[:,1:-1]`) that
`generate_op_tests` exports via `VulkanPartitioner` and compares to a
torch golden on Dawn. Also adds `test/ops/slice/test_slice.py`
(`SliceModule` + `CONFIGS` + export-delegation/eager smoke test) and the
`aten.slice_copy.Tensor` partitioner-allowlist entry in `tester.py`.
@exported-using-ghexport

Differential Revision:
[D108793151](https://our.internmc.facebook.com/intern/diff/D108793151/)

Differential Revision:
[D108793151](https://our.internmc.facebook.com/intern/diff/D108793151)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] Add permute_copy + IntList graph support (aten.p…
11567a3 
…ermute_copy.default) (#20396)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* #20397
* __->__ #20396
* #20395
* #20394
* #20393
* #20392
* #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Adds `aten.permute_copy.default` (a coordinate-reorder gather) to the
WebGPU delegate, and the `IntList` graph value type it needs to read its
`dims` argument.

Composition:
- `runtime/WebGPUGraph.{h,cpp}` — adds `ValueType::IntList` backed by
`std::vector<std::vector<int64_t>> int_lists_` + `get_int_list(int)`;
`build()` deserializes `vkgraph::GraphTypes::IntList` via
`value_as_IntList()->items()` (int64, matching the FlatBuffer `[long]`);
mirrors the existing scalar value plumbing.
- `runtime/ops/permute/Permute.cpp` — reads the permutation via
`get_int_list`, normalizes negative dims, validates it is a permutation
of `[0, ndim)`, builds two `TensorMeta` UBOs + a `PermuteParams{perm:
vec4<u32>}` uniform, guards fp32 + rank≤4, dispatches over
`compute_1d_workgroup_count(out.numel)` with `override wg_size`;
releases all uniforms after the bind group.
- `runtime/ops/permute/permute.wgsl` — delinearizes the output index
over the contiguous output strides, reads `input` at
`in.strides[perm[d]]` per dim (mirrors Vulkan `permute_buffer.glsl`).
- Registers both `aten.permute_copy.default` and `aten.permute.default`
to the same handler.
@exported-using-ghexport

Differential Revision:
[D108793162](https://our.internmc.facebook.com/intern/diff/D108793162/)

Differential Revision:
[D108793162](https://our.internmc.facebook.com/intern/diff/D108793162)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
@JulianCloudNTH
[ExecuTorch][WebGPU] permute_copy op test suite (cases.py op-test fra…
5668635 
…mework) (#20397)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20465
* #20464
* #20463
* #20435
* #20399
* #20398
* __->__ #20397
* #20396
* #20395
* #20394
* #20393
* #20392
* #20391
* #20390
* #20363
* #20362
* #20361
* #20360
* #20359



Registers `aten.permute_copy.default` in the `cases.py` op-test
framework: a `_permute_suite` of 4 configs (3D rotation, 4D middle-dim
transpose, 2D transpose, full 4D shuffle) that `generate_op_tests`
exports via `VulkanPartitioner` and compares to a torch golden on Dawn.
Also adds `test/ops/permute/test_permute.py` (`PermuteModule` +
`CONFIGS` + `_op_delegated` smoke test) and the
`aten.permute_copy.default` partitioner-allowlist entry in `tester.py`.
@exported-using-ghexport

Differential Revision:
[D108793156](https://our.internmc.facebook.com/intern/diff/D108793156/)

Differential Revision:
[D108793156](https://our.internmc.facebook.com/intern/diff/D108793156)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@psiddh psiddh psiddh approved these changes
@larryliu0820 larryliu0820 Awaiting requested review from larryliu0820 larryliu0820 is a code owner
@kirklandsign kirklandsign Awaiting requested review from kirklandsign kirklandsign is a code owner

Assignees

No one assigned

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

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@JulianCloudNTH @psiddh @nil-is-all