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

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

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.

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

---------

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

Author: trivedivivek

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -773,8 +773,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() {
@@ -795,7 +793,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;
@@ -805,6 +808,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);