Add aux functions to grow the staging buffer size (unused for now)

Summary: Add aux functions to the staging device to add the ability to grow the staging buffer.

Reviewed By: corporateshark

Differential Revision: D49329271

fbshipit-source-id: 89b7b5bd02864dd725eaeb2c0214d6dc2601e087

src/igl/vulkan/VulkanStagingDevice.cpp

@@ -28,14 +28,15 @@ VulkanStagingDevice::VulkanStagingDevice(VulkanContext& ctx) : ctx_(ctx) {
 
   const auto& limits = ctx_.getVkPhysicalDeviceProperties().limits;
 
-  // Use default value of 256 MB, and clamp it to the max limits
-  stagingBufferSize_ = std::min(limits.maxStorageBufferRange, 256u * 1024u * 1024u);
+  // Use value of 256MB (limited by some architectures), and clamp it to the max limits
+  maxBufferCapacity_ = std::min(limits.maxStorageBufferRange, 256u * 1024u * 1024u);
+
+  // Use default value maxBufferCapacity_
+  stagingBufferSize_ = maxBufferCapacity_;
 
   // Initialize the block list with a block that spans the entire staging buffer
   regions_.push_front(MemoryRegion{0, stagingBufferSize_, VulkanImmediateCommands::SubmitHandle()});
 
-  bufferCapacity_ = stagingBufferSize_;
-
   stagingBuffer_ =
       ctx_.createBuffer(stagingBufferSize_,
                         VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
@@ -479,5 +480,51 @@ void VulkanStagingDevice::waitAndReset() {
   regions_.push_front({0, stagingBufferSize_, VulkanImmediateCommands::SubmitHandle()});
 }
 
+bool VulkanStagingDevice::shouldGrowStagingBuffer(uint32_t sizeNeeded) const {
+  return !stagingBuffer_ || (sizeNeeded > stagingBufferSize_);
+}
+
+uint32_t VulkanStagingDevice::nextSize(uint32_t requestedSize) const {
+  return std::min(getAlignedSize(requestedSize), maxBufferCapacity_);
+}
+
+void VulkanStagingDevice::growStagingBuffer(uint32_t minimumSize) {
+  IGL_PROFILER_FUNCTION();
+
+  IGL_ASSERT(minimumSize <= maxBufferCapacity_);
+
+#if IGL_VULKAN_DEBUG_STAGING_DEVICE
+  IGL_LOG_INFO("Growing staging buffer from %u to %u bytes\n", stagingBufferSize_, minimumSize);
+#endif
+
+  waitAndReset();
+
+  // De-allocates the staging buffer
+  stagingBuffer_ = nullptr;
+
+  // If the size of the new staging buffer plus the size of the existing one is larger than the
+  // limit imposed by some architectures on buffers that are device and host visible, we need to
+  // wait for the current buffer to be destroyed before we can allocate a new one
+  if ((minimumSize + stagingBufferSize_) > maxBufferCapacity_) {
+    // Wait for the current buffer to be destroyed on the device
+    ctx_.waitDeferredTasks();
+  }
+
+  stagingBufferSize_ = minimumSize;
+
+  // Create a new staging buffer with the new size
+  stagingBuffer_ =
+      ctx_.createBuffer(stagingBufferSize_,
+                        VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
+                        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
+                        nullptr,
+                        "Buffer: staging buffer");
+  IGL_ASSERT(stagingBuffer_.get());
+
+  // Clear out the old regions and add one that represents the entire buffer
+  regions_.clear();
+  regions_.push_front({0, stagingBufferSize_, VulkanImmediateCommands::SubmitHandle()});
+}
+
 } // namespace vulkan
 } // namespace igl

src/igl/vulkan/VulkanStagingDevice.h

@@ -64,20 +64,32 @@ class VulkanStagingDevice final {
    * @return The offset of the free memory block on the staging buffer and the size of the block
    * found.
    */
-  MemoryRegion nextFreeBlock(uint32_t size);
+  [[nodiscard]] MemoryRegion nextFreeBlock(uint32_t size);
 
   uint32_t getAlignedSize(uint32_t size) const;
 
   /// @brief Waits for all memory blocks to become available and resets the staging device's
   /// internal state
   void waitAndReset();
 
+  /// @brief Returns true if the staging buffer cannot store the size requested
+  [[nodiscard]] bool shouldGrowStagingBuffer(uint32_t sizeNeeded) const;
+
+  /// @brief Returns the next size to allocate for the staging buffer given the requested size
+  [[nodiscard]] uint32_t nextSize(uint32_t requestedSize) const;
+
+  /// @brief Grows the staging buffer to a size that is at least as large as the requested size
+  void growStagingBuffer(uint32_t minimumSize);
+
  private:
   VulkanContext& ctx_;
   std::shared_ptr<VulkanBuffer> stagingBuffer_;
   std::unique_ptr<VulkanImmediateCommands> immediate_;
-  uint32_t stagingBufferSize_;
-  uint32_t bufferCapacity_;
+
+  /// @brief Current size of the staging buffer
+  uint32_t stagingBufferSize_ = 0;
+  /// @brief Maximum staging buffer capacity, limited by some architectures
+  uint32_t maxBufferCapacity_ = 0;
 
   /**
    * @brief Stores the used and unused blocks of memory in the staging buffer. There is no