From 89c83a384941e321297951f7fac18dd1168d88eb Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 16 Sep 2026 16:20:55 -0700 Subject: [PATCH 01/11] Limit AOT function expansion --- .../core/framework/graph_partitioner.cc | 106 +++++++++++++++++- onnxruntime/test/framework/function_test.cc | 74 ++++++++++++ 2 files changed, 177 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 2850b96025dcc..fcdf093d40913 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1277,6 +1277,78 @@ 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 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 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 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; +} + +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); + cost = {CountNodesIncludingSubgraphs(function_proto), function_proto.ByteSizeLong()}; + return Status::OK(); +} + +static size_t CountModelNodes(const ONNX_NAMESPACE::ModelProto& model) { + SafeInt node_count = CountNodesIncludingSubgraphs(model.graph()); + for (const auto& function : model.functions()) { + node_count += CountNodesIncludingSubgraphs(function); + } + + return node_count; +} + static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_providers, const KernelRegistryManager& kernel_registry_mgr, Graph& graph, @@ -1284,7 +1356,11 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide const logging::Logger& logger, const CheckLoadCancellationFn& check_load_cancellation_fn, InlinedHashSet& not_inlined, - size_t& inlined_count) { + 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 +1378,11 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide logger, check_load_cancellation_fn, not_inlined, - inlined_count)); + inlined_count, + expansion_node_budget, + expanded_node_count, + expansion_byte_budget, + expanded_proto_bytes)); } } @@ -1355,6 +1435,15 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide auto* node = graph.GetNode(node_index); if (node != nullptr) { if (claimed_by_ep.count(node_index) == 0) { + FunctionExpansionCost expansion_cost{}; + ORT_RETURN_IF_ERROR(GetFunctionExpansionCost(*node, expansion_cost)); + ORT_RETURN_IF(expansion_cost.node_count > expansion_node_budget - expanded_node_count, + "AOT function inlining exceeds the node expansion limit of ", expansion_node_budget, "."); + ORT_RETURN_IF(expansion_cost.proto_bytes > expansion_byte_budget - expanded_proto_bytes, + "AOT function inlining exceeds the protobuf expansion limit of ", expansion_byte_budget, + " bytes."); + 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,6 +1831,13 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, auto check_load_cancellation_fn = [this]() -> bool { return IsLoadCancellationFlagSet(); }; auto& graph = model.MainGraph(); + const auto model_proto = model.ToProto(); + const size_t expansion_node_budget = + static_cast(SafeInt(CountModelNodes(model_proto)) * kAotFunctionExpansionRatio); + const size_t expansion_byte_budget = + static_cast(SafeInt(model_proto.ByteSizeLong()) * kAotFunctionExpansionRatio); + size_t expanded_node_count = 0; + size_t expanded_proto_bytes = 0; InlinedHashSet not_inlined; do { size_t inlined_count = 0; @@ -1752,7 +1848,11 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, logger, check_load_cancellation_fn, not_inlined, - inlined_count)); + inlined_count, + expansion_node_budget, + expanded_node_count, + expansion_byte_budget, + expanded_proto_bytes)); if (inlined_count == 0) { break; diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 88089b6cbde8a..e45b4ad6103a0 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -120,6 +120,80 @@ static Status LoadModel(const char* source) { return session_object.Load(sstr); } +static ONNX_NAMESPACE::ModelProto CreateFunctionExpansionModel(size_t body_node_count, size_t call_count) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(8); + auto* onnx_opset = model.add_opset_import(); + onnx_opset->set_domain(""); + onnx_opset->set_version(13); + auto* local_opset = model.add_opset_import(); + local_opset->set_domain("local"); + local_opset->set_version(1); + + auto set_float_value = [](ONNX_NAMESPACE::ValueInfoProto& value, const std::string& name) { + value.set_name(name); + auto* tensor_type = value.mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor_type->mutable_shape()->add_dim()->set_dim_value(1); + }; + + auto* graph = model.mutable_graph(); + graph->set_name("function_expansion"); + set_float_value(*graph->add_input(), "input"); + + auto* function = model.add_functions(); + function->set_domain("local"); + function->set_name("Expand"); + function->add_input("function_input"); + function->add_output("function_output"); + auto* function_opset = function->add_opset_import(); + function_opset->set_domain(""); + function_opset->set_version(13); + + std::string previous_value = "function_input"; + for (size_t i = 0; i < body_node_count; ++i) { + auto* node = function->add_node(); + node->set_op_type("Identity"); + node->add_input(previous_value); + previous_value = i + 1 == body_node_count ? "function_output" : "body_" + std::to_string(i); + node->add_output(previous_value); + } + + previous_value = "input"; + for (size_t i = 0; i < call_count; ++i) { + auto* node = graph->add_node(); + node->set_domain("local"); + node->set_op_type("Expand"); + node->add_input(previous_value); + previous_value = "call_" + std::to_string(i); + node->add_output(previous_value); + } + set_float_value(*graph->add_output(), previous_value); + + return model; +} + +static Status InitializeFunctionExpansionModel(size_t body_node_count, size_t call_count) { + std::string serialized_model; + ORT_RETURN_IF_NOT(CreateFunctionExpansionModel(body_node_count, call_count).SerializeToString(&serialized_model), + "Failed to serialize function expansion model."); + + SessionOptions session_options; + InferenceSession session{session_options, GetEnvironment()}; + std::istringstream stream(serialized_model); + ORT_RETURN_IF_ERROR(session.Load(stream)); + return session.Initialize(); +} + +TEST(FunctionTest, AotInliningLimitsFunctionExpansion) { + constexpr size_t body_node_count = 21; + ASSERT_STATUS_OK(InitializeFunctionExpansionModel(body_node_count, 1)); + + const auto status = InitializeFunctionExpansionModel(body_node_count, 21); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("AOT function inlining exceeds the")); +} + // A recursive/cyclic chain of model-local functions can be rejected by either layer: // ONNX 1.22+ detects the cycle in its own model checker ("Cycle detected in model-local // function references"), which runs before ORT's equivalent check ("must not be recursive"). From a2b183793bc9d5487044008b5e3cd1d960a1930d Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 17 Sep 2026 16:01:06 -0700 Subject: [PATCH 02/11] Address comments --- .../core/framework/graph_partitioner.cc | 47 +++++++--- onnxruntime/core/graph/model.h | 2 + onnxruntime/test/framework/function_test.cc | 90 ++++++++++++++++--- 3 files changed, 117 insertions(+), 22 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index fcdf093d40913..8d237a8699288 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1318,6 +1318,17 @@ static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::FunctionProto& return node_count; } +static size_t CountNodesIncludingSubgraphs(const Graph& graph) { + SafeInt 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; @@ -1340,10 +1351,11 @@ static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& return Status::OK(); } -static size_t CountModelNodes(const ONNX_NAMESPACE::ModelProto& model) { - SafeInt node_count = CountNodesIncludingSubgraphs(model.graph()); - for (const auto& function : model.functions()) { - node_count += CountNodesIncludingSubgraphs(function); +static size_t CountModelNodes(const Model& model) { + SafeInt 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; @@ -1435,13 +1447,25 @@ 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 (not_inlined.count(function_id) != 0) { + continue; + } + FunctionExpansionCost expansion_cost{}; ORT_RETURN_IF_ERROR(GetFunctionExpansionCost(*node, expansion_cost)); - ORT_RETURN_IF(expansion_cost.node_count > expansion_node_budget - expanded_node_count, - "AOT function inlining exceeds the node expansion limit of ", expansion_node_budget, "."); - ORT_RETURN_IF(expansion_cost.proto_bytes > expansion_byte_budget - expanded_proto_bytes, - "AOT function inlining exceeds the protobuf expansion limit of ", expansion_byte_budget, - " bytes."); + 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(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(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)); @@ -1831,11 +1855,10 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, auto check_load_cancellation_fn = [this]() -> bool { return IsLoadCancellationFlagSet(); }; auto& graph = model.MainGraph(); - const auto model_proto = model.ToProto(); const size_t expansion_node_budget = - static_cast(SafeInt(CountModelNodes(model_proto)) * kAotFunctionExpansionRatio); + static_cast(SafeInt(CountModelNodes(model)) * kAotFunctionExpansionRatio); const size_t expansion_byte_budget = - static_cast(SafeInt(model_proto.ByteSizeLong()) * kAotFunctionExpansionRatio); + static_cast(SafeInt(model.ModelProtoByteSize()) * kAotFunctionExpansionRatio); size_t expanded_node_count = 0; size_t expanded_proto_bytes = 0; InlinedHashSet not_inlined; diff --git a/onnxruntime/core/graph/model.h b/onnxruntime/core/graph/model.h index a6ce2a3a0659c..eead39f1c78b0 100644 --- a/onnxruntime/core/graph/model.h +++ b/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; diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index e45b4ad6103a0..fa8fd787ac051 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -19,6 +19,7 @@ #include "core/session/inference_session.h" #include "test/common/tensor_op_test_utils.h" +#include "test/capturing_sink.h" #include "test/unittest_util/framework_test_utils.h" #include "test/internal_testing_ep/internal_testing_execution_provider.h" #include "test/test_environment.h" @@ -173,25 +174,94 @@ static ONNX_NAMESPACE::ModelProto CreateFunctionExpansionModel(size_t body_node_ return model; } -static Status InitializeFunctionExpansionModel(size_t body_node_count, size_t call_count) { +static ONNX_NAMESPACE::ModelProto CreateRecursiveFunctionExpansionModel(size_t branch_node_count, + size_t call_count) { + auto model = CreateFunctionExpansionModel(0, call_count); + model.set_doc_string(std::string(1024 * 1024, 'x')); + + auto* function = model.mutable_functions(0); + auto* condition = function->add_node(); + condition->set_op_type("Constant"); + condition->add_output("condition"); + auto* condition_value = condition->add_attribute(); + condition_value->set_name("value"); + condition_value->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + condition_value->mutable_t()->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_BOOL); + condition_value->mutable_t()->add_int32_data(1); + + auto* if_node = function->add_node(); + if_node->set_op_type("If"); + if_node->add_input("condition"); + if_node->add_output("function_output"); + + for (const auto* attribute_name : {"then_branch", "else_branch"}) { + auto* attribute = if_node->add_attribute(); + attribute->set_name(attribute_name); + attribute->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH); + auto* branch = attribute->mutable_g(); + branch->set_name(attribute_name); + + std::string previous_value = "function_input"; + for (size_t i = 0; i < branch_node_count; ++i) { + auto* node = branch->add_node(); + node->set_op_type("Identity"); + node->add_input(previous_value); + previous_value = "branch_" + std::to_string(i); + node->add_output(previous_value); + } + + auto* output = branch->add_output(); + output->set_name(previous_value); + auto* tensor_type = output->mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor_type->mutable_shape()->add_dim()->set_dim_value(1); + } + + return model; +} + +static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, + std::vector& log_messages) { std::string serialized_model; - ORT_RETURN_IF_NOT(CreateFunctionExpansionModel(body_node_count, call_count).SerializeToString(&serialized_model), + ORT_RETURN_IF_NOT(model.SerializeToString(&serialized_model), "Failed to serialize function expansion model."); SessionOptions session_options; - InferenceSession session{session_options, GetEnvironment()}; + auto capturing_sink = std::make_unique(); + auto* capturing_sink_ptr = capturing_sink.get(); + auto logging_manager = std::make_unique( + std::move(capturing_sink), logging::Severity::kWARNING, false, + logging::LoggingManager::InstanceType::Temporal); + std::unique_ptr environment; + ORT_RETURN_IF_ERROR(Environment::Create(std::move(logging_manager), environment)); + + InferenceSession session{session_options, *environment}; std::istringstream stream(serialized_model); ORT_RETURN_IF_ERROR(session.Load(stream)); - return session.Initialize(); + const auto status = session.Initialize(); + log_messages = capturing_sink_ptr->Messages(); + return status; } -TEST(FunctionTest, AotInliningLimitsFunctionExpansion) { - constexpr size_t body_node_count = 21; - ASSERT_STATUS_OK(InitializeFunctionExpansionModel(body_node_count, 1)); +TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { + auto model = CreateRecursiveFunctionExpansionModel(20, 14); + std::vector log_messages; + ASSERT_STATUS_OK(InitializeFunctionExpansionModel(std::move(model), log_messages)); + EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); + EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); +} - const auto status = InitializeFunctionExpansionModel(body_node_count, 21); - ASSERT_FALSE(status.IsOK()); - EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("AOT function inlining exceeds the")); +TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { + auto model = CreateFunctionExpansionModel(2, 20); + const std::string large_value_name(256 * 1024, 'x'); + auto* function = model.mutable_functions(0); + function->mutable_node(0)->set_output(0, large_value_name); + function->mutable_node(1)->set_input(0, large_value_name); + + std::vector log_messages; + ASSERT_STATUS_OK(InitializeFunctionExpansionModel(std::move(model), log_messages)); + EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("protobuf expansion limit"))); + EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("node expansion limit")))); } // A recursive/cyclic chain of model-local functions can be rejected by either layer: From c7fe54cd20bd840b76c407eb200794bb5c9f855e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 21 Sep 2026 19:18:32 +0000 Subject: [PATCH 03/11] Fix AOT expansion limit tests Include the complete Environment definition and avoid using an oversized graph value name when exercising the protobuf byte budget. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/framework/function_test.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index fa8fd787ac051..729c7597f139a 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -16,6 +16,7 @@ #include "core/graph/model.h" #include "core/graph/model_helpers.h" #include "core/providers/cpu/cpu_execution_provider.h" +#include "core/session/environment.h" #include "core/session/inference_session.h" #include "test/common/tensor_op_test_utils.h" @@ -253,10 +254,8 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { auto model = CreateFunctionExpansionModel(2, 20); - const std::string large_value_name(256 * 1024, 'x'); auto* function = model.mutable_functions(0); - function->mutable_node(0)->set_output(0, large_value_name); - function->mutable_node(1)->set_input(0, large_value_name); + function->mutable_node(0)->set_doc_string(std::string(256 * 1024, 'x')); std::vector log_messages; ASSERT_STATUS_OK(InitializeFunctionExpansionModel(std::move(model), log_messages)); From a8480ef192814b59d316e844e818983d2d584534 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 21 Sep 2026 21:54:09 +0000 Subject: [PATCH 04/11] Fix AOT protobuf budget test payload Use a Constant tensor payload to exercise protobuf expansion accounting without relying on an oversized node doc string that aborts graph processing across CI configurations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/framework/function_test.cc | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 7d748184f2910..5625d26092338 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -255,7 +255,16 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { auto model = CreateFunctionExpansionModel(2, 20); auto* function = model.mutable_functions(0); - function->mutable_node(0)->set_doc_string(std::string(256 * 1024, 'x')); + auto* payload_node = function->add_node(); + payload_node->set_op_type("Constant"); + payload_node->add_output("payload"); + auto* payload_attribute = payload_node->add_attribute(); + payload_attribute->set_name("value"); + payload_attribute->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* payload_tensor = payload_attribute->mutable_t(); + payload_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8); + payload_tensor->add_dims(256 * 1024); + payload_tensor->set_raw_data(std::string(256 * 1024, 'x')); std::vector log_messages; ASSERT_STATUS_OK(InitializeFunctionExpansionModel(std::move(model), log_messages)); From bc505642d540b493fc3c426f023095765f8d49e1 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 21 Sep 2026 22:08:44 +0000 Subject: [PATCH 05/11] Enforce AOT expansion limits at initialization Track functions retained because they exceed an AOT expansion budget and fail initialization after the AOT pass, before fallback partitioning can inline them without limits. Update both budget regressions to require the explicit failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/framework/graph_partitioner.cc | 18 ++++++++++++++++-- onnxruntime/test/framework/function_test.cc | 8 ++++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index af8d51389b578..3b3c14ad644d5 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1368,6 +1368,7 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide const logging::Logger& logger, const CheckLoadCancellationFn& check_load_cancellation_fn, InlinedHashSet& not_inlined, + InlinedHashSet& budget_limited_functions, size_t& inlined_count, size_t expansion_node_budget, size_t& expanded_node_count, @@ -1390,6 +1391,7 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide logger, check_load_cancellation_fn, not_inlined, + budget_limited_functions, inlined_count, expansion_node_budget, expanded_node_count, @@ -1457,13 +1459,15 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide 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(std::move(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(std::move(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; @@ -1862,6 +1866,7 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, size_t expanded_node_count = 0; size_t expanded_proto_bytes = 0; InlinedHashSet not_inlined; + InlinedHashSet budget_limited_functions; do { size_t inlined_count = 0; ORT_RETURN_IF_ERROR(InlineFunctionsAOTImpl(execution_providers, @@ -1871,6 +1876,7 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, logger, check_load_cancellation_fn, not_inlined, + budget_limited_functions, inlined_count, expansion_node_budget, expanded_node_count, @@ -1883,6 +1889,14 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, 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) diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 5625d26092338..2a0b96d3e74a2 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -247,7 +247,9 @@ static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { auto model = CreateRecursiveFunctionExpansionModel(20, 14); std::vector log_messages; - ASSERT_STATUS_OK(InitializeFunctionExpansionModel(std::move(model), log_messages)); + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("fallback inlining")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); } @@ -267,7 +269,9 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { payload_tensor->set_raw_data(std::string(256 * 1024, 'x')); std::vector log_messages; - ASSERT_STATUS_OK(InitializeFunctionExpansionModel(std::move(model), log_messages)); + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("fallback inlining")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("protobuf expansion limit"))); EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("node expansion limit")))); } From 2ca655765473c9879b6086abfd76563ced05b89a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 21 Sep 2026 23:40:17 +0000 Subject: [PATCH 06/11] Budget unclaimed calls of claimed functions Keep EP-retained function protos separate from functions rejected by the expansion budget, so unclaimed calls sharing an EP-claimed function still consume the AOT budget. Add selective-EP regression coverage for the mixed-call case. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/framework/graph_partitioner.cc | 2 +- onnxruntime/test/framework/function_test.cc | 43 ++++++++++++++++++- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 3b3c14ad644d5..361571af733ad 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1450,7 +1450,7 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide if (node != nullptr) { if (claimed_by_ep.count(node_index) == 0) { auto function_id = function_utils::GetFunctionIdentifier(node->Domain(), node->OpType(), node->Overload()); - if (not_inlined.count(function_id) != 0) { + if (budget_limited_functions.count(function_id) != 0) { continue; } diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 2a0b96d3e74a2..524313e6ba27a 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -16,6 +16,7 @@ #include "core/graph/model.h" #include "core/graph/model_helpers.h" #include "core/providers/cpu/cpu_execution_provider.h" +#include "core/providers/partitioning_utils.h" #include "core/session/environment.h" #include "core/session/inference_session.h" @@ -222,7 +223,8 @@ static ONNX_NAMESPACE::ModelProto CreateRecursiveFunctionExpansionModel(size_t b } static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, - std::vector& log_messages) { + std::vector& log_messages, + bool claim_first_function_call = false) { std::string serialized_model; ORT_RETURN_IF_NOT(model.SerializeToString(&serialized_model), "Failed to serialize function expansion model."); @@ -237,6 +239,36 @@ static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, ORT_RETURN_IF_ERROR(Environment::Create(std::move(logging_manager), environment)); InferenceSession session{session_options, *environment}; + if (claim_first_function_call) { + class FirstFunctionCallExecutionProvider final + : public internal_testing_ep::InternalTestingExecutionProvider { + public: + FirstFunctionCallExecutionProvider() + : InternalTestingExecutionProvider({}, {}, DataLayout::NCHW) {} + + std::vector> GetCapability( + const GraphViewer& graph_view, + const IKernelLookup&, + const GraphOptimizerRegistry&, + IResourceAccountant*) const override { + for (const auto node_index : graph_view.GetNodesInTopologicalOrder()) { + const auto* node = graph_view.GetNode(node_index); + if (node != nullptr && node->CanBeInlined()) { + return {utils::MakeComputeCapability( + graph_view, std::vector{node}, + [node_index]() { return "FirstFunctionCall_" + std::to_string(node_index); }, + Type(), false)}; + } + } + + return {}; + } + }; + + ORT_RETURN_IF_ERROR(session.RegisterExecutionProvider( + std::make_unique())); + } + std::istringstream stream(serialized_model); ORT_RETURN_IF_ERROR(session.Load(stream)); const auto status = session.Initialize(); @@ -254,6 +286,15 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); } +TEST(FunctionTest, AotInliningLimitsUnclaimedCallsSharingClaimedFunction) { + auto model = CreateRecursiveFunctionExpansionModel(20, 14); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages, true); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("fallback inlining")); + EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); +} + TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { auto model = CreateFunctionExpansionModel(2, 20); auto* function = model.mutable_functions(0); From 3184001d09d794a62bf5aff482568a95f9a96599 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 22 Sep 2026 00:13:11 +0000 Subject: [PATCH 07/11] Fix selective EP capability construction Build the capability vector explicitly so its unique_ptr element is moved instead of copied through an initializer_list. This restores compilation for full test targets across platforms. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/framework/function_test.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 524313e6ba27a..46a10c488e90c 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -254,10 +254,12 @@ static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, for (const auto node_index : graph_view.GetNodesInTopologicalOrder()) { const auto* node = graph_view.GetNode(node_index); if (node != nullptr && node->CanBeInlined()) { - return {utils::MakeComputeCapability( + std::vector> capabilities; + capabilities.push_back(utils::MakeComputeCapability( graph_view, std::vector{node}, [node_index]() { return "FirstFunctionCall_" + std::to_string(node_index); }, - Type(), false)}; + Type(), false)); + return capabilities; } } From e3a3241989259bdb7b19fb5943f27c98a2572774 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 22 Sep 2026 01:21:24 +0000 Subject: [PATCH 08/11] Fix selective function expansion regression Avoid colliding the retained local function with the standard Expand operator so optimizer passes do not interpret its one-input call as ONNX Expand. Increase the call count so unclaimed expansions exceed the shared node budget. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/framework/function_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 46a10c488e90c..23c630aa0b83d 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -146,7 +146,7 @@ static ONNX_NAMESPACE::ModelProto CreateFunctionExpansionModel(size_t body_node_ auto* function = model.add_functions(); function->set_domain("local"); - function->set_name("Expand"); + function->set_name("FunctionToExpand"); function->add_input("function_input"); function->add_output("function_output"); auto* function_opset = function->add_opset_import(); @@ -166,7 +166,7 @@ static ONNX_NAMESPACE::ModelProto CreateFunctionExpansionModel(size_t body_node_ for (size_t i = 0; i < call_count; ++i) { auto* node = graph->add_node(); node->set_domain("local"); - node->set_op_type("Expand"); + node->set_op_type("FunctionToExpand"); node->add_input(previous_value); previous_value = "call_" + std::to_string(i); node->add_output(previous_value); @@ -289,7 +289,7 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { } TEST(FunctionTest, AotInliningLimitsUnclaimedCallsSharingClaimedFunction) { - auto model = CreateRecursiveFunctionExpansionModel(20, 14); + auto model = CreateRecursiveFunctionExpansionModel(20, 15); std::vector log_messages; const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages, true); EXPECT_FALSE(status.IsOK()); From 7ad53b5fac05c2ebb466bd0eb981a10b9a69c27b Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 23 Sep 2026 01:29:18 +0000 Subject: [PATCH 09/11] Measure specialized AOT function payload Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/framework/graph_partitioner.cc | 10 +++++++++- onnxruntime/test/framework/function_test.cc | 10 ++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 361571af733ad..0a37f3f20ebc5 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1347,7 +1347,15 @@ static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& std::string accounting_prefix = "_inlfunc_" + node.OpType(); accounting_prefix.append(32, '_'); function_utils::Specialize(function_proto, node, accounting_prefix); - cost = {CountNodesIncludingSubgraphs(function_proto), function_proto.ByteSizeLong()}; + SafeInt node_count = function_proto.node_size(); + SafeInt 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(); } diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 23c630aa0b83d..2e4fc810e5182 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -319,6 +319,16 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("node expansion limit")))); } +TEST(FunctionTest, AotInliningIgnoresFunctionMetadataForProtoBytes) { + auto model = CreateFunctionExpansionModel(1, 11); + model.mutable_functions(0)->set_doc_string(std::string(1024 * 1024, 'x')); + + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); + EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); +} + // A recursive/cyclic chain of model-local functions can be rejected by either layer: // ONNX 1.22+ detects the cycle in its own model checker ("Cycle detected in model-local // function references"), which runs before ORT's equivalent check ("must not be recursive"). From e0d4ebfd26f958401ab984b9ebbcfed77af14b5d Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 23 Sep 2026 18:45:24 +0000 Subject: [PATCH 10/11] Preserve models under function expansion limits Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/workflows/android.yml | 4 +- .../onnxruntime_session_options_config_keys.h | 5 + .../core/framework/graph_partitioner.cc | 224 +++++++++++------- .../core/framework/graph_partitioner.h | 9 + onnxruntime/core/graph/model.h | 2 - onnxruntime/core/session/inference_session.cc | 1 + onnxruntime/test/framework/function_test.cc | 66 +++++- 7 files changed, 211 insertions(+), 100 deletions(-) diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml index 64f5de7e78dcc..b3d6b895a97fb 100644 --- a/.github/workflows/android.yml +++ b/.github/workflows/android.yml @@ -79,8 +79,8 @@ jobs: run: | set -e -x BINARY_SIZE_THRESHOLD_ARGS="" - echo "Binary size threshold in bytes: 1595392" - BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1595392" + echo "Binary size threshold in bytes: 1596416" + BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1596416" # Ensure ANDROID_NDK_HOME is available and get its real path if [ -z "$ANDROID_NDK_HOME" ]; then diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 977fd1c54571b..bc7237ee570ff 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -99,6 +99,11 @@ static const char* const kOrtSessionOptionsEnableCastChainElimination = "optimiz // Its default value is "0". static const char* const kOrtSessionOptionsDisableAheadOfTimeFunctionInlining = "session.disable_aot_function_inlining"; +// Limits cumulative model-local function expansion across AOT and fallback inlining. +// Values must be positive decimal integers. Defaults are 1,000,000 nodes and 1 GiB of serialized node payload. +static const char* const kOrtSessionOptionsFunctionExpansionNodeLimit = "session.function_expansion_node_limit"; +static const char* const kOrtSessionOptionsFunctionExpansionByteLimit = "session.function_expansion_byte_limit"; + #ifdef ENABLE_TRAINING // Specifies a path of the file containing a list of memory optimization configurations. // The value should be a string indicating the file path of the config file. diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 0a37f3f20ebc5..821c9c1bc4e0b 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1207,13 +1207,41 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr, return Status::OK(); } +struct FunctionExpansionCost { + size_t node_count; + size_t proto_bytes; +}; + +enum class FunctionExpansionLimit { + kNone, + kNodes, + kProtoBytes, +}; + +static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& cost); + +static FunctionExpansionLimit TryChargeFunctionExpansion(const FunctionExpansionCost& cost, + size_t node_limit, + size_t& expanded_node_count, + size_t byte_limit, + size_t& expanded_proto_bytes); + // expand any nodes that have an ONNX function definition but no matching ORT kernel -static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* layering_index) { +static Status InlineNodes(Graph& graph, + bool& modified_graph, + LayeringIndex* layering_index, + const logging::Logger& logger, + size_t expansion_node_limit, + size_t& expanded_node_count, + size_t expansion_byte_limit, + size_t& expanded_proto_bytes) { // recurse into nested graphs first so we process from bottom up for (auto& node : graph.Nodes()) { for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) { Graph* subgraph = entry.second; - ORT_RETURN_IF_ERROR(InlineNodes(*subgraph, modified_graph, layering_index)); + ORT_RETURN_IF_ERROR(InlineNodes(*subgraph, modified_graph, layering_index, logger, + expansion_node_limit, expanded_node_count, + expansion_byte_limit, expanded_proto_bytes)); } } @@ -1233,6 +1261,23 @@ static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* lay InlinedVector new_node_indices; for (auto* node : nodes_to_inline) { + FunctionExpansionCost expansion_cost{}; + ORT_RETURN_IF_ERROR(GetFunctionExpansionCost(*node, expansion_cost)); + const auto limit_exceeded = TryChargeFunctionExpansion(expansion_cost, + expansion_node_limit, + expanded_node_count, + expansion_byte_limit, + expanded_proto_bytes); + if (limit_exceeded != FunctionExpansionLimit::kNone) { + const auto function_id = + function_utils::GetFunctionIdentifier(node->Domain(), node->OpType(), node->Overload()); + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Function inlining exceeded the configured cumulative ", + limit_exceeded == FunctionExpansionLimit::kNodes ? "node" : "protobuf", + " expansion limit while expanding '", function_id, "'."); + } + // Check for an effective layering assignment: either from an explicit annotation // on the node, or from an inherited assignment via the LayeringIndex (e.g., a function // call node inside an annotated If/Loop subgraph that inherited its parent's rule). @@ -1277,7 +1322,8 @@ static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* lay return Status::OK(); } -constexpr size_t kAotFunctionExpansionRatio = 10; +constexpr size_t kDefaultFunctionExpansionNodeLimit = 1'000'000; +constexpr size_t kDefaultFunctionExpansionByteLimit = 1024ULL * 1024ULL * 1024ULL; static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::GraphProto& graph); @@ -1304,40 +1350,14 @@ static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::GraphProto& gra return node_count; } -static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::FunctionProto& function) { - SafeInt 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 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()}; + SafeInt proto_bytes = 0; + for (const auto& function_node : graph_proto.node()) { + proto_bytes += function_node.ByteSizeLong(); + } + cost = {CountNodesIncludingSubgraphs(graph_proto), proto_bytes}; return Status::OK(); } @@ -1359,14 +1379,49 @@ static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& return Status::OK(); } -static size_t CountModelNodes(const Model& model) { - SafeInt 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_); +static FunctionExpansionLimit TryChargeFunctionExpansion(const FunctionExpansionCost& cost, + size_t node_limit, + size_t& expanded_node_count, + size_t byte_limit, + size_t& expanded_proto_bytes) { + if (cost.node_count > node_limit - expanded_node_count) { + return FunctionExpansionLimit::kNodes; + } + if (cost.proto_bytes > byte_limit - expanded_proto_bytes) { + return FunctionExpansionLimit::kProtoBytes; } - return node_count; + expanded_node_count += cost.node_count; + expanded_proto_bytes += cost.proto_bytes; + return FunctionExpansionLimit::kNone; +} + +static Status InitializeFunctionExpansionLimits(const ConfigOptions& config_options, + bool& initialized, + size_t& node_limit, + size_t& byte_limit) { + if (initialized) { + return Status::OK(); + } + + const auto parse_limit = [&config_options](const char* config_key, + size_t default_value, + size_t& value) -> Status { + const auto config_value = config_options.GetConfigOrDefault(config_key, std::to_string(default_value)); + if (!TryParseStringWithClassicLocale(config_value, value) || value == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid positive integer value '", config_value, + "' for session configuration '", config_key, "'."); + } + return Status::OK(); + }; + + ORT_RETURN_IF_ERROR(parse_limit(kOrtSessionOptionsFunctionExpansionNodeLimit, + kDefaultFunctionExpansionNodeLimit, node_limit)); + ORT_RETURN_IF_ERROR(parse_limit(kOrtSessionOptionsFunctionExpansionByteLimit, + kDefaultFunctionExpansionByteLimit, byte_limit)); + initialized = true; + return Status::OK(); } static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_providers, @@ -1376,11 +1431,10 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide const logging::Logger& logger, const CheckLoadCancellationFn& check_load_cancellation_fn, InlinedHashSet& not_inlined, - InlinedHashSet& budget_limited_functions, size_t& inlined_count, - size_t expansion_node_budget, + size_t expansion_node_limit, size_t& expanded_node_count, - size_t expansion_byte_budget, + size_t expansion_byte_limit, 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 @@ -1399,11 +1453,10 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide logger, check_load_cancellation_fn, not_inlined, - budget_limited_functions, inlined_count, - expansion_node_budget, + expansion_node_limit, expanded_node_count, - expansion_byte_budget, + expansion_byte_limit, expanded_proto_bytes)); } } @@ -1458,28 +1511,22 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide 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 << "'."; + const auto limit_exceeded = TryChargeFunctionExpansion(expansion_cost, + expansion_node_limit, + expanded_node_count, + expansion_byte_limit, + expanded_proto_bytes); + if (limit_exceeded != FunctionExpansionLimit::kNone) { + LOGS(logger, WARNING) << "AOT function inlining reached the cumulative " + << (limit_exceeded == FunctionExpansionLimit::kNodes ? "node" : "protobuf") + << " expansion limit. " + << "Retaining function call '" << function_id + << "' for execution-provider partitioning."; 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 { @@ -1648,7 +1695,11 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params, KernelRegistryManager& kernel_registry_manager, const std::optional& acc_map, const GraphOptimizerRegistry& graph_optimizer_registry, - const logging::Logger& logger, bool disable_model_compile) { // Added arg + const logging::Logger& logger, bool disable_model_compile, + size_t expansion_node_limit, + size_t& expanded_node_count, + size_t expansion_byte_limit, + size_t& expanded_proto_bytes) { // Added arg bool modified_graph = false; auto& graph = partition_params.graph.get(); @@ -1695,7 +1746,9 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params, // expand any nodes that have an ONNX function definition but no matching ORT kernel. modified_graph = false; - ORT_RETURN_IF_ERROR(InlineNodes(graph, modified_graph, partition_params.layering_index)); + ORT_RETURN_IF_ERROR(InlineNodes(graph, modified_graph, partition_params.layering_index, logger, + expansion_node_limit, expanded_node_count, + expansion_byte_limit, expanded_proto_bytes)); // Resolve and rerun graph partitioning and inlining if there was a change if (modified_graph) { @@ -1855,6 +1908,7 @@ static Status PartitionOrtFormatModel(const PartitionParams& partition_params, Status GraphPartitioner::InlineFunctionsAOT(Model& model, const ExecutionProviders& execution_providers, const KernelRegistryManager& kernel_registry_manager, + const ConfigOptions& config_options, const logging::Logger& logger) const { const auto local_functions_num = model.GetModelLocalFunctionTemplates().size(); const bool is_there_local_functions = local_functions_num > 0; @@ -1865,16 +1919,14 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, } auto check_load_cancellation_fn = [this]() -> bool { return IsLoadCancellationFlagSet(); }; + ORT_RETURN_IF_ERROR(InitializeFunctionExpansionLimits( + config_options, + function_expansion_limits_initialized_, + function_expansion_node_limit_, + function_expansion_byte_limit_)); auto& graph = model.MainGraph(); - const size_t expansion_node_budget = - static_cast(SafeInt(CountModelNodes(model)) * kAotFunctionExpansionRatio); - const size_t expansion_byte_budget = - static_cast(SafeInt(model.ModelProtoByteSize()) * kAotFunctionExpansionRatio); - size_t expanded_node_count = 0; - size_t expanded_proto_bytes = 0; InlinedHashSet not_inlined; - InlinedHashSet budget_limited_functions; do { size_t inlined_count = 0; ORT_RETURN_IF_ERROR(InlineFunctionsAOTImpl(execution_providers, @@ -1884,12 +1936,11 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, logger, check_load_cancellation_fn, not_inlined, - budget_limited_functions, inlined_count, - expansion_node_budget, - expanded_node_count, - expansion_byte_budget, - expanded_proto_bytes)); + function_expansion_node_limit_, + expanded_function_node_count_, + function_expansion_byte_limit_, + expanded_function_proto_bytes_)); if (inlined_count == 0) { break; @@ -1897,14 +1948,6 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, 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) @@ -1993,11 +2036,20 @@ Status GraphPartitioner::Partition(Graph& graph, FuncManager& func_mgr, // The map is empty if not created if not enabled std::optional ep_acc_map; ORT_RETURN_IF_ERROR(CreateAccountants(config_options, graph.ModelPath(), ep_acc_map)); + ORT_RETURN_IF_ERROR(InitializeFunctionExpansionLimits( + config_options, + function_expansion_limits_initialized_, + function_expansion_node_limit_, + function_expansion_byte_limit_)); bool disable_model_compile = config_options.GetConfigOrDefault(kOrtSessionOptionsDisableModelCompile, "0") == "1"; ORT_RETURN_IF_ERROR(PartitionOnnxFormatModel(partition_params, mode, providers_, kernel_registry_mgr_, ep_acc_map, *graph_optimizer_registry_, logger, - disable_model_compile)); // Pass param + disable_model_compile, + function_expansion_node_limit_, + expanded_function_node_count_, + function_expansion_byte_limit_, + expanded_function_proto_bytes_)); // Pass param if (ep_acc_map.has_value()) { for (const auto& [ep_type, accountant] : *ep_acc_map) { diff --git a/onnxruntime/core/framework/graph_partitioner.h b/onnxruntime/core/framework/graph_partitioner.h index 6cf547c31107a..ad91b23282cea 100644 --- a/onnxruntime/core/framework/graph_partitioner.h +++ b/onnxruntime/core/framework/graph_partitioner.h @@ -103,6 +103,7 @@ class GraphPartitioner { Status InlineFunctionsAOT(Model& model, const ExecutionProviders& execution_providers, const KernelRegistryManager& kernel_registry_manager, + const ConfigOptions& config_options, const logging::Logger& logger) const; #endif @@ -114,6 +115,14 @@ class GraphPartitioner { std::unique_ptr graph_optimizer_registry_; CheckLoadCancellationFn check_load_cancellation_fn_; OnPartitionAssignmentFunction on_partition_assignment_fn_; +#ifndef ORT_MINIMAL_BUILD + // Shared by AOT and fallback inlining so neither path can bypass the cumulative limit. + mutable bool function_expansion_limits_initialized_ = false; + mutable size_t function_expansion_node_limit_ = 0; + mutable size_t function_expansion_byte_limit_ = 0; + mutable size_t expanded_function_node_count_ = 0; + mutable size_t expanded_function_proto_bytes_ = 0; +#endif }; } // namespace onnxruntime diff --git a/onnxruntime/core/graph/model.h b/onnxruntime/core/graph/model.h index eead39f1c78b0..a6ce2a3a0659c 100644 --- a/onnxruntime/core/graph/model.h +++ b/onnxruntime/core/graph/model.h @@ -199,8 +199,6 @@ 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; diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 16202678e0153..6ced42d6fd229 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1625,6 +1625,7 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool ORT_RETURN_IF_ERROR_SESSIONID_(partitioner.InlineFunctionsAOT(*model_, execution_providers_, kernel_registry_manager_, + session_options_.config_options, *session_logger_)); } diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 2e4fc810e5182..86d3b376b729e 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -4,6 +4,7 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include #include "core/graph/onnx_protobuf.h" @@ -19,6 +20,7 @@ #include "core/providers/partitioning_utils.h" #include "core/session/environment.h" #include "core/session/inference_session.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/tensor_op_test_utils.h" #include "test/capturing_sink.h" @@ -222,14 +224,37 @@ static ONNX_NAMESPACE::ModelProto CreateRecursiveFunctionExpansionModel(size_t b return model; } -static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, - std::vector& log_messages, - bool claim_first_function_call = false) { +struct FunctionExpansionTestOptions { + bool claim_first_function_call = false; + bool disable_aot_inlining = false; + std::optional node_limit; + std::optional byte_limit; +}; + +static Status InitializeFunctionExpansionModel( + ONNX_NAMESPACE::ModelProto model, + std::vector& log_messages, + const FunctionExpansionTestOptions& options = {}) { std::string serialized_model; ORT_RETURN_IF_NOT(model.SerializeToString(&serialized_model), "Failed to serialize function expansion model."); SessionOptions session_options; + if (options.disable_aot_inlining) { + ORT_RETURN_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsDisableAheadOfTimeFunctionInlining, "1")); + } + if (options.node_limit.has_value()) { + ORT_RETURN_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsFunctionExpansionNodeLimit, + std::to_string(*options.node_limit).c_str())); + } + if (options.byte_limit.has_value()) { + ORT_RETURN_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsFunctionExpansionByteLimit, + std::to_string(*options.byte_limit).c_str())); + } + auto capturing_sink = std::make_unique(); auto* capturing_sink_ptr = capturing_sink.get(); auto logging_manager = std::make_unique( @@ -239,7 +264,7 @@ static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, ORT_RETURN_IF_ERROR(Environment::Create(std::move(logging_manager), environment)); InferenceSession session{session_options, *environment}; - if (claim_first_function_call) { + if (options.claim_first_function_call) { class FirstFunctionCallExecutionProvider final : public internal_testing_ep::InternalTestingExecutionProvider { public: @@ -281,9 +306,10 @@ static Status InitializeFunctionExpansionModel(ONNX_NAMESPACE::ModelProto model, TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { auto model = CreateRecursiveFunctionExpansionModel(20, 14); std::vector log_messages; - const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, {.node_limit = 500}); EXPECT_FALSE(status.IsOK()); - EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("fallback inlining")); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); } @@ -291,9 +317,11 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { TEST(FunctionTest, AotInliningLimitsUnclaimedCallsSharingClaimedFunction) { auto model = CreateRecursiveFunctionExpansionModel(20, 15); std::vector log_messages; - const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages, true); + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, + {.claim_first_function_call = true, .node_limit = 500}); EXPECT_FALSE(status.IsOK()); - EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("fallback inlining")); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); } @@ -312,13 +340,31 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { payload_tensor->set_raw_data(std::string(256 * 1024, 'x')); std::vector log_messages; - const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, {.byte_limit = 1024 * 1024}); EXPECT_FALSE(status.IsOK()); - EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("fallback inlining")); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("protobuf expansion limit")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("protobuf expansion limit"))); EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("node expansion limit")))); } +TEST(FunctionTest, FallbackInliningEnforcesExpansionLimitWhenAotIsDisabled) { + auto model = CreateRecursiveFunctionExpansionModel(20, 14); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, + {.disable_aot_inlining = true, .node_limit = 500}); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); +} + +TEST(FunctionTest, DefaultExpansionLimitPreservesOrdinaryLargeFunctions) { + auto model = CreateFunctionExpansionModel(200, 24); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); +} + TEST(FunctionTest, AotInliningIgnoresFunctionMetadataForProtoBytes) { auto model = CreateFunctionExpansionModel(1, 11); model.mutable_functions(0)->set_doc_string(std::string(1024 * 1024, 'x')); From fd540e597711fdaaef0e8d71f0a69ef818c2e013 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 23 Sep 2026 19:30:49 +0000 Subject: [PATCH 11/11] Initialize function expansion test options Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/framework/function_test.cc | 22 +++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index 86d3b376b729e..ad6ccfbe38fc2 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -307,7 +307,11 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { auto model = CreateRecursiveFunctionExpansionModel(20, 14); std::vector log_messages; const auto status = InitializeFunctionExpansionModel( - std::move(model), log_messages, {.node_limit = 500}); + std::move(model), log_messages, + {.claim_first_function_call = false, + .disable_aot_inlining = false, + .node_limit = 500, + .byte_limit = std::nullopt}); EXPECT_FALSE(status.IsOK()); EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); @@ -319,7 +323,10 @@ TEST(FunctionTest, AotInliningLimitsUnclaimedCallsSharingClaimedFunction) { std::vector log_messages; const auto status = InitializeFunctionExpansionModel( std::move(model), log_messages, - {.claim_first_function_call = true, .node_limit = 500}); + {.claim_first_function_call = true, + .disable_aot_inlining = false, + .node_limit = 500, + .byte_limit = std::nullopt}); EXPECT_FALSE(status.IsOK()); EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); @@ -341,7 +348,11 @@ TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { std::vector log_messages; const auto status = InitializeFunctionExpansionModel( - std::move(model), log_messages, {.byte_limit = 1024 * 1024}); + std::move(model), log_messages, + {.claim_first_function_call = false, + .disable_aot_inlining = false, + .node_limit = std::nullopt, + .byte_limit = 1024 * 1024}); EXPECT_FALSE(status.IsOK()); EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("protobuf expansion limit")); EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("protobuf expansion limit"))); @@ -353,7 +364,10 @@ TEST(FunctionTest, FallbackInliningEnforcesExpansionLimitWhenAotIsDisabled) { std::vector log_messages; const auto status = InitializeFunctionExpansionModel( std::move(model), log_messages, - {.disable_aot_inlining = true, .node_limit = 500}); + {.claim_first_function_call = false, + .disable_aot_inlining = true, + .node_limit = 500, + .byte_limit = std::nullopt}); EXPECT_FALSE(status.IsOK()); EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); }