From 8ed7b2828a69a337a98429fedb0faba8b5414cde Mon Sep 17 00:00:00 2001
From: Vivek Trivedi <vivektrivedi@meta.com>
Date: Mon, 24 Nov 2025 11:03:05 -0800
Subject: [PATCH] Adding compile option warmup_execute_after_compile to
 optionally run execute rigtht after compilation to create command buffers.
 (#15962)

Summary:

This diff introduces a new compile option, `warmup_execute_after_compile`, which allows the graph to be executed once immediately after it is compiled. This option is disabled by default and can be enabled by setting the `warmup_execute_after_compile` flag to `true` in the `GraphConfig` object. When enabled, the `optional_warmup_execute` method makes a dummy call to `execute` function after the graph is prepacked.

Reviewed By: yipjustin

Differential Revision: D87781471
---
 backends/vulkan/runtime/VulkanBackend.cpp      | 8 ++++++++
 backends/vulkan/runtime/graph/ComputeGraph.cpp | 6 ++++++
 backends/vulkan/runtime/graph/ComputeGraph.h   | 6 ++++++
 backends/vulkan/runtime/graph/GraphConfig.h    | 4 ++++
 4 files changed, 24 insertions(+)

diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp
index cfa1242fbbf..fecef2598c7 100644
--- a/backends/vulkan/runtime/VulkanBackend.cpp
+++ b/backends/vulkan/runtime/VulkanBackend.cpp
@@ -179,6 +179,12 @@ GraphConfig get_graph_config(ArrayRef<CompileSpec>& compile_specs) {
         config.expect_dynamic_shapes = true;
       }
     }
+    if (strcmp(spec.key, "warmup_execute_after_compile") == 0) {
+      ET_CHECK_MSG(value_size == sizeof(uint8_t), "Unexpected value size!");
+      bool value = getBool(value_data);
+
+      config.warmup_execute_after_compile = value;
+    }
   }
 #ifdef ET_EVENT_TRACER_ENABLED
   config.enable_querypool = true;
@@ -579,6 +585,8 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
 
     compute_graph->prepack();
 
+    compute_graph->optional_warmup_execute();
+
     return Error::Ok;
   }
 
diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp
index d7f98d3244f..f96dbd6848f 100644
--- a/backends/vulkan/runtime/graph/ComputeGraph.cpp
+++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp
@@ -1107,6 +1107,12 @@ void ComputeGraph::prepack() {
   }
 }
 
+void ComputeGraph::optional_warmup_execute() {
+  if (config_.warmup_execute_after_compile) {
+    execute();
+  }
+}
+
 void ComputeGraph::execute() {
   if (deferred_cmd_list_.empty()) {
     context_->flush();
diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h
index b61bd4a51c0..7415a9dd2df 100644
--- a/backends/vulkan/runtime/graph/ComputeGraph.h
+++ b/backends/vulkan/runtime/graph/ComputeGraph.h
@@ -1033,6 +1033,12 @@ class ComputeGraph final {
    */
   void prepack();
 
+  //
+  // Optional Graph Execution
+  //
+
+  void optional_warmup_execute();
+
   //
   // Graph Execution
   //
diff --git a/backends/vulkan/runtime/graph/GraphConfig.h b/backends/vulkan/runtime/graph/GraphConfig.h
index 9a753775650..20d01362ef1 100644
--- a/backends/vulkan/runtime/graph/GraphConfig.h
+++ b/backends/vulkan/runtime/graph/GraphConfig.h
@@ -71,6 +71,10 @@ struct GraphConfig final {
   // many command buffers.
   size_t execute_max_cmds = 0;
 
+  // If true, then the graph will be executed once immediately after it is
+  // compiled.
+  bool warmup_execute_after_compile = false;
+
   vkapi::Adapter* external_adapter;
 
   // Generate a default graph config with pre-configured settings