[ET-VK] Using push constants for convolution prepack dispatches.

Pull Request resolved: #11306

This diff switches the `conv2d_prepack_weights`, `conv2d_dw_prepack_weights`, and `conv_transpose2d_prepack_weights` shaders to use push constants instead of uniform buffers for the `sizes` and `original_sizes` variables.

ghstack-source-id: 288178302
@exported-using-ghexport

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

Author: trivedivivek

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

@@ -26,11 +26,8 @@ layout(set = 0, binding = 1) buffer  PRECISION restrict readonly Buffer {
   BUF_T buffer_in[];
 };
 
-layout(set = 0, binding = 2) uniform PRECISION restrict Sizes {
+layout(push_constant) uniform PRECISION restrict Block {
   ivec4 sizes;
-};
-
-layout(set = 0, binding = 3) uniform PRECISION restrict OriginalSizes {
   ivec4 original_sizes;
 };
 

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

@@ -26,11 +26,8 @@ layout(set = 0, binding = 1) buffer  PRECISION restrict readonly Buffer {
   BUF_T buffer_in[];
 };
 
-layout(set = 0, binding = 2) uniform PRECISION restrict Sizes {
+layout(push_constant) uniform PRECISION restrict Block {
   ivec4 sizes;
-};
-
-layout(set = 0, binding = 3) uniform PRECISION restrict OriginalSizes {
   ivec4 original_sizes;
 };
 

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

@@ -26,11 +26,8 @@ layout(set = 0, binding = 1) buffer  PRECISION restrict readonly Buffer {
   BUF_T buffer_in[];
 };
 
-layout(set = 0, binding = 2) uniform PRECISION restrict Sizes {
+layout(push_constant) uniform PRECISION restrict Block {
   ivec4 sizes;
-};
-
-layout(set = 0, binding = 3) uniform PRECISION restrict OriginalSizes {
   ivec4 original_sizes;
 };
 

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

@@ -211,18 +211,20 @@ ValueRef prepack_weights(
   vkapi::ShaderInfo shader =
       get_conv2d_shader(graph, *t, /*prepack_weights = */ true, method, vref);
 
+  const auto original_sizes_pc =
+      utils::make_ivec4(original_sizes, /*reverse = */ true);
   graph.prepack_nodes().emplace_back(new PrepackNode(
       graph,
       shader,
       graph.create_global_wg_size(v),
       graph.create_local_wg_size(v),
       vref,
       v,
-      {t->sizes_ubo(),
-       graph.create_params_buffer(
-           utils::make_ivec4(original_sizes, /*reverse = */ true))},
+      {},
       // Specialization constants
-      {SV(t->packed_dim())}));
+      {SV(t->packed_dim())},
+      {graph.sizes_pc_of(v),
+       PushConstantDataInfo(&original_sizes_pc, sizeof(original_sizes_pc))}));
 
   return v;
 }