Skip to content

[ET-VK] Add ANY_STORAGE support to repeat#18056

Open
SS-JIA wants to merge 1 commit intogh/SS-JIA/473/basefrom
gh/SS-JIA/473/head
Open

[ET-VK] Add ANY_STORAGE support to repeat#18056
SS-JIA wants to merge 1 commit intogh/SS-JIA/473/basefrom
gh/SS-JIA/473/head

Conversation

@SS-JIA
Copy link
Contributor

@SS-JIA SS-JIA commented Mar 10, 2026

Stack from ghstack (oldest at bottom):

Add a buffer shader variant for the repeat operator and rewrite the C++
dispatch code to support both buffer and texture storage types. Key changes:

  • Add repeat_buffer.glsl/.yaml: a new buffer-path compute shader that uses
    BufferMetadata UBOs and linear index arithmetic to map each output element
    to its source element via modulo on each dimension.
  • Rename repeat.glsl/.yaml to repeat_texture.glsl/.yaml (shader variant name
    repeat_texture3d) and update the push constant from dst_repeats to out_dims,
    computing out_channel_size directly from output dimensions rather than from
    input size * repeat count.
  • Rewrite Repeat.cpp: remove the old check_args() function and pre-computed
    ivec4 push constants. Add resize_repeat_node() that uses
    extract_int_or_symint_list() to dynamically compute output sizes from input
    sizes and repeats, enabling dynamic shape support. Dispatch uses
    add_storage_type_suffix for shader selection, with the buffer path using
    meta_ubo UBOs and the texture path using push constants from
    logical_limits_pc_of/sizes_pc_of. Wire resize_repeat_node into
    DynamicDispatchNode.
  • Update op_registry.py: change inputs_storage from ANY_TEXTURE to ANY_STORAGE.
  • Update test cases: add utils::kBuffer to storage_types for both 2d and 3d
    repeat test suites.

Differential Revision: D95970170

Add a buffer shader variant for the repeat operator and rewrite the C++
dispatch code to support both buffer and texture storage types. Key changes:

- Add repeat_buffer.glsl/.yaml: a new buffer-path compute shader that uses
  BufferMetadata UBOs and linear index arithmetic to map each output element
  to its source element via modulo on each dimension.
- Rename repeat.glsl/.yaml to repeat_texture.glsl/.yaml (shader variant name
  repeat_texture3d) and update the push constant from dst_repeats to out_dims,
  computing out_channel_size directly from output dimensions rather than from
  input size * repeat count.
- Rewrite Repeat.cpp: remove the old check_args() function and pre-computed
  ivec4 push constants. Add resize_repeat_node() that uses
  extract_int_or_symint_list() to dynamically compute output sizes from input
  sizes and repeats, enabling dynamic shape support. Dispatch uses
  add_storage_type_suffix for shader selection, with the buffer path using
  meta_ubo UBOs and the texture path using push constants from
  logical_limits_pc_of/sizes_pc_of. Wire resize_repeat_node into
  DynamicDispatchNode.
- Update op_registry.py: change inputs_storage from ANY_TEXTURE to ANY_STORAGE.
- Update test cases: add utils::kBuffer to storage_types for both 2d and 3d
  repeat test suites.

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

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Mar 10, 2026

🔗 Helpful Links

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

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

❌ 4 New Failures, 2 Cancelled Jobs, 1 Unrelated Failure

As of commit c7bcdd3 with merge base f09bd55 (image):

NEW FAILURES - The following jobs have failed:

CANCELLED JOBS - The following jobs were cancelled. Please retry:

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.

@github-actions
Copy link

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.

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. fb-exported meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant