Update on "[ET-VK] 5/n Split dispatches between multiple command buffers. Track previously submitted command buffers in context and add function to execute all previous command buffers."

The diff adds changes to store command buffers submitted with final_use set to false.

Storing these buffers is necessary for `execute()` function. Since, `encode_execute()` function is typically called once but `execute()` can be called multiple times, `submit_all_non_final_cmds` function is added so all recorded command buffers with `final_use = False` can be called multiple times in `execute()`.

#### Key Changes
*   Added a flag `execute_pending_first_submission` to the `ComputeGraph` class to track whether execute nodes have been freshly encoded and need to be submitted first.
*   Added a new function `submit_all_non_final_cmds` to the `Context` class, which submits all non-final command buffers to the GPU.
*   Modified the `submit_cmd_to_gpu` function to add the submitted command buffer to the `non_final_cmds_` list if it's not marked as final use.
*   Updated the `execute` function in `ComputeGraph` to submit all non-final command buffers before executing the graph.

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

[ghstack-poisoned]

Author: trivedivivek

backends/vulkan/runtime/api/Context.cpp

@@ -214,44 +214,14 @@ void Context::submit_cmd_to_gpu(VkFence fence_handle, const bool final_use) {
 
     submit_count_ = 0u;
   }
-
-  if (!final_use) {
-    non_final_cmds_.emplace_back(std::move(cmd_));
-  }
-}
-
-void Context::submit_all_non_final_cmds(VkFence fence_handle) {
-  VkSemaphore local_prev_semaphore = VK_NULL_HANDLE;
-
-  for (uint32_t i = 0; i < non_final_cmds_.size(); i++) {
-    auto& cmd = non_final_cmds_[i];
-    VkSemaphore wait_semaphore = local_prev_semaphore;
-    VkSemaphore signal_semaphore = cmd.get_signal_semaphore();
-    local_prev_semaphore = signal_semaphore;
-
-    if (cmd) {
-      cmd.end();
-      adapter_p_->submit_cmd(
-          queue_,
-          cmd.get_submit_handle(false),
-          i == (non_final_cmds_.size() - 1) ? fence_handle : VK_NULL_HANDLE,
-          wait_semaphore,
-          signal_semaphore);
-    }
-  }
 }
 
 void Context::flush() {
-  VK_CHECK(vkQueueWaitIdle(queue()));
+  VK_CHECK(vkQueueWaitIdle(queue().handle));
 
   command_pool_.flush();
   descriptor_pool_.flush();
 
-  for (auto& cmd : non_final_cmds_) {
-    cmd.invalidate();
-  }
-  non_final_cmds_.clear();
-
   // If there is an existing command buffer, invalidate it
   if (cmd_) {
     cmd_.invalidate();

backends/vulkan/runtime/api/Context.h

@@ -68,8 +68,6 @@ class Context final {
   // Command buffers submission
   std::mutex cmd_mutex_;
   vkapi::CommandBuffer cmd_;
-  // List of submitted command buffers, not marked as final use.
-  std::vector<vkapi::CommandBuffer> non_final_cmds_;
   // Semaphore for the previously submitted command buffer, if any
   VkSemaphore prev_semaphore_;
   uint32_t submit_count_;
@@ -92,8 +90,8 @@ class Context final {
     return device_;
   }
 
-  inline VkQueue queue() {
-    return queue_.handle;
+  inline vkapi::Adapter::Queue& queue() {
+    return queue_;
   }
 
   // Device Caches
@@ -232,7 +230,9 @@ class Context final {
       VkFence fence_handle = VK_NULL_HANDLE,
       const bool final_use = false);
 
-  void submit_all_non_final_cmds(VkFence fence_handle = VK_NULL_HANDLE);
+  vkapi::CommandBuffer& extract_cmd() {
+    return cmd_;
+  }
 
   void flush();
 

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -158,6 +158,7 @@ ComputeGraph::~ComputeGraph() {
 
   prepack_nodes_.clear();
   execute_nodes_.clear();
+  deferred_cmd_list_.clear();
 
   context_->flush();
 }
@@ -767,6 +768,30 @@ void ComputeGraph::submit_current_cmd_and_wait(const bool final_use) {
   context_->fences().return_fence(fence);
 }
 
+void ComputeGraph::submit_deferred_cmds() {
+  VkSemaphore prev_semaphore = VK_NULL_HANDLE;
+  vkapi::VulkanFence fence = context_->fences().get_fence();
+
+  for (uint32_t i = 0; i < deferred_cmd_list_.size(); i++) {
+    auto& cmd = deferred_cmd_list_[i];
+    VkSemaphore wait_semaphore = prev_semaphore;
+    VkSemaphore signal_semaphore = cmd.get_signal_semaphore();
+    prev_semaphore = signal_semaphore;
+
+    if (cmd) {
+      cmd.end();
+      context_->adapter_ptr()->submit_cmd(
+          context_->queue(),
+          cmd.get_submit_handle(false),
+          i == (deferred_cmd_list_.size() - 1) ? fence.get_submit_handle() : VK_NULL_HANDLE,
+          wait_semaphore,
+          signal_semaphore);
+    }
+  }
+  fence.wait();
+  context_->fences().return_fence(fence);
+}
+
 void ComputeGraph::prepack() {
   int i = 0;
   bool submitted = false;
@@ -805,6 +830,7 @@ void ComputeGraph::prepack() {
 }
 
 void ComputeGraph::encode_execute() {
+  deferred_cmd_list_.clear();
   context_->flush();
   context_->set_cmd(/*reusable = */ true);
 
@@ -814,21 +840,11 @@ void ComputeGraph::encode_execute() {
     node->encode(this);
   }
 
-  // Indicate execute nodes have been freshly encoded and needs to be submitted
-  // first
-  execute_pending_first_submission = true;
+  deferred_cmd_list_.emplace_back(std::move(context_->extract_cmd()));
 }
 
 void ComputeGraph::execute() {
-  if (execute_pending_first_submission) {
-    submit_current_cmd_and_wait(/*final_use=*/false);
-    execute_pending_first_submission = false;
-  } else {
-    vkapi::VulkanFence fence = context_->fences().get_fence();
-    context_->submit_all_non_final_cmds(fence.get_submit_handle());
-    fence.wait();
-    context_->fences().return_fence(fence);
-  }
+  submit_deferred_cmds();
   execute_count_++;
 }
 

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -193,6 +193,9 @@ class ComputeGraph final {
   // Utility constexpr to express byte quantities
   constexpr static size_t MB = 1024 * 1024;
 
+  // List of command buffers deferred for submission
+  std::vector<vkapi::CommandBuffer> deferred_cmd_list_;
+
  protected:
   size_t values_in_use_ = 0;
   size_t execute_count_ = 0;
@@ -204,10 +207,6 @@ class ComputeGraph final {
   // current Context's command buffer is submitted now.
   size_t staging_nbytes_in_cmd_ = 0;
 
-  // Flag to indicate if execute nodes have been freshly encoded and have not
-  // been submitted yet.
-  bool execute_pending_first_submission = false;
-
  public:
   //
   // Accessors
@@ -855,6 +854,11 @@ class ComputeGraph final {
    */
   void submit_current_cmd_and_wait(const bool final_use = false);
 
+  /*
+   * Submits all the commands gathered in deferred_cmd_bufs_ to the GPU.
+   */
+  void submit_deferred_cmds();
+
  public:
   //
   // Graph Prepacking