Skip to content

[ExecuTorch][WebGPU] Add select_copy op (aten.select_copy.int)#20362

Merged
JulianCloudNTH merged 7 commits into
gh/JulianCloudNTH/37/basefrom
gh/JulianCloudNTH/37/head
Jun 26, 2026
Merged

[ExecuTorch][WebGPU] Add select_copy op (aten.select_copy.int)#20362
JulianCloudNTH merged 7 commits into
gh/JulianCloudNTH/37/basefrom
gh/JulianCloudNTH/37/head

Conversation

@JulianCloudNTH

@JulianCloudNTH JulianCloudNTH commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

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

Differential Revision: D108793166

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jun 17, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure, 1 Pending, 2 Unrelated Failures

As of commit 94dd90d with merge base e03f777 (image):

NEW FAILURE - The following job has failed:

BROKEN TRUNK - The following jobs 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.

@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 18, 2026
@JulianCloudNTH JulianCloudNTH requested review from SS-JIA and psiddh June 18, 2026 17:41
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
@JulianCloudNTH JulianCloudNTH merged commit 01ffde5 into gh/JulianCloudNTH/37/base Jun 26, 2026
179 of 183 checks passed
@JulianCloudNTH JulianCloudNTH deleted the gh/JulianCloudNTH/37/head branch June 26, 2026 15:14
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
…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
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
…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
…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
….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
Pull Request resolved: #20362

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.
ghstack-source-id: 397026510
@exported-using-ghexport

Differential Revision: [D108793166](https://our.internmc.facebook.com/intern/diff/D108793166/)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
Pull Request resolved: #20362

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.
ghstack-source-id: 397026510
@exported-using-ghexport

Differential Revision: [D108793166](https://our.internmc.facebook.com/intern/diff/D108793166/)
JulianCloudNTH added a commit that referenced this pull request Jun 26, 2026
Pull Request resolved: #20362

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.
ghstack-source-id: 397026510
@exported-using-ghexport

Differential Revision: [D108793166](https://our.internmc.facebook.com/intern/diff/D108793166/)
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.

3 participants