From 48c882e236d00ddcbd05d2ad169c0e11f3d13cdb Mon Sep 17 00:00:00 2001
From: Jingyan Wang <jingywa@microsoft.com>
Date: Wed, 8 Jan 2025 21:47:02 +0000
Subject: [PATCH] Updated session option merge and unit test

---
 .../core/session/provider_bridge_ort.cc       |  6 +++++
 .../providers/tensorrt/tensorrt_basic_test.cc | 24 +++++++++----------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc
index 0aa93bce35..b6aba2b926 100644
--- a/onnxruntime/core/session/provider_bridge_ort.cc
+++ b/onnxruntime/core/session/provider_bridge_ort.cc
@@ -2233,6 +2233,7 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_TensorRT_V2,
     std::string context_cache_path = "";
     std::string embed_mode = "";
     if (options) {
+      LOGS_DEFAULT(WARNING) << "Propagating EP context options from session option since tensorrt options are empty";
       context_cache_enabled = (options->value).config_options.GetConfigOrDefault(kOrtSessionOptionEpContextEnable, "0") != "0";
       new_tensorrt_options.trt_dump_ep_context_model = context_cache_enabled;
       LOGS_DEFAULT(VERBOSE) << "Context cache enable: " << context_cache_enabled;
@@ -2246,13 +2247,18 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_TensorRT_V2,
         new_tensorrt_options.trt_ep_context_embed_mode = 1;
       } else if ("0" == embed_mode) {
         new_tensorrt_options.trt_ep_context_embed_mode = 0;
+        new_tensorrt_options.trt_engine_cache_enable = 1; // Enable engine cache if not embedded mode
       } else {
         LOGS_DEFAULT(VERBOSE) << "Invalid ep.context_embed_mode: " << embed_mode << " only 0 or 1 allowed. Set to 1.";
       }
       LOGS_DEFAULT(VERBOSE) << "User specified context cache embed mode: " << embed_mode;
+      
     }
     factory = onnxruntime::TensorrtProviderFactoryCreator::Create(&new_tensorrt_options);
   } else {
+    if (ep_context_cache_enabled_from_provider_options && ep_context_cache_enabled_from_sess_options) {
+      LOGS_DEFAULT(WARNING) << "Provider options provided EP context configs, the EP Context configs from session options will be ignored";
+    }
     factory = onnxruntime::TensorrtProviderFactoryCreator::Create(tensorrt_options);
   }
 #else
diff --git a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc
index 3355d91a43..4ac61adf2c 100644
--- a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc
+++ b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc
@@ -459,10 +459,6 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) {
   so.session_logid = sess_log_id;
   RunOptions run_options;
   run_options.run_tag = so.session_logid;
-  so.config_options.AddConfigEntry("ep.context_enable", "1");
-  so.config_options.AddConfigEntry("ep.context_file_path", "EP_Context_model.onnx");
-  so.config_options.AddConfigEntry("ep.context_embed_mode", "0"); // Default 1
-  InferenceSession session_object{so, GetEnvironment()};
   auto cuda_provider = DefaultCudaExecutionProvider();
   auto cpu_allocator = cuda_provider->CreatePreferredAllocators()[1];
   std::vector<int64_t> dims_mul_x = {1, 3, 2};
@@ -488,13 +484,7 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) {
 
   /*
    * Test case 1: Dump context model
-   *
-   * provider options=>
-   *   trt_ep_context_file_path = "EP_Context_model.onnx"
-   *
-   * expected result =>
-   *   context model "EP_Context_model.onnx" should be created in current directory
-   * --- New UT
+   * Note: *.engine file need to be removed before running this test case
    * session options =>
    *  ep.context_enable = "1"
    *  ep.context_file_path = "EP_Context_model.onnx"
@@ -506,10 +496,16 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) {
    * expected result =>
    *  context model "EP_Context_model.onnx" should be created in current directory
    */
+  so.config_options.AddConfigEntry("ep.context_enable", "1");
+  so.config_options.AddConfigEntry("ep.context_file_path", "EP_Context_model.onnx");
+  so.config_options.AddConfigEntry("ep.context_embed_mode", "0");
+  so.session_log_severity_level = 0;
+  InferenceSession session_object{so, GetEnvironment()};
   OrtTensorRTProviderOptionsV2 params;
   params.trt_engine_cache_enable = 1;
   params.trt_dump_ep_context_model = 1;
-  // params.trt_ep_context_file_path = "EP_Context_model.onnx";
+  params.trt_ep_context_file_path = "EP_Context_model.onnx";
+  params.trt_ep_context_embed_mode = 0;
   std::unique_ptr<IExecutionProvider> execution_provider = TensorrtExecutionProviderWithOptions(&params);
   EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
   auto status = session_object.Load(model_name);
@@ -552,10 +548,12 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) {
   OrtTensorRTProviderOptionsV2 params2;
   params2.trt_engine_cache_enable = 1;
   params2.trt_dump_ep_context_model = 1;
+  params2.trt_ep_context_file_path = "context_model_folder/EPContextNode_test_ctx.onnx";
+  params2.trt_ep_context_embed_mode = 0;
   // params2.trt_engine_cache_prefix = "TRT_engine_cache";
   // params2.trt_engine_cache_path = "engine_cache_folder";  // due to dump_ep_context_model = 1, the new cache path is ./context_model_folder/engine_cache_folder
   // params2.trt_ep_context_file_path = "context_model_folder";
-  params2.trt_engine_cache_path = "context_model_folder";
+  // params2.trt_engine_cache_path = "context_model_folder";
   execution_provider = TensorrtExecutionProviderWithOptions(&params2);
   EXPECT_TRUE(session_object2.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
   status = session_object2.Load(model_name);