[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,6 +214,31 @@ 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() {
@@ -222,6 +247,11 @@ void Context::flush() {
   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,6 +68,8 @@ 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_;
@@ -230,6 +232,8 @@ 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);
+
   void flush();
 
 #ifdef VULKAN_DEBUG

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -815,13 +815,22 @@ void ComputeGraph::encode_execute() {
   for (std::unique_ptr<ExecuteNode>& node : execute_nodes_) {
     node->encode(this);
   }
+
+  // Indicate execute nodes have been freshly encoded and needs to be submitted
+  // first
+  execute_pending_first_submission = true;
 }
 
 void ComputeGraph::execute() {
-  vkapi::VulkanFence fence = context_->fences().get_fence();
-  context_->submit_cmd_to_gpu(fence.get_submit_handle());
-  fence.wait();
-  context_->fences().return_fence(fence);
+  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);
+  }
   execute_count_++;
 }
 

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -198,6 +198,10 @@ 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 = true;
+
  public:
   //
   // Accessors