[ET-VK] Add ANY_STORAGE support to repeat

backends/vulkan/op_registry.py

@@ -1378,7 +1378,7 @@ def register_grid_priors():
 @update_features(exir_ops.edge.aten.repeat.default)
 def register_repeat():
     return OpFeatures(
-        inputs_storage=utils.ANY_TEXTURE,
+        inputs_storage=utils.ANY_STORAGE,
         inputs_dtypes=utils.FP_INT_BOOL_T,
     )
 

backends/vulkan/runtime/graph/ops/glsl/repeat_buffer.glsl

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+${define_required_extensions("buffer", DTYPE)}
+
+#define PRECISION ${PRECISION}
+
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type("buffer")}
+
+layout(std430) buffer;
+
+#include "indexing.glslh"
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer")}
+
+${layout_declare_ubo(B, "BufferMetadata", "out_meta")}
+${layout_declare_ubo(B, "BufferMetadata", "in_meta")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const uint out_bufi = gl_GlobalInvocationID.x;
+  if (out_of_bounds(out_bufi, out_meta)) {
+    return;
+  }
+
+  TensorIndex out_tidx = linear_idx_to_tensor_idx(out_meta, out_bufi);
+
+  TensorIndex in_tidx;
+  initialize(in_tidx);
+
+  const int n = int_ndim(out_meta);
+  for (int d = 0; d < n; d++) {
+    in_tidx.data[div_4(d)][mod_4(d)] =
+        idx_at(out_tidx, d) % size_at(in_meta, d);
+  }
+
+  const uint in_bufi = tensor_idx_to_linear_idx(in_meta, in_tidx);
+
+  t_out[out_bufi] = t_in[in_bufi];
+}

backends/vulkan/runtime/graph/ops/glsl/repeat.yaml → backends/vulkan/runtime/graph/ops/glsl/repeat_buffer.yaml

@@ -1,8 +1,7 @@
-repeat:
+repeat_buffer:
   parameter_names_with_default_values:
     DTYPE: float
-    NDIM: 3
-    STORAGE: texture3d
+    STORAGE: buffer
   generate_variant_forall:
     DTYPE:
       - VALUE: half
@@ -11,4 +10,4 @@ repeat:
       - VALUE: int8
       - VALUE: uint8
   shader_variants:
-    - NAME: repeat
+    - NAME: repeat_buffer

backends/vulkan/runtime/graph/ops/glsl/repeat_texture.glsl

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+${define_required_extensions("texture3d", DTYPE)}
+
+#define PRECISION ${PRECISION}
+
+#define VEC4_T ${texel_load_type(DTYPE, "texture3d")}
+
+${define_active_storage_type("texture3d")}
+
+#extension GL_EXT_control_flow_attributes : require
+
+layout(std430) buffer;
+
+#include "common.glslh"
+#include "indexing.glslh"
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
+${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")}
+
+${layout_declare_ubo(B, "TextureMetadata", "out_meta")}
+${layout_declare_ubo(B, "TextureMetadata", "in_meta")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
+const int packed_dim = get_packed_dim(out_layout);
+
+void main() {
+  const ivec3 out_pos = ivec3(gl_GlobalInvocationID);
+
+  if (out_of_bounds(out_pos, out_meta)) {
+    return;
+  }
+
+  TensorIndex4D out_tidx = texture_pos_to_tensor4d_idx_simple(out_meta, out_pos);
+
+  VEC4_T out_texel = VEC4_T(0);
+
+  int limit = min(
+      4, out_meta.sizes[packed_dim] - out_tidx.data[packed_dim]);
+  for (int comp = 0; comp < 4; comp++) {
+    TensorIndex4D in_tidx = out_tidx;
+    in_tidx.data = ivec4(
+        out_tidx.data.x % in_meta.sizes.x,
+        out_tidx.data.y % in_meta.sizes.y,
+        out_tidx.data.z % in_meta.sizes.z,
+        out_tidx.data.w % in_meta.sizes.w);
+
+    TextureElementIndex in_elem =
+        tensor4d_idx_to_texture_element_idx_simple(in_meta, in_tidx);
+
+    VEC4_T in_texel = texelFetch(t_in, in_elem.pos, 0);
+    out_texel[comp] = in_texel[in_elem.comp];
+
+    out_tidx.data[packed_dim]++;
+  }
+
+  imageStore(t_out, out_pos, out_texel);
+}

backends/vulkan/runtime/graph/ops/glsl/repeat_texture.yaml

@@ -0,0 +1,12 @@
+repeat_texture:
+  parameter_names_with_default_values:
+    DTYPE: float
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+      - VALUE: int32
+      - VALUE: int8
+      - VALUE: uint8
+  shader_variants:
+    - NAME: repeat_texture3d