Limit AOT function expansion - #32662
Open
Akshay Sonawane (apsonawane) wants to merge 1 commit into
Open
Akshay Sonawane (apsonawane) wants to merge 1 commit into
Akshay Sonawane (apsonawane) wants to merge 1 commit into
Conversation
Akshay Sonawane (apsonawane)
enabled auto-merge (squash)
September 16, 2026 23:23
Copilot started reviewing on behalf of
Akshay Sonawane (apsonawane)
September 16, 2026 23:24
View session
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Budget calculation duplicates large models in memory, and the test does not independently validate both limits.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds safeguards against excessive AOT function-inlining growth.
Changes:
- Adds recursive node and protobuf-size accounting.
- Enforces node and byte expansion budgets.
- Adds an expansion-limit unit test.
File summaries
| File | Description |
|---|---|
onnxruntime/core/framework/graph_partitioner.cc |
Calculates and enforces AOT expansion budgets. |
onnxruntime/test/framework/function_test.cc |
Adds function-expansion test coverage. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 2
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| auto check_load_cancellation_fn = [this]() -> bool { return IsLoadCancellationFlagSet(); }; | ||
|
|
||
| auto& graph = model.MainGraph(); | ||
| const auto model_proto = model.ToProto(); |
Comment on lines
+192
to
+194
| const auto status = InitializeFunctionExpansionModel(body_node_count, 21); | ||
| ASSERT_FALSE(status.IsOK()); | ||
| EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("AOT function inlining exceeds the")); |
| 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, |
There was a problem hiding this comment.
Is it ok to return FAIL here? If the function exceeds the budget, that does not mean none of the following functions should be rejected.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a node and byte budget for Ahead-Of-Time (AOT) function inlining in the ONNX Runtime graph partitioner, limiting the expansion of functions during inlining to prevent excessive growth in model size and complexity. It adds utility methods to count nodes (including within subgraphs and functions), calculates expansion budgets based on the original model size, and enforces these limits during inlining. Additionally, a new unit test verifies the enforcement of these limits.
AOT Function Inlining Limits
CountNodesIncludingSubgraphs,CountModelNodes, etc.) to recursively count nodes and protobuf size in graphs, attributes, and functions to track potential expansion during inlining.kAotFunctionExpansionRatio) and used it to compute node and byte budgets for inlining, based on the original model's size.InlineFunctionsAOTImplto check and enforce node and byte budgets before inlining each function, returning an error if the limits would be exceeded. [1] [2] [3]Testing
AotInliningLimitsFunctionExpansion) infunction_test.ccto verify that the inlining limits are enforced and that exceeding them produces the expected error.