[7/n][ET-VK] Support texture2D in etvk.copy_offset

Pull Request resolved: #5933

To use `copy_offset` with 2D textures in AHB experimentation.
ghstack-source-id: 246657862
@exported-using-ghexport

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

Author: jorgep31415

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

@@ -10,6 +10,8 @@
 
 #define PRECISION ${PRECISION}
 
+${define_active_storage_type(STORAGE)}
+
 #include "indexing_utils.h"
 
 layout(std430) buffer;

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

@@ -8,5 +8,10 @@ copy_offset:
       - VALUE: half
       - VALUE: float
       - VALUE: int
+      - VALUE: int8
+      - VALUE: uint8
+    STORAGE:
+      - VALUE: texture3d
+      - VALUE: texture2d
   shader_variants:
     - NAME: copy_offset

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -212,6 +212,8 @@ ivec3 lpos_to_pos(const ivec3 lpos, const ivec4 axis_map) {
 #define load_texel(buf, idx) buf[idx]
 #elif defined(USING_TEXTURE2D)
 #define load_texel(im, pos) texelFetch(im, pos.xy, 0)
+#define load_texel_lpos(im, lpos, axis_map) \
+  texelFetch(im, lpos_to_pos(lpos, axis_map).xy, 0)
 #else // defined(USING_TEXTURE3D)
 #define load_texel(im, pos) texelFetch(im, pos, 0)
 #define load_texel_lpos(im, lpos, axis_map) \
@@ -222,6 +224,8 @@ ivec3 lpos_to_pos(const ivec3 lpos, const ivec4 axis_map) {
 #define write_texel(buf, idx, texel) buf[idx] = texel
 #elif defined(USING_TEXTURE2D)
 #define write_texel(im, pos, texel) imageStore(im, pos.xy, texel)
+#define write_texel_lpos(im, lpos, texel, axis_map) \
+  imageStore(im, lpos_to_pos(lpos, axis_map).xy, texel)
 #else // defined(USING_TEXTURE3D)
 #define write_texel(im, pos, texel) imageStore(im, pos, texel)
 #define write_texel_lpos(im, lpos, texel, axis_map) \

backends/vulkan/runtime/graph/ops/impl/Copy.cpp

@@ -31,6 +31,7 @@ void add_copy_offset_node(
   std::string kernel_name = "copy_offset";
   kernel_name.reserve(kShaderNameReserve);
   add_dtype_suffix(kernel_name, *t_out);
+  add_storage_type_suffix(kernel_name, *t_out);
 
   const struct Block final {
     alignas(16) ivec3 range;
@@ -115,6 +116,7 @@ void add_copy_channel_offset_node(
   std::string kernel_name = "copy_channel_offset";
   kernel_name.reserve(kShaderNameReserve);
   add_dtype_suffix(kernel_name, *t_out);
+  add_storage_type_suffix(kernel_name, *t_out);
 
   int32_t out_channels = dim_at<kChannel4D>(out_sizes);
 

backends/vulkan/test/op_tests/cases.py

@@ -704,7 +704,30 @@ def get_clone_inputs():
 
 @register_test_suite("aten.repeat.default")
 def get_repeat_inputs():
-    test_suite = VkTestSuite(
+    test_suite_2d = VkTestSuite(
+        [
+            # Repeat channels only (most challenging case)
+            ((3, XS, S), [2, 1, 1]),
+            ((7, XS, S), [4, 1, 1]),
+            # More other cases
+            ((2, 3), [1, 4]),
+            ((2, 3), [4, 1]),
+            ((2, 3), [4, 4]),
+            ((S1, S2, S2), [1, 3, 1]),
+            ((S1, S2, S2), [1, 3, 3]),
+            ((S1, S2, S2), [3, 3, 1]),
+            ((S1, S2, S2), [3, 3, 3]),
+            # Expanding cases
+            ((2, 3), [3, 1, 4]),
+        ]
+    )
+    test_suite_2d.layouts = ["utils::kChannelsPacked"]
+    test_suite_2d.storage_types = ["utils::kTexture2D"]
+    test_suite_2d.data_gen = "make_seq_tensor"
+    test_suite_2d.dtypes = ["at::kFloat"]
+    test_suite_2d.test_name_suffix = "2d"
+
+    test_suite_3d = VkTestSuite(
         [
             # Repeat channels only (most challenging case)
             ((3, XS, S), [2, 1, 1]),
@@ -739,12 +762,13 @@ def get_repeat_inputs():
             ((2, 3), [3, 3, 2, 4]),
         ]
     )
-    test_suite.layouts = [
-        "utils::kChannelsPacked",
-    ]
-    test_suite.data_gen = "make_seq_tensor"
-    test_suite.dtypes = ["at::kFloat"]
-    return test_suite
+    test_suite_3d.layouts = ["utils::kChannelsPacked"]
+    test_suite_3d.storage_types = ["utils::kTexture3D"]
+    test_suite_3d.data_gen = "make_seq_tensor"
+    test_suite_3d.dtypes = ["at::kFloat"]
+    test_suite_2d.test_name_suffix = "3d"
+
+    return [test_suite_2d, test_suite_3d]
 
 
 @register_test_suite("aten.repeat_interleave.self_int")