[ET-VK] Adding PushConstantDataInfo input to PrepackNode class. (#9048)

* [ET-VK] Moving PushConstantData class implementation a separate h and cpp file.

Pull Request resolved: #8648

This diff moves the implementation of the `PushConstantData` class from the `DispatchNode` file to a separate `PushConstantData.h` and `PushConstantData.cpp` files. This is done to enable push constants for PrepackNodes in follow up changes.
ghstack-source-id: 270375637
@exported-using-ghexport

Differential Revision: [D70102031](https://our.internmc.facebook.com/intern/diff/D70102031/)

* [ET-VK] Adding PushConstantDataInfo input to PrepackNode class.

Pull Request resolved: #8649

This diff adds a new input to the PrepackNode class called PushConstantDataInfo. This input is used to pass push constant data to the shader.
ghstack-source-id: 270375636
@exported-using-ghexport

Differential Revision: [D70102043](https://our.internmc.facebook.com/intern/diff/D70102043/)

---------

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

Author: pytorchbot

backends/vulkan/runtime/graph/ops/PrepackNode.cpp

@@ -32,15 +32,17 @@ PrepackNode::PrepackNode(
     const ValueRef tref,
     const ValueRef packed,
     const vkapi::ParamsBindList& params,
-    const vkapi::SpecVarList& spec_vars)
+    const vkapi::SpecVarList& spec_vars,
+    const std::vector<PushConstantDataInfo>& push_constants)
     : shader_(shader),
       noop_shader_(get_noop_shader(graph, packed)),
       global_workgroup_size_(global_workgroup_size),
       local_workgroup_size_(local_workgroup_size),
       tref_(tref),
       packed_(packed),
       params_(params),
-      spec_vars_(spec_vars) {
+      spec_vars_(spec_vars),
+      push_constants_(push_constants) {
   graph.update_descriptor_counts(shader, /*execute = */ false);
   graph.update_descriptor_counts(noop_shader_, /*execute = */ false);
 }
@@ -75,10 +77,20 @@ void PrepackNode::encode(ComputeGraph* graph) {
 
   std::unique_lock<std::mutex> cmd_lock = context->dispatch_lock();
 
+  std::array<uint8_t, kMaxPushConstantSize> push_constants_data;
+  uint32_t push_constants_offset = 0;
+
+  for (const auto& push_constant : push_constants_) {
+    push_constants_offset += push_constant.write(
+        push_constants_data.data(),
+        push_constants_offset,
+        kMaxPushConstantSize);
+  }
+
   {
     vkapi::PipelineBarrier pipeline_barrier{};
     vkapi::DescriptorSet descriptor_set = context->get_descriptor_set(
-        shader_, local_workgroup_size_, spec_vars_, 0u);
+        shader_, local_workgroup_size_, spec_vars_, push_constants_offset);
 
     uint32_t idx = 0;
     bind_tensor_to_descriptor_set(
@@ -91,7 +103,12 @@ void PrepackNode::encode(ComputeGraph* graph) {
     bind_params_to_descriptor_set(params_, descriptor_set, idx);
 
     context->register_shader_dispatch(
-        descriptor_set, pipeline_barrier, shader_, global_workgroup_size_);
+        descriptor_set,
+        pipeline_barrier,
+        shader_,
+        global_workgroup_size_,
+        push_constants_data.data(),
+        push_constants_offset);
   }
 
   // Submit a compute shader that performs a no-op with the packed tensor in

backends/vulkan/runtime/graph/ops/PrepackNode.h

@@ -10,6 +10,7 @@
 
 #include <executorch/backends/vulkan/runtime/api/api.h>
 
+#include <executorch/backends/vulkan/runtime/graph/containers/PushConstantData.h>
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
 namespace vkcompute {
@@ -34,7 +35,8 @@ class PrepackNode final {
       const ValueRef tref,
       const ValueRef packed,
       const vkapi::ParamsBindList& params,
-      const vkapi::SpecVarList& spec_vars = {});
+      const vkapi::SpecVarList& spec_vars = {},
+      const std::vector<PushConstantDataInfo>& push_constants = {});
 
   ~PrepackNode() = default;
 
@@ -54,6 +56,7 @@ class PrepackNode final {
   const ValueRef packed_;
   const vkapi::ParamsBindList params_;
   const vkapi::SpecVarList spec_vars_;
+  const std::vector<PushConstantDataInfo> push_constants_;
 
  private:
   api::StagingBuffer create_staging_buffer(ComputeGraph* graph);