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

Pull Request resolved: #9831

This diff replaces Uniform buffers with push constants for the native layer norm op in the Vulkan backend of Executorch. The changes include updating the shader 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: [D70943355](https://our.internmc.facebook.com/intern/diff/D70943355/)
ghstack-source-id: 275813550

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) {