[SS-JIA]

Stack from ghstack (oldest at bottom):

• [ET-VK][ops] Consolidate binary scalar shaders behind op_type #24756
• -> [ET-VK][runtime] Preserve binary and permute zero padding #24755
• [ET-VK][runtime] Preserve zero padding in unary operators #24754
• [ET-VK][runtime] Align buffer packed dimensions to four #24753
• [ET-VK][ops] Pass op_type to unary and binary op helpers #24752

Ulterior Motive

Make buffer storage behave more like texture storage, so one shader can service
both storage types and buffer tensors need only a single vectorized shader.
This extends the unary zero-padding work to binary and permute outputs.

Rationale

What: Binary tensor-tensor and tensor-scalar ops, and texture permute, keep
the padding of their outputs zero. Each binary op family now has one vectorized
shader for buffers and textures.

Why: Consumers can load whole vec4s only if every producer leaves padding
zero. With buffers padded to four, the buffer binary shaders can also use vec4
loads, which makes them much faster.

Details

• Binary shaders
  • Before: binary_op_{buffer,texture} and
    binary_scalar_{buffer,texture}. Buffers ran one element per thread and
    computed in the storage type, so fp16 floor division rounded its quotient
    to fp16 before floor.
  • After: binary and binary_scalar, one vec4 per thread for both
    storage types, computing in texel_type (fp32 for half). Shader names are
    unchanged. alpha is cast to the component type, so int32 add and sub
    stay exact.
• Floor division
  • Before: int32 floor division truncated toward zero, and
    div.Tensor_mode ran floor division for every rounding mode.
  • After: int32 floor division rounds toward negative infinity. Only
    rounding_mode="floor" is supported: the partitioner keeps other modes on
    CPU, and dispatch checks the mode.
• Broadcast detection
  • Before: should_broadcast and the texture packed-dim broadcast flags
    were spec constants set at build time. Resizing an input to or from a
    broadcast shape gave wrong results.
  • After: both are computed at run time from tensor metadata. When both
    inputs match the output's sizes and strides (buffers) or sizes and layout
    (textures), texels load at the output's index. An input with
    size_at(in, packed_dim) < size_at(out, packed_dim) broadcasts along the
    packed dim; its lane 0 is copied to all lanes.
• Buffer broadcasting
  • The texel's first element is decoded to a tensor index and clamped per
    input. The packed dim has stride 1 and is padded to four, so the clamped
    index starts an input texel.
• Padding
  • binary_op_preserves_zero(op_type, alpha) in binary_op.glslh lists the
    tensor-tensor ops that map two zeros to zero. Like apply_binary_op, it
    takes the op type as an argument. add and sub qualify only for a
    finite alpha, since 0 * inf is NaN.
  • Those ops mask padding lanes only when an input broadcasts along the
    packed dim, since the copied lane fills padding lanes. Other tensor-tensor
    ops and all tensor-scalar ops always mask; mul by an infinite scalar
    maps zero to NaN. This replaces the texture shader's division-only
    padding fix. The mask returns after one check when the packed dim needs
    no padding.
• Permute
  • Textures that keep the packed dim still copy whole texels. Otherwise each
    output lane is gathered as before. Padding lanes re-read the last valid
    element, so every fetch stays in bounds, and are zeroed before the store.
• Dispatch
  • Binary ops use default_pick_vec4_gwg. Texture dispatch for binary,
    tensor-scalar and permute checks standard axis maps. Tensor-scalar
    dispatch checks matching packed dims and passes the output's hashed
    layout as a spec constant.

Smaller changes:

• size_at(TextureMetadata, dim) mirrors the buffer accessor.
• The unused is_packed_dim_broadcasted and create_broadcast_params are
  removed.
• Tensor views set legacy uniforms only for ranks up to four, and derive
  logical limits from the view's axis map.
• Generated tests:
  • Binary, comparison, bitwise, tensor-scalar and permute registrations check
    zero padding.
  • Comparison tests are re-enabled with int32 inputs.
  • Tensor pow and floor division use the binary case matrix. Floor
    division also runs int32, with operands in [-16, 17) from a new
    make_nonzero_randint_tensor generator, so floor and truncation differ.
  • New int32 add/sub suite with values in [2^24, 2^25), which fp32
    cannot all represent.
  • New add/sub suite with an infinite alpha, and mul.Scalar suite
    with an infinite scalar.
  • Tensor-scalar cases cover zero and negative scalars.
  • logical_and uses the bitwise suite.
  • Optional string arguments, such as floor division's rounding mode, are
    generated.
• test_vulkan_shader_names.py checks the merged binary template.
• New binary_packed_dim_broadcast_after_resize test.
• New test_vulkan_backend_trunc_div_is_not_delegated delegate test.

Review binary.glsl, binary_op.glslh and binary_scalar.glsl first, then
permute_texture.glsl, dispatch, Tensor.cpp and tests.

Authored with Codex and Claude Code.

Differential Revision: D130297306

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 4 Pending, 2 Unrelated Failures

As of commit 419b803 with merge base 9517b6b ():

NEW FAILURE - The following job has failed:

• pull / test-llama-runner-linux (fp32, xnnpack+custom+quantize_kv, mt-l-arm64g4-16-62, executorch-ubuntu-... / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Test Vulkan Backend (specialized runners) / test-model-vulkan-e2e / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Test Vulkan Backend (specialized runners) / test-vulkan-nvidia / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

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.