Using push constants for add mm native op.

Differential Revision: D88097601

Pull Request resolved: #16115

Author: trivedivivek

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

@@ -21,13 +21,17 @@ ${layout_declare_tensor(B, "r", "mat1_tensor", DTYPE, "texture3d")}
 ${layout_declare_tensor(B, "r", "mat2_tensor", DTYPE, "texture3d")}
 $if HAS_BIAS:
   ${layout_declare_tensor(B, "r", "bias_tensor", DTYPE, "texture3d")}
-${layout_declare_ubo(B, "ivec4", "out_sizes")}
-${layout_declare_ubo(B, "ivec3", "out_limits")}
-${layout_declare_ubo(B, "ivec4", "mat1_sizes")}
-${layout_declare_ubo(B, "ivec4", "mat2_sizes")}
-$if HAS_BIAS:
-  ${layout_declare_ubo(B, "ivec4", "bias_sizes")}
-  ${layout_declare_ubo(B, "float", "alpha", "float", "beta")}
+
+layout(push_constant) uniform restrict Block {
+  ivec4 out_sizes;
+  ivec4 mat1_sizes;
+  ivec4 mat2_sizes;
+  ivec3 out_limits;
+  $if HAS_BIAS:
+    ivec4 bias_sizes;
+    float alpha;
+    float beta;
+};
 
 #include "indexing_utils.h"
 

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

@@ -182,16 +182,14 @@ void add_addmm_naive_texture_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {{mat1, mat2, self}, vkapi::kRead}},
       // Shader params buffers
-      {
-          graph.sizes_ubo(out),
-          graph.logical_limits_ubo(out),
-          graph.sizes_ubo(mat1),
-          graph.sizes_ubo(mat2),
-          graph.sizes_ubo(self),
-          graph.create_params_buffer(params),
-      },
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(out),
+       graph.sizes_pc_of(mat1),
+       graph.sizes_pc_of(mat2),
+       graph.logical_limits_pc_of(out),
+       graph.sizes_pc_of(self),
+       PushConstantDataInfo(&params, sizeof(params))},
       // Specialization Constants
       {graph.hashed_layout_of(out),
        graph.hashed_layout_of(mat1),

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

@@ -162,14 +162,12 @@ void add_matmul_naive_texture3d_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}},
       // Shader params buffers
-      {
-          graph.sizes_ubo(out),
-          graph.logical_limits_ubo(out),
-          graph.sizes_ubo(mat1),
-          graph.sizes_ubo(mat2),
-      },
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(out),
+       graph.sizes_pc_of(mat1),
+       graph.sizes_pc_of(mat2),
+       graph.logical_limits_pc_of(out)},
       // Specialization Constants
       {graph.hashed_layout_of(out),
        graph.hashed_layout_of(mat1),