From 0346905ed94614eeb3da22d5724c1bdb3aa1f0c8 Mon Sep 17 00:00:00 2001
From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com>
Date: Mon, 2 Jun 2025 19:10:29 -0700
Subject: [PATCH] [ET-VK] Using push constants for convolution prepack
 dispatches.

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

[ghstack-poisoned]
---
 .../graph/ops/glsl/conv2d_dw_prepack_weights.glsl        | 5 +----
 .../runtime/graph/ops/glsl/conv2d_prepack_weights.glsl   | 5 +----
 .../graph/ops/glsl/conv_transpose2d_prepack_weights.glsl | 5 +----
 backends/vulkan/runtime/graph/ops/impl/Convolution.cpp   | 9 +++++----
 4 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.glsl
index 49ce76423d5..f5361d40b66 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.glsl
+++ b/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;
 };
 
diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.glsl
index 4e8bff94947..d2f3f615f74 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.glsl
+++ b/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;
 };
 
diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.glsl
index df8589e737f..0b10683cee4 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.glsl
+++ b/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;
 };
 
diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp
index 32f478fa5bd..cdce33affb3 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp
@@ -211,6 +211,7 @@ ValueRef prepack_weights(
   vkapi::ShaderInfo shader =
       get_conv2d_shader(graph, *t, /*prepack_weights = */ true, method, vref);
 
+  const auto sizes_pc = utils::make_ivec4(original_sizes, /*reverse = */ true);
   graph.prepack_nodes().emplace_back(new PrepackNode(
       graph,
       shader,
@@ -218,11 +219,11 @@ ValueRef prepack_weights(
       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(&sizes_pc, sizeof(sizes_pc))}));
 
   return v;
 }