From 2a36fd4f6e0aabfb4a00c4c20ab5b9bd474c60b8 Mon Sep 17 00:00:00 2001
From: Hector Li <hecli@microsoft.com>
Date: Thu, 12 Dec 2024 21:12:02 -0800
Subject: [PATCH] Fix the ctx_gen tool to make sure all generated ctx.onnx have
 max_size (#23097)

### Description
Fix the qnn_ctx_gen tool to make sure all generated ctx.onnx have max_size
---
 onnxruntime/test/qnn_ctx_gen/main.cc | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/onnxruntime/test/qnn_ctx_gen/main.cc b/onnxruntime/test/qnn_ctx_gen/main.cc
index d568d5e786..3be0bd253c 100644
--- a/onnxruntime/test/qnn_ctx_gen/main.cc
+++ b/onnxruntime/test/qnn_ctx_gen/main.cc
@@ -33,8 +33,11 @@ static void CheckStatus(const Status& status) {
 
 // from the last context cache Onnx model, find the EPContext node with main_context=1,
 // and get the QNN context binary file name, this context binary contains all graphs from all Onnx models
+// get the max spill fill buffer size
 static void GetLastContextBinaryFileName(const std::basic_string<ORTCHAR_T> last_onnx_ctx_file,
-                                         std::string& last_ctx_bin_file) {
+                                         std::string& last_ctx_bin_file,
+                                         int64_t& max_size) {
+  max_size = 0;
   std::shared_ptr<Model> ctx_model;
   CheckStatus(Model::Load(ToPathString(last_onnx_ctx_file), ctx_model, nullptr,
                           (*((OrtEnv*)*ort_env.get())->GetEnvironment().GetLoggingManager()).DefaultLogger()));
@@ -43,6 +46,7 @@ static void GetLastContextBinaryFileName(const std::basic_string<ORTCHAR_T> last
     if (node.OpType() == "EPContext") {
       NodeAttrHelper node_helper(node);
       int64_t is_main_context = node_helper.Get("main_context", static_cast<int64_t>(0));
+      max_size = node_helper.Get("max_size", static_cast<int64_t>(0));
       if (1 == is_main_context) {
         last_ctx_bin_file = node_helper.Get("ep_cache_context", "");
         return;
@@ -55,7 +59,8 @@ static void GetLastContextBinaryFileName(const std::basic_string<ORTCHAR_T> last
 // the last QNN context binary file
 // Remove not used QNN context binary file, only keep the last one which contains all graphs
 static void UpdateEpContextModel(const std::vector<std::basic_string<ORTCHAR_T>>& ep_ctx_files,
-                                 const std::string& last_qnn_ctx_binary_file_name) {
+                                 const std::string& last_qnn_ctx_binary_file_name,
+                                 int64_t max_size) {
   for (auto ep_ctx_file : ep_ctx_files) {
     std::shared_ptr<Model> ctx_model;
     auto path_str = ToPathString(ep_ctx_file);
@@ -75,6 +80,8 @@ static void UpdateEpContextModel(const std::vector<std::basic_string<ORTCHAR_T>>
           std::remove(file_path.string().c_str());
           node.ClearAttribute("ep_cache_context");
           node.AddAttribute("ep_cache_context", last_qnn_ctx_binary_file_name);
+          node.ClearAttribute("max_size");
+          node.AddAttribute("max_size", max_size);
         }
       }
     }
@@ -181,7 +188,8 @@ int real_main(int argc, char* argv[]) {
 
     // Get the last context binary file name
     std::string last_qnn_ctx_binary_file_name;
-    GetLastContextBinaryFileName(ep_ctx_files.back(), last_qnn_ctx_binary_file_name);
+    int64_t max_size = 0;
+    GetLastContextBinaryFileName(ep_ctx_files.back(), last_qnn_ctx_binary_file_name, max_size);
     std::cout << "The last context binary file: " << last_qnn_ctx_binary_file_name << std::endl;
     if (last_qnn_ctx_binary_file_name.empty()) {
       throw Ort::Exception("Can't find QNN context binary file from the Onnx model.", OrtErrorCode::ORT_FAIL);
@@ -191,7 +199,7 @@ int real_main(int argc, char* argv[]) {
     // Update generated context cache Onnx model to make the main EPContext node point to
     // the last QNN context binary file
     // Remove not used QNN context binary file, only keep the last one which contains all graphs
-    UpdateEpContextModel(ep_ctx_files, last_qnn_ctx_binary_file_name);
+    UpdateEpContextModel(ep_ctx_files, last_qnn_ctx_binary_file_name, max_size);
   }
   ORT_CATCH(const Ort::Exception& e) {
     fprintf(stderr, "Failed to generate context cache file: %s \n", e.what());