Support texture2D in etvk.copy_offset (#5933)

Summary:
Pull Request resolved: #5933

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

Reviewed By: SS-JIA

Differential Revision: D63905519

fbshipit-source-id: 03b1cdf90bd9db299003a6152648f08705197f72

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;

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

@@ -26,6 +26,11 @@ void check_args(
   VK_CHECK_COND(check_packed_dim_is(in, WHCN::kChannelsDim));
   VK_CHECK_COND(check_packed_dim_is(out, WHCN::kChannelsDim));
 
+  VK_CHECK_COND(in.storage_type() == out.storage_type());
+  if (in.storage_type() == utils::kTexture2D) {
+    VK_CHECK_COND(in.dim() <= 2);
+  }
+
   int64_t in_dim = in.dim();
   VK_CHECK_COND(
       in_dim <= repeats.size(),

backends/vulkan/test/op_tests/cases.py

@@ -704,7 +704,21 @@ def get_clone_inputs():
 
 @register_test_suite("aten.repeat.default")
 def get_repeat_inputs():
-    test_suite = VkTestSuite(
+    test_suite_2d = VkTestSuite(
+        [
+            ((2, 3), [1, 4]),
+            ((2, 3), [4, 1]),
+            ((2, 3), [4, 4]),
+            ((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 +753,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_3d.test_name_suffix = "3d"
+
+    return [test_suite_2d, test_suite_3d]
 
 
 @register_test_suite("aten.repeat_interleave.self_int")