[ET-VK] Using push constants for convolution prepack dispatches. (#11372)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #11306 by
@trivedivivek
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/110/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/110/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/109/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/110/orig
@diff-train-skip-merge

Co-authored-by: Vivek Trivedi <[email protected]>

Author: pytorchbot

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;
 }