From 096307c64b789354e2c20e915e918e66d911dd3e Mon Sep 17 00:00:00 2001
From: Dmitri Smirnov <yuslepukhin@users.noreply.github.com>
Date: Tue, 7 Nov 2023 13:46:42 -0800
Subject: [PATCH] Do not run AOT function inlining when the model does not
 define any local functions (#18302)

### Description
Check if the model defines any local functions.
if not, skip AOT inlining including any schema based functions.
The latter would be inlined during partitioning.

### Motivation and Context
This prevents calls GetCapability() to EPs and enhahces  compatibility.
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: Pranav Sharma <prs@microsoft.com>
---
 onnxruntime/core/framework/graph_partitioner.cc | 17 ++++++++++++++++-
 onnxruntime/core/framework/graph_partitioner.h  |  4 +++-
 onnxruntime/core/session/inference_session.cc   |  4 +++-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc
index b028596fe4e6d..e4fe0c7564548 100644
--- a/onnxruntime/core/framework/graph_partitioner.cc
+++ b/onnxruntime/core/framework/graph_partitioner.cc
@@ -798,7 +798,16 @@ static Status PartitionOrtFormatModel(const PartitionParams& partition_params,
 
 Status GraphPartitioner::InlineFunctionsAOT(Model& model,
                                             const ExecutionProviders& execution_providers,
-                                            const KernelRegistryManager& kernel_registry_manager) const {
+                                            const KernelRegistryManager& kernel_registry_manager,
+                                            const logging::Logger& logger) const {
+  const auto local_functions_num = model.GetModelLocalFunctionTemplates().size();
+  const bool is_there_local_functions = local_functions_num > 0;
+
+  if (!is_there_local_functions) {
+    LOGS(logger, INFO) << "This model does not have any local functions defined. AOT Inlining is not performed";
+    return Status::OK();
+  }
+
   auto& graph = model.MainGraph();
   InlinedHashSet<std::string> not_inlined;
   do {
@@ -818,6 +827,12 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model,
 
   model.RemoveLocalFunctionsProtos(not_inlined);
 
+  LOGS(logger, INFO)
+      << "AOT inlining completed. (" << (local_functions_num - model.GetModelLocalFunctionTemplates().size())
+      << ") functions of ("
+      << local_functions_num
+      << ") pruned.";
+
   return Status::OK();
 }
 
diff --git a/onnxruntime/core/framework/graph_partitioner.h b/onnxruntime/core/framework/graph_partitioner.h
index c1fa46de9145d..4fc85c2588260 100644
--- a/onnxruntime/core/framework/graph_partitioner.h
+++ b/onnxruntime/core/framework/graph_partitioner.h
@@ -48,10 +48,12 @@ class GraphPartitioner {
   /// <param name="model">model instance</param>
   /// <param name="execution_providers">execution providers considered</param>
   /// <param name="kernel_registry_manager">registry manager</param>
+  /// <param name="logger">session logger</param>
   /// <returns></returns>
   Status InlineFunctionsAOT(Model& model,
                             const ExecutionProviders& execution_providers,
-                            const KernelRegistryManager& kernel_registry_manager) const;
+                            const KernelRegistryManager& kernel_registry_manager,
+                            const logging::Logger& logger) const;
 #endif
 
  private:
diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc
index 1163be27b1685..ccedc71b9119a 100644
--- a/onnxruntime/core/session/inference_session.cc
+++ b/onnxruntime/core/session/inference_session.cc
@@ -1029,7 +1029,9 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
               kOrtSessionOptionsDisableAheadOfTimeFunctionInlining, "0") == "1";
       !disable_aot_function_inlining) {
     ORT_RETURN_IF_ERROR_SESSIONID_(partitioner.InlineFunctionsAOT(*model_,
-                                                                  execution_providers_, kernel_registry_manager_));
+                                                                  execution_providers_,
+                                                                  kernel_registry_manager_,
+                                                                  *session_logger_));
   }
 
   auto apply_transformer_once = [](const GraphTransformer& transformer, const logging::Logger& logger,