From 40fd5742d6939a4b91b7edb207b282abded12beb Mon Sep 17 00:00:00 2001
From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com>
Date: Fri, 6 Dec 2024 12:38:47 -0800
Subject: [PATCH] [ET-VK] Replace Uniform buffers with push constants for view
 op

This diff replaces uniform buffers with push constants for view op in the Vulkan backend of Executorch. The changes include updating the GLSL code to use push constants instead of uniform buffers and updating the C++ code to pass the sizes as push constants to the shader.

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

[ghstack-poisoned]
---
 backends/vulkan/runtime/graph/ops/glsl/view.glsl | 6 ++++--
 backends/vulkan/runtime/graph/ops/impl/View.cpp  | 6 ++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/backends/vulkan/runtime/graph/ops/glsl/view.glsl b/backends/vulkan/runtime/graph/ops/glsl/view.glsl
index 8d45e65b396..599879514e3 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/view.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/view.glsl
@@ -19,8 +19,10 @@ layout(std430) buffer;
 ${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)}
 ${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)}
 
-${layout_declare_ubo(2, "ivec4", "out_sizes")}
-${layout_declare_ubo(3, "ivec4", "in_sizes")}
+layout(push_constant) uniform PRECISION restrict Block {
+  ivec4 out_sizes;
+  ivec4 in_sizes;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
diff --git a/backends/vulkan/runtime/graph/ops/impl/View.cpp b/backends/vulkan/runtime/graph/ops/impl/View.cpp
index 060696a4fa6..fc5c7075222 100644
--- a/backends/vulkan/runtime/graph/ops/impl/View.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/View.cpp
@@ -76,12 +76,14 @@ void add_view_node(
       {{out, vkapi::MemoryAccessType::WRITE},
        {in, vkapi::MemoryAccessType::READ}},
       // Parameter Buffers
-      {t_out->sizes_ubo(), t_in->sizes_ubo()},
+      {},
       // Specialization Constants
       {SV(t_in->packed_dim()), SV(t_out->packed_dim())},
       // Resizing Logic
       resize_view_node,
-      {sizes}));
+      {sizes},
+      // Push Constants
+      {{graph.sizes_pc_of(out), graph.sizes_pc_of(in)}}));
 }
 
 void view(ComputeGraph& graph, const std::vector<ValueRef>& args) {