[ET-VK] Add PushConstantDataInfo and vector to hold push constants data in DispatchNode.

Pull Request resolved: #7223

This diff adds a new class called `PushConstantDataInfo` to the `DispatchNode` class in the Vulkan backend for Executorch.
This class represents a push constant data entry, which can either be a shared pointer to a tensor's uniform data with an attribute or data with a maximum size of 16 bytes.
The `write` method is also added to this class, which writes the data to a destination buffer.
ghstack-source-id: 257899587
@exported-using-ghexport

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

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

Author: pytorchbot

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

@@ -14,6 +14,22 @@
 
 namespace vkcompute {
 
+uint32_t PushConstantDataInfo::write(
+    void* dst,
+    const uint32_t dst_offset,
+    const uint32_t max_dst_size) const {
+  if (tensorUniformData != nullptr) {
+    return tensorUniformData->write_attribute(
+        dst, dst_offset, max_dst_size, payload_.attr);
+  }
+
+  VK_CHECK_COND(
+      (dst_offset + payload_.dataSize) <= max_dst_size,
+      "Attempting to write push constant data outside data boundary.");
+  memcpy((uint8_t*)dst + dst_offset, payload_.data, payload_.dataSize);
+  return payload_.dataSize;
+}
+
 DispatchNode::DispatchNode(
     ComputeGraph& graph,
     const vkapi::ShaderInfo& shader,
@@ -23,13 +39,15 @@ DispatchNode::DispatchNode(
     const vkapi::ParamsBindList& params,
     const vkapi::SpecVarList& spec_vars,
     const ResizeFunction& resize_fn,
-    const std::vector<ValueRef>& resize_args)
+    const std::vector<ValueRef>& resize_args,
+    const std::vector<PushConstantDataInfo>& push_constants)
     : ExecuteNode(resize_fn, resize_args, args, shader.kernel_name),
       shader_(shader),
       global_workgroup_size_(global_workgroup_size),
       local_workgroup_size_(local_workgroup_size),
       params_(params),
-      spec_vars_(spec_vars) {
+      spec_vars_(spec_vars),
+      push_constants_(push_constants) {
   graph.update_descriptor_counts(shader, /*execute = */ true);
 }
 
@@ -57,8 +75,22 @@ void DispatchNode::encode(ComputeGraph* graph) {
 
   bind_params_to_descriptor_set(params_, descriptor_set, idx);
 
+  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);
+  }
   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);
 
   context->report_shader_dispatch_end();
 }

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

@@ -18,6 +18,51 @@ namespace vkcompute {
 
 class ComputeGraph;
 
+constexpr uint32_t kMaxPushConstantSize = 128;
+/*
+ * Represents a push constant data entry
+ * Which is either shared pointer to a tensor's uniform data with an attribute
+ * Or data with a maximum size of 16 bytes
+ */
+class PushConstantDataInfo {
+  std::shared_ptr<api::vTensor::UniformData> tensorUniformData;
+  union Payload {
+    struct {
+      api::vTensor::Attribute attr;
+    };
+    struct {
+      uint8_t data[16];
+      uint32_t dataSize;
+    };
+  };
+
+  Payload payload_;
+
+ public:
+  explicit PushConstantDataInfo(
+      const std::shared_ptr<api::vTensor::UniformData>& tensorUniformData,
+      api::vTensor::Attribute attr)
+      : tensorUniformData(tensorUniformData) {
+    payload_.attr = attr;
+  }
+
+  explicit PushConstantDataInfo(const void* data, uint32_t dataLen)
+      : tensorUniformData(nullptr) {
+    VK_CHECK_COND(
+        dataLen <= 16, "Single push constant data size must be <= 16 bytes");
+    payload_.dataSize = dataLen;
+    memcpy(payload_.data, data, payload_.dataSize);
+  }
+
+  /*
+   * Function writes push constant data to the destination buffer
+   */
+  uint32_t write(
+      void* dst,
+      const uint32_t dst_offset,
+      const uint32_t max_dst_size) const;
+};
+
 /*
  * Represents a single shader execution op in a ML model.
  */
@@ -34,7 +79,8 @@ class DispatchNode final : public ExecuteNode {
       const vkapi::ParamsBindList& params,
       const vkapi::SpecVarList& spec_vars = {},
       const ResizeFunction& resize_fn = nullptr,
-      const std::vector<ValueRef>& resize_args = {});
+      const std::vector<ValueRef>& resize_args = {},
+      const std::vector<PushConstantDataInfo>& push_constants = {});
 
   ~DispatchNode() override = default;
 
@@ -46,6 +92,7 @@ class DispatchNode final : public ExecuteNode {
   const utils::uvec3 local_workgroup_size_;
   const vkapi::ParamsBindList params_;
   const vkapi::SpecVarList spec_vars_;
+  const std::vector<PushConstantDataInfo> push_constants_;
 
  public:
   operator bool() const {