[ET-VK] Use push constants for image and buffer to nchw prepack nodes.

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

[ghstack-poisoned]

Author: trivedivivek

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

@@ -20,8 +20,15 @@ layout(std430) buffer;
 
 ${layout_declare_buffer(B, "w", "nchw_out", "int")}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
-${layout_declare_ubo(B, "ivec4", "tensor_sizes")}
-${layout_declare_ubo(B, "int", "out_numel")}
+
+$if USE_PUSH_CONST:
+  layout(push_constant) uniform restrict Block {
+    ivec4 tensor_sizes;
+    int out_numel;
+  };
+$else:
+  ${layout_declare_ubo(B, "ivec4", "tensor_sizes")}
+  ${layout_declare_ubo(B, "int", "out_numel")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -8,6 +8,7 @@ bitw8_image_to_nchw_nobitw8buffer:
   parameter_names_with_default_values:
     STORAGE: texture3d
     DTYPE: int8
+    USE_PUSH_CONST: True
   generate_variant_forall:
     STORAGE:
       - VALUE: texture2d
@@ -17,3 +18,5 @@ bitw8_image_to_nchw_nobitw8buffer:
       - VALUE: uint8
   shader_variants:
     - NAME: bitw8_image_to_nchw_nobitw8buffer
+    - NAME: bitw8_image_to_nchw_nobitw8buffer_no_pc
+      USE_PUSH_CONST: False

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

@@ -12,9 +12,17 @@ layout(std430) buffer;
 
 ${layout_declare_tensor(0, "w", "nchw_buf", DTYPE, STORAGE)}
 ${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)}
-${layout_declare_ubo(2, "ivec4", "in_sizes")}
-${layout_declare_ubo(3, "ivec4", "in_strides")}
-${layout_declare_ubo(4, "int", "numel")}
+
+$if USE_PUSH_CONST:
+  layout(push_constant) uniform restrict Block {
+    ivec4 in_sizes;
+    ivec4 in_strides;
+    int numel;
+  };
+$else:
+  ${layout_declare_ubo(2, "ivec4", "in_sizes")}
+  ${layout_declare_ubo(3, "ivec4", "in_strides")}
+  ${layout_declare_ubo(4, "int", "numel")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -8,6 +8,7 @@ buffer_to_nchw:
   parameter_names_with_default_values:
     DTYPE: float
     STORAGE: buffer
+    USE_PUSH_CONST: True
   generate_variant_forall:
     DTYPE:
       - VALUE: half
@@ -17,3 +18,5 @@ buffer_to_nchw:
       - VALUE: uint8
   shader_variants:
     - NAME: buffer_to_nchw
+    - NAME: buffer_to_nchw_no_pc
+      USE_PUSH_CONST: False

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

@@ -21,9 +21,17 @@ layout(std430) buffer;
 
 ${layout_declare_buffer(B, "w", "buf_out", DTYPE)}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
-${layout_declare_ubo(B, "ivec4", "sizes")}
-$if not TO_STAGING:
-  ${layout_declare_ubo(B, "ivec4", "buf_strides")}
+
+$if USE_PUSH_CONST:
+  layout(push_constant) uniform restrict Block {
+    ivec4 sizes;
+  $if not TO_STAGING:
+    ivec4 buf_strides;
+  };
+$else:
+  ${layout_declare_ubo(B, "ivec4", "sizes")}
+  $if not TO_STAGING:
+    ${layout_declare_ubo(B, "ivec4", "buf_strides")}
 
 #include "indexing_utils.h"
 

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

@@ -9,6 +9,7 @@ image_to_nchw:
     DTYPE: float
     STORAGE: texture3d
     TO_STAGING: True
+    USE_PUSH_CONST: True
   generate_variant_forall:
     DTYPE:
       - VALUE: half
@@ -22,3 +23,11 @@ image_to_nchw:
       STORAGE: texture2d
     - NAME: clone_image_to_buffer
       TO_STAGING: False
+    - NAME: image_to_nchw_no_pc_texture3d
+      USE_PUSH_CONST: False
+    - NAME: image_to_nchw_no_pc_texture2d
+      STORAGE: texture2d
+      USE_PUSH_CONST: False
+    - NAME: clone_image_to_buffer_no_pc
+      TO_STAGING: False
+      USE_PUSH_CONST: False

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

@@ -77,9 +77,9 @@ void add_image_to_buffer_node(
       // Input and Outputs
       {{buffer, vkapi::kWrite}, {image, vkapi::kRead}},
       // Parameter Buffers
-      {graph.sizes_ubo(image), graph.strides_ubo(buffer)},
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(image), graph.strides_pc_of(buffer)},
       // Specialization Constants
       {graph.hashed_layout_of(image)},
       // Resize Args

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

@@ -76,14 +76,14 @@ void add_tensor_to_staging_node(
 
   utils::uvec3 global_wg_size = graph.create_global_wg_size(in_tensor);
 
-  vkapi::ParamsBindList ubos;
+  std::vector<PushConstantDataInfo> pcs;
   if (graph.is_buffer_storage(in_tensor)) {
-    ubos.append(
-        {graph.sizes_ubo(in_tensor),
-         graph.strides_ubo(in_tensor),
-         graph.numel_ubo(in_tensor)});
+    pcs = {
+        graph.sizes_pc_of(in_tensor),
+        graph.strides_pc_of(in_tensor),
+        graph.numel_pc_of(in_tensor)};
   } else {
-    ubos.append({graph.sizes_ubo(in_tensor)});
+    pcs = {graph.sizes_pc_of(in_tensor)};
   }
 
   // Normally, the image_to_nchw shader is structured so that each thread reads
@@ -97,7 +97,7 @@ void add_tensor_to_staging_node(
   if (is_bitw8_shader(shader)) {
     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)});
+    pcs.push_back(graph.numel_pc_of(in_tensor));
   }
 
   graph.execute_nodes().emplace_back(new DispatchNode(
@@ -108,9 +108,9 @@ void add_tensor_to_staging_node(
       // Input and Outputs
       {{out_staging, vkapi::kWrite}, {in_tensor, vkapi::kRead}},
       // Parameter Buffers
-      ubos,
-      // Push Constants
       {},
+      // Push Constants
+      pcs,
       // Specialization Constants
       {graph.hashed_layout_of(in_tensor)},
       // Resize Args

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

@@ -59,25 +59,35 @@ vkapi::ShaderInfo get_nchw_to_tensor_shader(
 
 vkapi::ShaderInfo get_tensor_to_nchw_shader(
     const api::vTensor& v_src,
-    bool int8_buffer_enabled) {
+    bool int8_buffer_enabled,
+    bool push_constant_variant) {
   std::string kernel_name;
   kernel_name.reserve(kShaderNameReserve);
 
   if (is_bitw8(v_src.dtype()) && v_src.storage_type() != utils::kBuffer &&
       !int8_buffer_enabled) {
     kernel_name = "bitw8_image_to_nchw_nobitw8buffer";
+    if (!push_constant_variant) {
+      kernel_name += "_no_pc";
+    }
     add_storage_type_suffix(kernel_name, v_src);
     add_dtype_suffix(kernel_name, v_src);
     return VK_KERNEL_FROM_STR(kernel_name);
   }
 
   if (v_src.storage_type() == utils::kBuffer) {
     kernel_name = "buffer_to_nchw";
+    if (!push_constant_variant) {
+      kernel_name += "_no_pc";
+    }
     add_dtype_suffix(kernel_name, v_src);
     return VK_KERNEL_FROM_STR(kernel_name);
   }
 
   kernel_name = "image_to_nchw";
+  if (!push_constant_variant) {
+    kernel_name += "_no_pc";
+  }
   add_storage_type_suffix(kernel_name, v_src);
   add_dtype_suffix(kernel_name, v_src);
 

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

@@ -18,6 +18,7 @@ vkapi::ShaderInfo get_nchw_to_tensor_shader(
     bool push_constant_variant = true);
 vkapi::ShaderInfo get_tensor_to_nchw_shader(
     const api::vTensor& v_src,
-    bool int8_buffer_enabled = true);
+    bool int8_buffer_enabled = true,
+    bool push_constant_variant = true);
 
 } // namespace vkcompute