Skip to content

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

Open
SS-JIA wants to merge 4 commits 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 4 commits 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.

❌ 3 New Failures, 2 Unrelated Failures

As of commit c087102 with merge base 8a285b7 (image):

NEW FAILURES - The following jobs have failed:

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.

@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.

ssjia added 3 commits March 11, 2026 09:52
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]
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]
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]
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