Repository navigation
[ET-VK][unary] Select unary ops with a specialization constant - #24718
Merged
Merged
Conversation
## Ulterior Motive
Preparing ET-VK unary ops for operator fusion, such as applying an activation
inside the shader that produces its input. A fused shader needs a shared
implementation of every unary op it can apply, instead of a separate shader per
op.
## Rationale
**What**: Every unary op of a storage type and dtype now dispatches one shader,
and the `op_type` specialization constant selects the op. The op
implementations move to `unary_op.glslh` as `apply_unary_op()`. `tan` joins the
shared shader. relu gets its own op, `max(X, 0)`, instead of clamp with bounds
0 and +inf. `q8ta_relu` uses the same header and constant. Unary sources move
into `unary/` family directories, with the q8ta unary op under `unary/q8ta/`.
**Why**: Fusion shaders can include `unary_op.glslh` to apply any unary op.
Generated unary shaders drop from 86 to 14.
## Details
- Shader variants
- **Before**: one per op, dtype and storage (`exp_half_texture3d`,
`clamp_int32_buffer`, `clamp_dynamic_float_texture3d`, ...), each with the
op expression baked in by codegen; `tan` had its own template.
- **After**: `unary_op_{half,float,int32}_{texture3d,buffer}` (int32
implements clamp only), `unary_op_uint8_{texture3d,buffer}` (bitwise_not
only), and `unary_op_dynamic_{half,float,int32}_{texture3d,buffer}` for
clamp bounds that are only known at runtime.
- Op selection
- **Before**: the kernel name (op + dtype + storage).
- **After**: `op_type` (specialization constant 3), from `UnaryOpType` in
`UnaryOp.h`; values must match `unary_op.glslh`. relu6 and hardtanh still
use clamp. The driver folds the switch at pipeline creation; every
op/variant pair checked compiles to the same Adreno ISA as the old per-op
shader, except the cases below.
- relu
- **Before**: the clamp shader with bounds 0 and +inf from push constants.
The compiler cannot see that the upper bound is +inf, so every element
also pays a `min`.
- **After**: `UNARY_OP_RELU` computes `max(X, 0)`, so fp relu drops the
`min`. Without it, q8ta relu would have gained the `min`: its old shader
had `max(X, 0)` baked in.
- `tan`
- **Before**: its own shader, reading the output limits from a UBO.
- **After**: the shared shader, reading them from push constants like every
other unary op.
- Runtime clamp bounds on fp16 buffers
- **Before**: clamped in fp16.
- **After**: clamped in fp32, then stored as fp16. The result is unchanged:
rounding to fp16 is monotonic and the input is already fp16. This lets the
float-only ops compile into the runtime-bounds shader.
- q8ta relu
- **Before**: `q8ta_relu_buffer`, with `max(X, vec4(0.0))` baked in by
codegen.
- **After**: `q8ta_unary_buffer`, which calls `apply_unary_op()` with
`op_type` (specialization constant 6). `add_q8ta_unary_node()` takes a
`UnaryOpType` instead of an op name; the enum moves to `UnaryOp.h` so both
files share it.
- Files
- `unary_op.glsl` / `.yaml` become `unary/unary.glsl` / `.yaml`: codegen keys
sources by name without extension, so the template cannot share the
header's name.
- The `activations.h` helpers move into `unary_op.glslh`; `Tan.cpp`,
`tan.glsl` and `tan.yaml` are removed.
- `UnaryOp.cpp` moves to `unary/`; `q8ta_unary.*` and `Q8taUnary.*` move to
`unary/q8ta/`.
- Kernel names in profiles: `unary_op_*` for every op, and `q8ta_unary_buffer`
for q8ta relu; the op is no longer part of the name.
A fusion shader declares `op_type`, defines the operand types, and includes the
header (`unary/unary_op.glslh` from outside the unary family):
```glsl
#define IN_T VEC4_T
#define OUT_T VEC4_T
#define PARAM_T float
#define UNARY_OP_IS_FLOAT
${layout_declare_spec_const(C, "int", "op_type", "0")}
#include "unary/unary_op.glslh"
...
out_texel = apply_unary_op(conv_texel, minimum, maximum);
```
Differential Revision: [D130014629](https://our.internmc.facebook.com/intern/diff/D130014629/)
ghstack-source-id: 466197031
Pull-Request: #24696
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24718
Note: Links to docs will display an error until the docs builds have been completed. ❌ 135 New Failures, 7 Unrelated Failures, 10 Unclassified FailuresAs of commit c242387 with merge base 728e13f ( 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:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
SS-JIA
approved these changes
Nov 6, 2026
This branch had an error being deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #24696 by @SS-JIA
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/SS-JIA/691/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/SS-JIA/691/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/SS-JIA/690/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/SS-JIA/691/orig
Differential Revision: D130014629
@diff-train-skip-merge