[ET-VK] Replace Uniform buffers with push constants for native layer norm op

Differential Revision: D70943355

Pull Request resolved: #9831

Author: trivedivivek

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

@@ -27,9 +27,11 @@ ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "t_weight", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "t_bias", DTYPE, STORAGE)}
 
-${layout_declare_ubo(B, "ivec3", "out_limits")}
-${layout_declare_ubo(B, "ivec4", "sizes")}
-${layout_declare_ubo(B, "float", "epsilon")}
+layout(push_constant) uniform PRECISION restrict Block {
+  ivec3 out_limits;
+  ivec4 sizes;
+  float epsilon;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -101,19 +101,20 @@ void add_native_layer_norm_node(
         vkapi::MemoryAccessType::WRITE},
        {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
-      {
-          t_out->logical_limits_ubo(),
-          t_out->sizes_ubo(),
-          graph.create_params_buffer(epsilon),
-      },
+      {},
       // Specialization Constants
       {
           t_input->hashed_layout(),
           t_out->hashed_layout(),
       },
       // Resizing Logic
       resize_native_layer_norm_node,
-      {normalized_shape}));
+      {normalized_shape},
+      {
+          graph.logical_limits_pc_of(out_val->at(0)),
+          graph.sizes_pc_of(out_val->at(0)),
+          PushConstantDataInfo(&epsilon, sizeof(epsilon)),
+      }));
 }
 
 void native_layer_norm(ComputeGraph& graph, const std::vector<ValueRef>& args) {