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

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #9831 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/73/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/73/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/72/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/73/orig
@diff-train-skip-merge

---------

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

Author: pytorchbot

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