[8/n][ET-VK] Support staging any 8-bit texture

Pull Request resolved: #5934

Changes following from #4485 to support `texture2d` and support `uint8`, respectively.
ghstack-source-id: 246658454
@exported-using-ghexport

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

Author: jorgep31415

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

@@ -10,14 +10,16 @@
 
 #define PRECISION ${PRECISION}
 
+${define_active_storage_type(STORAGE)}
+
 #include "indexing_utils.h"
 
 layout(std430) buffer;
 
 #extension GL_EXT_control_flow_attributes : require
 
 ${layout_declare_buffer(B, "w", "nchw_out", "int")}
-${layout_declare_tensor(B, "r", "t_in", "int8", "texture3d")}
+${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 ${layout_declare_ubo(B, "ivec4", "tensor_sizes")}
 ${layout_declare_ubo(B, "ivec4", "axis_map")}
 ${layout_declare_ubo(B, "int", "out_numel")}
@@ -44,7 +46,7 @@ void main() {
     const ivec4 tidx = nchwi_to_tidx(in_buf_idx, tensor_sizes);
     const ivec4 texture_pos = to_texture_elem_pos(
         tidx, tensor_sizes, packed_dim);
-    values[i] = load_texel(t_in, texture_pos.xyz)[texture_pos.w];
+    values[i] = ivec4(load_texel(t_in, texture_pos.xyz))[texture_pos.w];
     in_buf_idx++;
   }
 

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

@@ -0,0 +1,19 @@
+# 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.
+
+int8_image_to_nchw_noint8:
+  parameter_names_with_default_values:
+    STORAGE: texture3d
+    DTYPE: int8
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: int8
+      - VALUE: uint8
+    STORAGE:
+      - VALUE: texture2d
+      - VALUE: texture3d
+  shader_variants:
+    - NAME: int8_image_to_nchw_noint8

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

@@ -10,13 +10,17 @@
 
 #define PRECISION ${PRECISION}
 
+#define VEC4_T ${texel_load_type(DTYPE, STORAGE)}
+
+${define_active_storage_type(STORAGE)}
+
 #include "indexing_utils.h"
 
 layout(std430) buffer;
 
 #extension GL_EXT_control_flow_attributes : require
 
-${layout_declare_tensor(B, "w", "t_out", "int8", "texture3d")}
+${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
 ${layout_declare_buffer(B, "r", "nchw_in", "int")}
 ${layout_declare_ubo(B, "ivec4", "sizes")}
 ${layout_declare_ubo(B, "ivec4", "axis_map")}
@@ -71,5 +75,5 @@ void main() {
     return;
   }
 
-  write_texel(t_out, lpos_to_pos(lpos, axis_map), read_texel(tidx));
+  write_texel(t_out, lpos_to_pos(lpos, axis_map), VEC4_T(read_texel(tidx)));
 }

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

@@ -0,0 +1,19 @@
+# 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.
+
+nchw_to_int8_image_noint8:
+  parameter_names_with_default_values:
+    STORAGE: texture3d
+    DTYPE: int8
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: int8
+      - VALUE: uint8
+    STORAGE:
+      - VALUE: texture2d
+      - VALUE: texture3d
+  shader_variants:
+    - NAME: nchw_to_int8_image_noint8

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

@@ -80,7 +80,7 @@ void add_tensor_to_staging_node(
   // output buffer. Therefore, the global work group size for this shader will
   // be the number of elements in the output buffer divided by 4, as opposed to
   // the extents of the input texture.
-  if (shader.kernel_name == "int8_image_to_nchw_noint8") {
+  if (shader.kernel_name.starts_with("int8_image_to_nchw_noint8")) {
     uint32_t buffer_len = graph.get_staging(out_staging)->numel() / 4;
     global_wg_size = {buffer_len, 1, 1};
     ubos.append({graph.numel_ubo(in_tensor)});

backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp

@@ -15,15 +15,23 @@
 
 namespace vkcompute {
 
+bool is_8bit(vkapi::ScalarType dtype) {
+  return dtype == vkapi::kByte || dtype == vkapi::kChar ||
+      dtype == vkapi::kQInt8 || dtype == vkapi::kQUInt8;
+}
+
 vkapi::ShaderInfo get_nchw_to_tensor_shader(
     const api::vTensor& v_dst,
     const bool int8_buffer_enabled) {
   std::string kernel_name;
   kernel_name.reserve(kShaderNameReserve);
 
-  if (v_dst.dtype() == vkapi::kChar &&
-      v_dst.storage_type() == utils::kTexture3D && !int8_buffer_enabled) {
-    return VK_KERNEL(nchw_to_int8_image_noint8);
+  if (is_8bit(v_dst.dtype()) && v_dst.storage_type() != utils::kBuffer &&
+      !int8_buffer_enabled) {
+    kernel_name = "nchw_to_int8_image_noint8";
+    add_dtype_suffix(kernel_name, v_dst);
+    add_storage_type_suffix(kernel_name, v_dst);
+    return VK_KERNEL_FROM_STR(kernel_name);
   }
 
   if (v_dst.storage_type() == utils::kBuffer) {
@@ -45,9 +53,12 @@ vkapi::ShaderInfo get_tensor_to_nchw_shader(
   std::string kernel_name;
   kernel_name.reserve(kShaderNameReserve);
 
-  if (v_src.dtype() == vkapi::kChar &&
-      v_src.storage_type() == utils::kTexture3D && !int8_buffer_enabled) {
-    return VK_KERNEL(int8_image_to_nchw_noint8);
+  if (is_8bit(v_src.dtype()) && v_src.storage_type() != utils::kBuffer &&
+      !int8_buffer_enabled) {
+    kernel_name = "int8_image_to_nchw_noint8";
+    add_dtype_suffix(kernel_name, v_src);
+    add_storage_type_suffix(kernel_name, v_src);
+    return VK_KERNEL_FROM_STR(kernel_name);
   }
 
   if (v_src.storage_type() == utils::kBuffer) {