[ET-VK] Add ANY_STORAGE support to repeat#18056
[ET-VK] Add ANY_STORAGE support to repeat#18056SS-JIA wants to merge 1 commit intogh/SS-JIA/473/basefrom
Conversation
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]
🔗 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 FailureAs of commit c7bcdd3 with merge base f09bd55 ( 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. |
This PR needs a
|
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:
BufferMetadata UBOs and linear index arithmetic to map each output element
to its source element via modulo on each dimension.
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.
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.
repeat test suites.
Differential Revision: D95970170