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

"bitw8" = bit width 8, which is equivalent to 8-bit. We use "bitw8" as the name since shader compilation disallows names starting with a digit.

Changes follow from #4485 to support `texture2d` and support `uint8`, respectively.

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

[ghstack-poisoned]

Author: jorgep31415

backends/vulkan/runtime/graph/ops/glsl/int8_image_to_nchw_noint8.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/bitw8_image_to_nchw_nobitw8buffer.glsl

@@ -4,7 +4,7 @@
 # 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:
+bitw8_image_to_nchw_nobitw8buffer:
   parameter_names_with_default_values:
     STORAGE: texture3d
     DTYPE: int8
@@ -16,4 +16,4 @@ int8_image_to_nchw_noint8:
       - VALUE: texture2d
       - VALUE: texture3d
   shader_variants:
-    - NAME: int8_image_to_nchw_noint8
+    - NAME: bitw8_image_to_nchw_nobitw8buffer

backends/vulkan/runtime/graph/ops/glsl/int8_image_to_nchw_noint8.yaml renamed to backends/vulkan/runtime/graph/ops/glsl/bitw8_image_to_nchw_nobitw8buffer.yaml

@@ -4,7 +4,7 @@
 # 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:
+nchw_to_bitw8_image_nobitw8buffer:
   parameter_names_with_default_values:
     STORAGE: texture3d
     DTYPE: int8
@@ -16,4 +16,4 @@ nchw_to_int8_image_noint8:
       - VALUE: texture2d
       - VALUE: texture3d
   shader_variants:
-    - NAME: nchw_to_int8_image_noint8
+    - NAME: nchw_to_bitw8_image_nobitw8buffer

backends/vulkan/runtime/graph/ops/glsl/nchw_to_int8_image_noint8.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.glsl

@@ -116,7 +116,6 @@ 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/runtime/graph/ops/glsl/nchw_to_int8_image_noint8.yaml renamed to backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.yaml

@@ -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.starts_with("int8_image_to_nchw_noint8")) {
+  if (shader.kernel_name.starts_with("bitw8_image_to_nchw_nobitw8buffer")) {
     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/impl/Copy.cpp

@@ -15,7 +15,7 @@
 
 namespace vkcompute {
 
-bool is_8bit(vkapi::ScalarType dtype) {
+bool is_bitw8(vkapi::ScalarType dtype) {
   return dtype == vkapi::kByte || dtype == vkapi::kChar ||
       dtype == vkapi::kQInt8 || dtype == vkapi::kQUInt8;
 }
@@ -26,9 +26,9 @@ vkapi::ShaderInfo get_nchw_to_tensor_shader(
   std::string kernel_name;
   kernel_name.reserve(kShaderNameReserve);
 
-  if (is_8bit(v_dst.dtype()) && v_dst.storage_type() != utils::kBuffer &&
+  if (is_bitw8(v_dst.dtype()) && v_dst.storage_type() != utils::kBuffer &&
       !int8_buffer_enabled) {
-    kernel_name = "nchw_to_int8_image_noint8";
+    kernel_name = "nchw_to_bitw8_image_nobitw8buffer";
     add_dtype_suffix(kernel_name, v_dst);
     add_storage_type_suffix(kernel_name, v_dst);
     return VK_KERNEL_FROM_STR(kernel_name);
@@ -53,9 +53,9 @@ vkapi::ShaderInfo get_tensor_to_nchw_shader(
   std::string kernel_name;
   kernel_name.reserve(kShaderNameReserve);
 
-  if (is_8bit(v_src.dtype()) && v_src.storage_type() != utils::kBuffer &&
+  if (is_bitw8(v_src.dtype()) && v_src.storage_type() != utils::kBuffer &&
       !int8_buffer_enabled) {
-    kernel_name = "int8_image_to_nchw_noint8";
+    kernel_name = "bitw8_image_to_nchw_nobitw8buffer";
     add_dtype_suffix(kernel_name, v_src);
     add_storage_type_suffix(kernel_name, v_src);
     return VK_KERNEL_FROM_STR(kernel_name);