Skip to content

[ET-VK][ops] Add eq.Scalar operator#20383

Open
SS-JIA wants to merge 2 commits into
gh/SS-JIA/562/basefrom
gh/SS-JIA/562/head
Open

[ET-VK][ops] Add eq.Scalar operator#20383
SS-JIA wants to merge 2 commits into
gh/SS-JIA/562/basefrom
gh/SS-JIA/562/head

Conversation

@SS-JIA

@SS-JIA SS-JIA commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Adds Vulkan support for aten.eq.Scalar. This is the second of two ops needed to collapse the Llama4-mini TISO en_US backbone export to a single Vulkan partition (after bitwise_or): the discrete-speech mask compares the int token-id tensor against scalar constants via aten.eq.Scalar, which previously had no Vulkan implementation and forced a CPU fallback that split the delegated graph.

Implemented by extending the existing tensor-scalar binary-op path with a comparison-output variant: binary_scalar_buffer.glsl / binary_scalar_texture.glsl gain an IS_COMPARISON_OP code path that writes a uint8 (bool) output while leaving the existing arithmetic (e.g. pow) path unchanged; binary_scalar_buffer.yaml / binary_scalar_texture.yaml add an eq_scalar variant (half/float/int32 — the texture variant uses equal(X, Y) for per-lane bvec4, the buffer variant uses scalar X == Y); BinaryScalarOp.cpp adds an eq_tensor_scalar dispatch and VK_REGISTER_OP(aten.eq.Scalar, eq_tensor_scalar); op_registry.py registers aten.eq.Scalar OpFeatures (FP/INT tensor input, bool output). The int64 token tensor is serialized to int32 via the existing downcast_64_bit path, so the dispatch resolves to the int32 shader variant; no dtype-conversion pass is added.

This change was authored with Claude.

Differential Revision: D108457791

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jun 18, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 5 New Failures, 3 Unrelated Failures

As of commit 05e3da6 with merge base 1621fa2 (image):

NEW FAILURES - The following jobs have 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.

@linux-foundation-easycla

linux-foundation-easycla Bot commented Jun 18, 2026

Copy link
Copy Markdown

CLA Missing ID

@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
@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]
SS-JIA pushed a commit that referenced this pull request Jun 24, 2026
Pull Request resolved: #20383

Adds Vulkan support for `aten.eq.Scalar`. This is the second of two ops needed to collapse the Llama4-mini TISO en_US backbone export to a single Vulkan partition (after `bitwise_or`): the discrete-speech mask compares the int token-id tensor against scalar constants via `aten.eq.Scalar`, which previously had no Vulkan implementation and forced a CPU fallback that split the delegated graph.

Implemented by extending the existing tensor-scalar binary-op path with a comparison-output variant: `binary_scalar_buffer.glsl` / `binary_scalar_texture.glsl` gain an `IS_COMPARISON_OP` code path that writes a `uint8` (bool) output while leaving the existing arithmetic (e.g. `pow`) path unchanged; `binary_scalar_buffer.yaml` / `binary_scalar_texture.yaml` add an `eq_scalar` variant (half/float/int32 — the texture variant uses `equal(X, Y)` for per-lane `bvec4`, the buffer variant uses scalar `X == Y`); `BinaryScalarOp.cpp` adds an `eq_tensor_scalar` dispatch and `VK_REGISTER_OP(aten.eq.Scalar, eq_tensor_scalar)`; `op_registry.py` registers `aten.eq.Scalar` `OpFeatures` (FP/INT tensor input, bool output). The int64 token tensor is serialized to int32 via the existing `downcast_64_bit` path, so the dispatch resolves to the int32 shader variant; no dtype-conversion pass is added.

This change was authored with Claude.
ghstack-source-id: 396618180
@exported-using-ghexport

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

1 participant