Limit AOT function expansion

onnxruntime/core/framework/graph_partitioner.cc

@@ -1277,14 +1277,111 @@ static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* lay
   return Status::OK();
 }
 
+constexpr size_t kAotFunctionExpansionRatio = 10;
+
+static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::GraphProto& graph);
+
+static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::AttributeProto& attribute) {
+  SafeInt<size_t> node_count = 0;
+  if (attribute.has_g()) {
+    node_count += CountNodesIncludingSubgraphs(attribute.g());
+  }
+  for (const auto& attribute_graph : attribute.graphs()) {
+    node_count += CountNodesIncludingSubgraphs(attribute_graph);
+  }
+
+  return node_count;
+}
+
+static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::GraphProto& graph) {
+  SafeInt<size_t> node_count = graph.node_size();
+  for (const auto& node : graph.node()) {
+    for (const auto& attribute : node.attribute()) {
+      node_count += CountNodesIncludingSubgraphs(attribute);
+    }
+  }
+
+  return node_count;
+}
+
+static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::FunctionProto& function) {
+  SafeInt<size_t> node_count = function.node_size();
+  for (const auto& node : function.node()) {
+    for (const auto& attribute : node.attribute()) {
+      node_count += CountNodesIncludingSubgraphs(attribute);
+    }
+  }
+  for (const auto& default_attribute : function.attribute_proto()) {
+    node_count += CountNodesIncludingSubgraphs(default_attribute);
+  }
+
+  return node_count;
+}
+
+static size_t CountNodesIncludingSubgraphs(const Graph& graph) {
+  SafeInt<size_t> node_count = graph.NumberOfNodes();
+  for (const auto& node : graph.Nodes()) {
+    for (const auto& subgraph : node.GetSubgraphs()) {
+      node_count += CountNodesIncludingSubgraphs(*subgraph);
+    }
+  }
+
+  return node_count;
+}
+
+struct FunctionExpansionCost {
+  size_t node_count;
+  size_t proto_bytes;
+};
+
+static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& cost) {
+  if (const auto* function_body = node.GetFunctionBody()) {
+    const auto graph_proto = function_body->Body().ToGraphProto();
+    cost = {CountNodesIncludingSubgraphs(graph_proto), graph_proto.ByteSizeLong()};
+    return Status::OK();
+  }
+
+  ONNX_NAMESPACE::FunctionProto function_proto;
+  ORT_RETURN_IF_NOT(node.TryGetFunctionProto(function_proto),
+                    "Unable to get function body for node '", node.Name(), "'.");
+  std::string accounting_prefix = "_inlfunc_" + node.OpType();
+  accounting_prefix.append(32, '_');
+  function_utils::Specialize(function_proto, node, accounting_prefix);
+  SafeInt<size_t> node_count = function_proto.node_size();
+  SafeInt<size_t> proto_bytes = 0;
+  for (const auto& function_node : function_proto.node()) {
+    proto_bytes += function_node.ByteSizeLong();
+    for (const auto& attribute : function_node.attribute()) {
+      node_count += CountNodesIncludingSubgraphs(attribute);
+    }
+  }
+  cost = {node_count, proto_bytes};
+  return Status::OK();
+}
+
+static size_t CountModelNodes(const Model& model) {
+  SafeInt<size_t> node_count = CountNodesIncludingSubgraphs(model.MainGraph());
+  for (const auto& [function_id, function_template] : model.GetModelLocalFunctionTemplates()) {
+    ORT_UNUSED_PARAMETER(function_id);
+    node_count += CountNodesIncludingSubgraphs(*function_template->onnx_func_proto_);
+  }
+
+  return node_count;
+}
+
 static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_providers,
                                      const KernelRegistryManager& kernel_registry_mgr,
                                      Graph& graph,
                                      const GraphOptimizerRegistry& graph_optimizer_registry,
                                      const logging::Logger& logger,
                                      const CheckLoadCancellationFn& check_load_cancellation_fn,
                                      InlinedHashSet<std::string>& not_inlined,
-                                     size_t& inlined_count) {
+                                     InlinedHashSet<std::string>& budget_limited_functions,
+                                     size_t& inlined_count,
+                                     size_t expansion_node_budget,
+                                     size_t& expanded_node_count,
+                                     size_t expansion_byte_budget,
+                                     size_t& expanded_proto_bytes) {
   // handle testing edge case where optimizers or constant lifting results in graph with no nodes.
   // doing it here saves all providers checking for this in GetCapability
   if (graph.NumberOfNodes() == 0) {
@@ -1302,7 +1399,12 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide
                                                  logger,
                                                  check_load_cancellation_fn,
                                                  not_inlined,
-                                                 inlined_count));
+                                                 budget_limited_functions,
+                                                 inlined_count,
+                                                 expansion_node_budget,
+                                                 expanded_node_count,
+                                                 expansion_byte_budget,
+                                                 expanded_proto_bytes));
     }
   }
 
@@ -1355,6 +1457,29 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide
     auto* node = graph.GetNode(node_index);
     if (node != nullptr) {
       if (claimed_by_ep.count(node_index) == 0) {
+        auto function_id = function_utils::GetFunctionIdentifier(node->Domain(), node->OpType(), node->Overload());
+        if (budget_limited_functions.count(function_id) != 0) {
+          continue;
+        }
+
+        FunctionExpansionCost expansion_cost{};
+        ORT_RETURN_IF_ERROR(GetFunctionExpansionCost(*node, expansion_cost));
+        if (expansion_cost.node_count > expansion_node_budget - expanded_node_count) {
+          LOGS(logger, WARNING) << "AOT function inlining exceeds the node expansion limit of "
+                                << expansion_node_budget << ". Retaining function '" << function_id << "'.";
+          ORT_IGNORE_RETURN_VALUE(not_inlined.insert(function_id));
+          ORT_IGNORE_RETURN_VALUE(budget_limited_functions.insert(std::move(function_id)));
+          continue;
+        }
+        if (expansion_cost.proto_bytes > expansion_byte_budget - expanded_proto_bytes) {
+          LOGS(logger, WARNING) << "AOT function inlining exceeds the protobuf expansion limit of "
+                                << expansion_byte_budget << " bytes. Retaining function '" << function_id << "'.";
+          ORT_IGNORE_RETURN_VALUE(not_inlined.insert(function_id));
+          ORT_IGNORE_RETURN_VALUE(budget_limited_functions.insert(std::move(function_id)));
+          continue;
+        }
+        expanded_node_count += expansion_cost.node_count;
+        expanded_proto_bytes += expansion_cost.proto_bytes;
         ORT_RETURN_IF_ERROR(graph.InlineFunction(*node));
         ++inlined_count;
       } else {
@@ -1742,7 +1867,14 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model,
   auto check_load_cancellation_fn = [this]() -> bool { return IsLoadCancellationFlagSet(); };
 
   auto& graph = model.MainGraph();
+  const size_t expansion_node_budget =
+      static_cast<size_t>(SafeInt<size_t>(CountModelNodes(model)) * kAotFunctionExpansionRatio);
+  const size_t expansion_byte_budget =
+      static_cast<size_t>(SafeInt<size_t>(model.ModelProtoByteSize()) * kAotFunctionExpansionRatio);
+  size_t expanded_node_count = 0;
+  size_t expanded_proto_bytes = 0;
   InlinedHashSet<std::string> not_inlined;
+  InlinedHashSet<std::string> budget_limited_functions;
   do {
     size_t inlined_count = 0;
     ORT_RETURN_IF_ERROR(InlineFunctionsAOTImpl(execution_providers,
@@ -1752,14 +1884,27 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model,
                                                logger,
                                                check_load_cancellation_fn,
                                                not_inlined,
-                                               inlined_count));
+                                               budget_limited_functions,
+                                               inlined_count,
+                                               expansion_node_budget,
+                                               expanded_node_count,
+                                               expansion_byte_budget,
+                                               expanded_proto_bytes));
 
     if (inlined_count == 0) {
       break;
     }
     ORT_RETURN_IF_ERROR(graph.Resolve());
   } while (true);
 
+  if (!budget_limited_functions.empty()) {
+    return ORT_MAKE_STATUS(
+        ONNXRUNTIME, FAIL,
+        "AOT function inlining exceeded an expansion limit for ",
+        budget_limited_functions.size(),
+        " function(s). Initialization cannot continue because fallback inlining would exceed the same limit.");
+  }
+
   model.RemoveLocalFunctionsProtos(not_inlined);
 
   LOGS(logger, INFO)

onnxruntime/core/graph/model.h

@@ -199,6 +199,8 @@ class Model {
   const Graph& MainGraph() const noexcept;
 
 #if !defined(ORT_MINIMAL_BUILD)
+  size_t ModelProtoByteSize() const noexcept { return model_proto_.ByteSizeLong(); }
+
   // Get model's serialization proto data.
   ONNX_NAMESPACE::ModelProto ToProto() const;