Introduce ParamsBindList to prevent needing to pass shared_ptr to bind parameter UBOs (#3150)

Summary:
Pull Request resolved: #3150

## Context

In keeping with the below changeset in this stack, this diff introduces the `ParamsBindList` structure to avoid storing shared pointers to `api::UniformParamsBuffer` objects in `ExecuteNode` and `PrepackNode`.

The idea is to store the binding information of each UPB instead of taking ownership of the UPB itself. There isn't really a need for `ExecuteNode` and `PrepackNode` to take ownership since `ComputeGraph` provides a guarantee that the UPBs will be in scope at the time of binding.

With this change, all `shared_ptr` members can be eliminated from `vTensor`, further reducing heap allocations and pointer chasing.

In the future I will change `prepack_nodes_` and `execute_nodes_` to store `PrepackNode` and `ExecuteNode` instances directly instead of storing unique pointers to them.

ghstack-source-id: 223225899
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D56357188

fbshipit-source-id: 5f4d1be900711753aa2cc035c044fe71f93d555b

Author: SS-JIA

backends/vulkan/runtime/api/Context.cpp

@@ -235,5 +235,11 @@ UniformParamsBuffer& UniformParamsBuffer::operator=(
   return *this;
 }
 
+ParamsBindList::ParamsBindList(
+    std::initializer_list<const api::BufferBindInfo> init_list) {
+  bind_infos.resize(init_list.size());
+  std::copy(init_list.begin(), init_list.end(), bind_infos.begin());
+}
+
 } // namespace api
 } // namespace vkcompute

backends/vulkan/runtime/api/Context.h

@@ -244,7 +244,7 @@ class UniformParamsBuffer final {
     }
   }
 
-  VulkanBuffer& buffer() {
+  const VulkanBuffer& buffer() const {
     return vulkan_buffer_;
   }
 
@@ -264,6 +264,12 @@ class UniformParamsBuffer final {
   }
 };
 
+struct ParamsBindList final {
+  std::vector<api::BufferBindInfo> bind_infos;
+
+  ParamsBindList(std::initializer_list<const api::BufferBindInfo> init_list);
+};
+
 class StorageBuffer final {
  private:
   Context* context_p_;
@@ -331,6 +337,10 @@ inline void arg_is_empty(bool& any_is_empty, const VulkanImage& image) {
   any_is_empty = any_is_empty || !image;
 }
 
+inline void arg_is_empty(bool& any_is_empty, const BufferBindInfo& bind_info) {
+  any_is_empty = any_is_empty || (bind_info.handle == VK_NULL_HANDLE);
+}
+
 /*
   Reports if any VulkanBuffer or VulkanImage argument in a variadic argument
   list does not have any memory associated with it.

backends/vulkan/runtime/api/Descriptor.cpp

@@ -15,6 +15,18 @@
 namespace vkcompute {
 namespace api {
 
+//
+// BufferBinding
+//
+
+BufferBindInfo::BufferBindInfo()
+    : handle(VK_NULL_HANDLE), offset(0u), range(0u) {}
+
+BufferBindInfo::BufferBindInfo(const VulkanBuffer& buffer_p)
+    : handle(buffer_p.handle()),
+      offset(buffer_p.mem_offset()),
+      range(buffer_p.mem_range()) {}
+
 //
 // DescriptorSet
 //
@@ -66,6 +78,21 @@ DescriptorSet& DescriptorSet::bind(
   return *this;
 }
 
+DescriptorSet& DescriptorSet::bind(
+    const uint32_t idx,
+    const BufferBindInfo& bind_info) {
+  DescriptorSet::ResourceBinding binder{};
+  binder.binding_idx = idx; // binding_idx
+  binder.descriptor_type = shader_layout_signature_[idx]; // descriptor_type
+  binder.is_image = false; // is_image
+  binder.resource_info.buffer_info.buffer = bind_info.handle; // buffer
+  binder.resource_info.buffer_info.offset = bind_info.offset; // offset
+  binder.resource_info.buffer_info.range = bind_info.range; // range
+  add_binding(binder);
+
+  return *this;
+}
+
 DescriptorSet& DescriptorSet::bind(
     const uint32_t idx,
     const VulkanImage& image) {

backends/vulkan/runtime/api/Descriptor.h

@@ -20,6 +20,20 @@
 namespace vkcompute {
 namespace api {
 
+/*
+ * Stores the binding information of a Vulkan Buffer so that the buffer can be
+ * bound at a later time. This struct should only be used if the buffer to be
+ * bound is guaranteed to be active at the time of binding.
+ */
+struct BufferBindInfo final {
+  VkBuffer handle;
+  VkDeviceSize offset;
+  VkDeviceSize range;
+
+  BufferBindInfo();
+  BufferBindInfo(const VulkanBuffer& buffer_p);
+};
+
 class DescriptorSet final {
  public:
   explicit DescriptorSet(VkDevice, VkDescriptorSet, ShaderLayout::Signature);
@@ -50,6 +64,7 @@ class DescriptorSet final {
   std::vector<ResourceBinding> bindings_;
 
  public:
+  DescriptorSet& bind(const uint32_t, const BufferBindInfo&);
   DescriptorSet& bind(const uint32_t, const VulkanBuffer&);
   DescriptorSet& bind(const uint32_t, const VulkanImage&);
 

backends/vulkan/runtime/api/Tensor.cpp

@@ -140,9 +140,9 @@ vTensor::vTensor(
       sizes_(sizes.begin(), sizes.end()),
       gpu_sizes_{calc_gpu_sizes(sizes, memory_layout_, storage_type)},
       // Utility Uniform Buffers that can be passed to shaders as arguments
-      cpu_sizes_uniform_(nullptr),
-      gpu_sizes_uniform_(nullptr),
-      extents_uniform_(nullptr),
+      cpu_sizes_uniform_(),
+      gpu_sizes_uniform_(),
+      extents_uniform_(),
       // Construct Tensor storage
       storage_(
           context,
@@ -189,33 +189,33 @@ api::VulkanBuffer& vTensor::buffer(
   return storage_.buffer_;
 }
 
-std::shared_ptr<api::UniformParamsBuffer> vTensor::cpu_sizes_ubo() {
-  if (!cpu_sizes_uniform_) {
-    cpu_sizes_uniform_.reset(new api::UniformParamsBuffer(
-        storage_.context_, api::utils::make_whcn_ivec4(sizes_)));
+const api::BufferBindInfo vTensor::cpu_sizes_ubo() {
+  if (!cpu_sizes_uniform_.buffer()) {
+    cpu_sizes_uniform_ = api::UniformParamsBuffer(
+        storage_.context_, api::utils::make_whcn_ivec4(sizes_));
   }
-  return cpu_sizes_uniform_;
+  return api::BufferBindInfo(cpu_sizes_uniform_.buffer());
 }
 
-std::shared_ptr<api::UniformParamsBuffer> vTensor::gpu_sizes_ubo() {
-  if (!gpu_sizes_uniform_) {
-    gpu_sizes_uniform_.reset(new api::UniformParamsBuffer(
-        storage_.context_, api::utils::make_whcn_ivec4(gpu_sizes_)));
+const api::BufferBindInfo vTensor::gpu_sizes_ubo() {
+  if (!gpu_sizes_uniform_.buffer()) {
+    gpu_sizes_uniform_ = api::UniformParamsBuffer(
+        storage_.context_, api::utils::make_whcn_ivec4(gpu_sizes_));
   }
-  return gpu_sizes_uniform_;
+  return api::BufferBindInfo(gpu_sizes_uniform_.buffer());
 }
 
-std::shared_ptr<api::UniformParamsBuffer> vTensor::extents_ubo() {
-  if (!extents_uniform_) {
-    extents_uniform_.reset(new api::UniformParamsBuffer(
+const api::BufferBindInfo vTensor::extents_ubo() {
+  if (!extents_uniform_.buffer()) {
+    extents_uniform_ = api::UniformParamsBuffer(
         storage_.context_,
         api::utils::uvec4(
             {storage_.extents_.data[0],
              storage_.extents_.data[1],
              storage_.extents_.data[2],
-             1u})));
+             1u}));
   }
-  return extents_uniform_;
+  return api::BufferBindInfo(extents_uniform_.buffer());
 }
 
 VmaAllocationCreateInfo vTensor::get_allocation_create_info() const {
@@ -258,16 +258,16 @@ void vTensor::update_size_metadata(const std::vector<int64_t>& new_sizes) {
   api::utils::uvec3 virtual_extents =
       create_image_extents(gpu_sizes_, storage_type(), memory_layout_);
 
-  if (cpu_sizes_uniform_) {
-    cpu_sizes_uniform_->update(api::utils::make_whcn_ivec4(sizes_));
+  if (cpu_sizes_uniform_.buffer()) {
+    cpu_sizes_uniform_.update(api::utils::make_whcn_ivec4(sizes_));
   }
 
-  if (gpu_sizes_uniform_) {
-    gpu_sizes_uniform_->update(api::utils::make_whcn_ivec4(gpu_sizes_));
+  if (gpu_sizes_uniform_.buffer()) {
+    gpu_sizes_uniform_.update(api::utils::make_whcn_ivec4(gpu_sizes_));
   }
 
-  if (extents_uniform_) {
-    extents_uniform_->update(api::utils::uvec4(
+  if (extents_uniform_.buffer()) {
+    extents_uniform_.update(api::utils::uvec4(
         {virtual_extents.data[0],
          virtual_extents.data[1],
          virtual_extents.data[2],

backends/vulkan/runtime/api/Tensor.h

@@ -118,17 +118,17 @@ class vTensor final {
 
   // A Vulkan uniform buffer containing the tensor sizes in WHCN that can be
   // passed into a shader.
-  std::shared_ptr<api::UniformParamsBuffer> cpu_sizes_uniform_;
+  api::UniformParamsBuffer cpu_sizes_uniform_;
 
   // A Vulkan uniform buffer containing the GPU tensor sizes in WHCN that can
   // be passed into a shader. GPU sizes refers to the sizes of the tensor after
   // padding has been applied to one dimension to align it to the next multiple
   // of 4.
-  std::shared_ptr<api::UniformParamsBuffer> gpu_sizes_uniform_;
+  api::UniformParamsBuffer gpu_sizes_uniform_;
 
   // A Vulkan uniform buffer containing the image extents of the underlying
   // image texture that can be passed into a shader.
-  std::shared_ptr<api::UniformParamsBuffer> extents_uniform_;
+  api::UniformParamsBuffer extents_uniform_;
 
   vTensorStorage storage_;
 
@@ -207,21 +207,21 @@ class vTensor final {
    * shader. Note that the UBO will be created the first time this function is
    * called.
    */
-  std::shared_ptr<api::UniformParamsBuffer> cpu_sizes_ubo();
+  const api::BufferBindInfo cpu_sizes_ubo();
 
   /*
    * Get a uniform buffer object containing the tensor GPU sizes to use in a
    * compute shader. Note that the UBO will be created the first time this
    * function is called.
    */
-  std::shared_ptr<api::UniformParamsBuffer> gpu_sizes_ubo();
+  const api::BufferBindInfo gpu_sizes_ubo();
 
   /*
    * Get a uniform buffer object containing the image extents to use in a
    * compute shader. Note that the UBO will be created the first time this
    * function is called.
    */
-  std::shared_ptr<api::UniformParamsBuffer> extents_ubo();
+  const api::BufferBindInfo extents_ubo();
 
   inline size_t numel() const {
     return api::utils::multiply_integers(sizes());

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -59,6 +59,7 @@ ComputeGraph::ComputeGraph(GraphConfig config)
           config_.contextConfig)},
       shared_objects_{},
       values_{},
+      param_ubos_{},
       prepack_nodes_{},
       execute_nodes_{},
       inputs_{},

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -93,6 +93,7 @@ class ComputeGraph final {
   std::unique_ptr<api::Context> context_;
   std::vector<SharedObject> shared_objects_;
   std::vector<Value> values_;
+  std::vector<api::UniformParamsBuffer> param_ubos_;
 
   std::vector<std::unique_ptr<PrepackNode>> prepack_nodes_;
   std::vector<std::unique_ptr<ExecuteNode>> execute_nodes_;
@@ -314,9 +315,9 @@ class ComputeGraph final {
   ValueRef set_output_tensor(const ValueRef idx, const bool use_staging = true);
 
   template <typename Block>
-  inline std::shared_ptr<api::UniformParamsBuffer> create_params_buffer(
-      const Block& data) {
-    return std::make_shared<api::UniformParamsBuffer>(context_.get(), data);
+  const api::BufferBindInfo create_params_buffer(const Block& data) {
+    param_ubos_.emplace_back(api::UniformParamsBuffer(context_.get(), data));
+    return api::BufferBindInfo(param_ubos_.back().buffer());
   }
 
   /*

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

@@ -20,7 +20,7 @@ ExecuteNode::ExecuteNode(
     const api::utils::uvec3& global_workgroup_size,
     const api::utils::uvec3& local_workgroup_size,
     const std::vector<ArgGroup>& args,
-    const std::vector<std::shared_ptr<api::UniformParamsBuffer>>& params,
+    const api::ParamsBindList& params,
     const ResizeFunction& resize_fn,
     const std::vector<ValueRef>& resize_args,
     const api::SpecVarList& spec_vars)
@@ -47,6 +47,7 @@ void ExecuteNode::encode(ComputeGraph* graph) {
   uint32_t idx = 0;
   idx = bind_values_to_descriptor_set(
       graph, args_, pipeline_barrier, descriptor_set, idx);
+
   bind_params_to_descriptor_set(params_, descriptor_set, idx);
 
   context->register_shader_dispatch(

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

@@ -54,7 +54,7 @@ class ExecuteNode final {
       const api::utils::uvec3& global_workgroup_size,
       const api::utils::uvec3& local_workgroup_size,
       const std::vector<ArgGroup>& args,
-      const std::vector<std::shared_ptr<api::UniformParamsBuffer>>& params,
+      const api::ParamsBindList& params,
       const ResizeFunction& resize_fn = nullptr,
       const std::vector<ValueRef>& resize_args = {},
       const api::SpecVarList& spec_vars = {});
@@ -74,7 +74,7 @@ class ExecuteNode final {
   const api::utils::uvec3 global_workgroup_size_;
   const api::utils::uvec3 local_workgroup_size_;
   const std::vector<ArgGroup> args_;
-  std::vector<std::shared_ptr<api::UniformParamsBuffer>> params_;
+  const api::ParamsBindList params_;
   const ResizeFunction resize_fn_;
   const std::vector<ValueRef> resize_args_;
   const api::SpecVarList spec_vars_;