[ET-VK] 4/n Split dispatches between multiple command buffers. Removing flush function call from submit_current_cmd_and_wait.

Pull Request resolved: #12525

This diff removes the `flush` function call from `submit_current_cmd_and_wait` and instead adds it to the `encode_prepack` method.

This change will allow better control over flush which is needed in follow up diffs.
ghstack-source-id: 296712817
@exported-using-ghexport

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

Author: trivedivivek

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -765,8 +765,6 @@ void ComputeGraph::submit_current_cmd_and_wait(const bool final_use) {
   context_->submit_cmd_to_gpu(fence.get_submit_handle(), final_use);
   fence.wait();
   context_->fences().return_fence(fence);
-
-  context_->flush();
 }
 
 void ComputeGraph::prepack() {
@@ -787,7 +785,12 @@ void ComputeGraph::prepack() {
       // Otherwise, just submit the current command buffer for execution and
       // proceed. This results in lower load latency at the cost of higher peak
       // memory usage.
-      reduce_peak_memory ? submit_current_cmd_and_wait() : submit_current_cmd();
+      if (reduce_peak_memory) {
+        submit_current_cmd_and_wait();
+        context_->flush();
+      } else {
+        submit_current_cmd();
+      }
       staging_nbytes_in_cmd_ = 0;
       context_->set_cmd();
       submitted = true;
@@ -797,6 +800,7 @@ void ComputeGraph::prepack() {
     i++;
   }
   submit_current_cmd_and_wait(/*final_use=*/true);
+  context_->flush();
   staging_nbytes_in_cmd_ = 0;
 }
 

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -847,8 +847,7 @@ class ComputeGraph final {
 
   /*
    * Submits the current command buffer in the Context to the GPU for execution,
-   * and wait for it to complete before returning. This function will also flush
-   * the Context after execution.
+   * and wait for it to complete before returning.
    */
   void submit_current_cmd_and_wait(const bool final_use = false);