Using push constants for softmax op. (#16112)

Summary:

The softmax operation has been updated to utilize push constants for efficient data transfer.

Differential Revision: D88097603

Author: trivedivivek

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

@@ -23,8 +23,10 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "tout", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "tin", DTYPE, STORAGE)}
 
-${layout_declare_ubo(B, "ivec3", "tout_limits")}
-${layout_declare_ubo(B, "ivec4", "tin_sizes")}
+layout(push_constant) uniform restrict Block {
+  ivec4 tin_sizes;
+  ivec3 tout_limits;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -139,9 +139,9 @@ void add_softmax_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {in, vkapi::kRead}},
       // Shader params buffers
-      {graph.logical_limits_ubo(out), graph.sizes_ubo(in)},
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(in), graph.logical_limits_pc_of(out)},
       // Specialization Constants
       {graph.packed_dim_of(out), reduce_dim_xyz, group_dim},
       // Resize Args