[ET-VK] 1/n Split dispatches between multiple command buffers. Add semaphore support in command buffer. (#12519)

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* #12525
* #12523
* #12522
* __->__ #12519

execute stage
## Context

This following diffs aims to improve the performance of the Executorch
Vulkan backend by adding a mechanism to issue multiple command buffers
in prepack and execute function, so GPU work is issues while CPU is
still working on issuing new work.

## This Diff

### Summary

This diff is the first in a series of diffs that aim to split dispatches
between multiple command buffers and add semaphore support in the
command buffer.

The changes in this diff include:

* Adding a `VkSemaphore` parameter to the `CommandBuffer` constructor in
`vk_api/Command.cpp` and `vk_api/Command.h` to support signaling when
the command buffer has completed execution.
* Modifying the `CommandBuffer` constructor in `vk_api/Command.h` and
`vk_api/Command.cpp` to include the `VkSemaphore` parameter.
* Updating the `CommandBuffer` object in `api/Context.cpp` to include
the `VkSemaphore` parameter.

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

---------

Co-authored-by: Stephen Jia <[email protected]>

Author: trivedivivek

backends/vulkan/runtime/api/Context.cpp

@@ -38,7 +38,7 @@ Context::Context(vkapi::Adapter* adapter, const ContextConfig& config)
       querypool_(config_.query_pool_config, nullptr),
       // Command buffer submission
       cmd_mutex_{},
-      cmd_(VK_NULL_HANDLE, 0u),
+      cmd_(VK_NULL_HANDLE, VK_NULL_HANDLE, 0u),
       submit_count_{0u},
       // Memory Management
       buffer_clearlist_mutex_{},

backends/vulkan/runtime/vk_api/Command.cpp

@@ -20,28 +20,34 @@ namespace vkapi {
 
 CommandBuffer::CommandBuffer(
     VkCommandBuffer handle,
+    VkSemaphore semaphore,
     const VkCommandBufferUsageFlags flags)
     : handle_(handle),
+      signal_semaphore_(semaphore),
       flags_(flags),
       state_(CommandBuffer::State::NEW),
       bound_{} {}
 
 CommandBuffer::CommandBuffer(CommandBuffer&& other) noexcept
     : handle_(other.handle_),
+      signal_semaphore_(other.signal_semaphore_),
       flags_(other.flags_),
-      state_(CommandBuffer::State::INVALID),
+      state_(other.state_),
       bound_(other.bound_) {
   other.handle_ = VK_NULL_HANDLE;
+  other.signal_semaphore_ = VK_NULL_HANDLE;
   other.bound_.reset();
 }
 
 CommandBuffer& CommandBuffer::operator=(CommandBuffer&& other) noexcept {
   handle_ = other.handle_;
+  signal_semaphore_ = other.signal_semaphore_;
   flags_ = other.flags_;
   state_ = other.state_;
   bound_ = other.bound_;
 
   other.handle_ = VK_NULL_HANDLE;
+  other.signal_semaphore_ = VK_NULL_HANDLE;
   other.bound_.reset();
   other.state_ = CommandBuffer::State::INVALID;
 
@@ -304,6 +310,12 @@ CommandPool::~CommandPool() {
   if (pool_ == VK_NULL_HANDLE) {
     return;
   }
+  for (auto& semaphore : semaphores_) {
+    if (semaphore != VK_NULL_HANDLE) {
+      vkDestroySemaphore(device_, semaphore, nullptr);
+    }
+  }
+
   vkDestroyCommandPool(device_, pool_, nullptr);
 }
 
@@ -314,14 +326,15 @@ CommandBuffer CommandPool::get_new_cmd(bool reusable) {
   allocate_new_batch(config_.cmd_pool_batch_size);
 
   VkCommandBuffer handle = buffers_[in_use_];
+  VkSemaphore semaphore = semaphores_[in_use_];
 
   VkCommandBufferUsageFlags cmd_flags = 0u;
   if (!reusable) {
     cmd_flags |= VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
   }
 
   in_use_++;
-  return CommandBuffer(handle, cmd_flags);
+  return CommandBuffer(handle, semaphore, cmd_flags);
 }
 
 void CommandPool::flush() {
@@ -337,6 +350,7 @@ void CommandPool::allocate_new_batch(const uint32_t count) {
   }
 
   buffers_.resize(buffers_.size() + count);
+  semaphores_.resize(buffers_.size() + count);
 
   const VkCommandBufferAllocateInfo allocate_info{
       VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO, // sType
@@ -348,6 +362,17 @@ void CommandPool::allocate_new_batch(const uint32_t count) {
 
   VK_CHECK(vkAllocateCommandBuffers(
       device_, &allocate_info, buffers_.data() + in_use_));
+
+  const VkSemaphoreCreateInfo semaphoreCreateInfo = {
+      VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO, nullptr, 0};
+
+  for (uint32_t i = 0; i < count; i++) {
+    VK_CHECK(vkCreateSemaphore(
+        device_,
+        &semaphoreCreateInfo,
+        nullptr,
+        semaphores_.data() + in_use_ + i));
+  }
 }
 
 } // namespace vkapi

backends/vulkan/runtime/vk_api/Command.h

@@ -26,7 +26,10 @@ namespace vkapi {
 
 class CommandBuffer final {
  public:
-  explicit CommandBuffer(VkCommandBuffer, const VkCommandBufferUsageFlags);
+  explicit CommandBuffer(
+      VkCommandBuffer,
+      VkSemaphore,
+      const VkCommandBufferUsageFlags);
 
   CommandBuffer(const CommandBuffer&) = delete;
   CommandBuffer& operator=(const CommandBuffer&) = delete;
@@ -70,6 +73,8 @@ class CommandBuffer final {
 
  private:
   VkCommandBuffer handle_;
+  // Semaphore to signal when the command buffer has completed execution
+  VkSemaphore signal_semaphore_;
   VkCommandBufferUsageFlags flags_;
   State state_;
   Bound bound_;
@@ -81,6 +86,7 @@ class CommandBuffer final {
 
   inline void invalidate() {
     handle_ = VK_NULL_HANDLE;
+    signal_semaphore_ = VK_NULL_HANDLE;
     bound_.reset();
   }
 
@@ -100,6 +106,10 @@ class CommandBuffer final {
 
   VkCommandBuffer get_submit_handle(const bool final_use = false);
 
+  VkSemaphore get_signal_semaphore() const {
+    return signal_semaphore_;
+  }
+
   inline operator bool() const {
     return handle_ != VK_NULL_HANDLE;
   }
@@ -130,6 +140,8 @@ class CommandPool final {
   // New Buffers
   std::mutex mutex_;
   std::vector<VkCommandBuffer> buffers_;
+  // Semaphores corresponding to the command buffers
+  std::vector<VkSemaphore> semaphores_;
   size_t in_use_;
 
  public: