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

To use `copy_offset` with 2D textures in AHB experimentation.

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

ghstack-source-id: 246607067
Pull Request resolved: #5933

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

@@ -742,6 +742,7 @@ def get_repeat_inputs():
     test_suite.layouts = [
         "utils::kChannelsPacked",
     ]
+    test_suite.storage_types = ["utils::kTexture2D", "utils::kTexture3D"]
     test_suite.data_gen = "make_seq_tensor"
     test_suite.dtypes = ["at::kFloat"]
     return test_suite