Using push constants for binary scalar op parameter.

Differential Revision: D88097606

Pull Request resolved: pytorch#16114

Author: trivedivivek

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

@@ -29,7 +29,9 @@ ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 ${layout_declare_ubo(B, "BufferMetadata", "outp")}
 ${layout_declare_ubo(B, "BufferMetadata", "inp")}
 
-${layout_declare_ubo(B, "float", "scalar_value")}
+layout(push_constant) uniform restrict Block {
+  float scalar_value;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -30,7 +30,9 @@ ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 ${layout_declare_ubo(B, "TextureMetadata", "outp")}
 ${layout_declare_ubo(B, "TextureMetadata", "inp")}
 
-${layout_declare_ubo(B, "float", "scalar_value")}
+layout(push_constant) uniform restrict Block {
+  float scalar_value;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -48,10 +48,7 @@ void add_binary_scalar_op_node(
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
   add_dtype_suffix(kernel_name, graph.dtype_of(in));
 
-  vkapi::ParamsBindList param_ubos = {
-      graph.meta_ubo(out),
-      graph.meta_ubo(in),
-      graph.create_params_buffer(scalar_val)};
+  vkapi::ParamsBindList param_ubos = {graph.meta_ubo(out), graph.meta_ubo(in)};
 
   graph.execute_nodes().emplace_back(new DynamicDispatchNode(
       graph,
@@ -63,7 +60,7 @@ void add_binary_scalar_op_node(
       // Shader params buffers
       param_ubos,
       // Push Constants
-      {},
+      {PushConstantDataInfo(&scalar_val, sizeof(scalar_val))},
       // Specialization Constants
       {},
       // Resize Args