Add an option to not inline a function when building the graph#2851
Open
justinchuby wants to merge 4 commits intomainfrom
Open
Add an option to not inline a function when building the graph#2851justinchuby wants to merge 4 commits intomainfrom
justinchuby wants to merge 4 commits intomainfrom
Conversation
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds support in the internal GraphBuilder to optionally not inline an ONNXScript function call, emitting a single function-call node and tracking the referenced function definitions for later export/attachment.
Changes:
- Add
Op.domainconvenience property to expose an op’s opset domain. - Extend
GraphBuilder.call/OpBuilder.callwith_inlineflag; when_inline=False, emit a single call node and register the callee inGraphBuilder.functions. - Add unit tests covering
_inline=Falsebehavior (single node emission, output renaming, prefix handling, function registration).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
| onnxscript/_internal/values.py | Adds Op.domain property forwarding to the bound opset domain. |
| onnxscript/_internal/builder.py | Introduces function registry and _inline option to either inline or emit a function-call node. |
| onnxscript/_internal/builder_test.py | Adds tests verifying non-inlined call behavior and registration. |
onnxscript/_internal/builder.py
Outdated
Comment on lines
+594
to
+606
| node = ir.node( | ||
| op_type=function.name, | ||
| inputs=args, | ||
| attributes=kwargs or None, | ||
| outputs=[ | ||
| ir.Value(name=output_renaming[output.name]) for output in graph.outputs | ||
| ], | ||
| domain=function.domain, | ||
| name=self._qualify_node_name(function.name), | ||
| ) | ||
| outputs = node.outputs | ||
| self.add_node(node) | ||
| self._functions[function.identifier] = function |
Comment on lines
+594
to
+603
| node = ir.node( | ||
| op_type=function.name, | ||
| inputs=args, | ||
| attributes=kwargs or None, | ||
| outputs=[ | ||
| ir.Value(name=output_renaming[output.name]) for output in graph.outputs | ||
| ], | ||
| domain=function.domain, | ||
| name=self._qualify_node_name(function.name), | ||
| ) |
Comment on lines
+598
to
+604
| outputs=[ | ||
| ir.Value(name=output_renaming[output.name]) for output in graph.outputs | ||
| ], | ||
| domain=function.domain, | ||
| name=self._qualify_node_name(function.name), | ||
| ) | ||
| outputs = node.outputs |
Comment on lines
+734
to
+735
| called as a separate node. When False, the function will be added | ||
| to the ``.functions`` dictionary. Defaults to True. |
|
|
||
| # The function should be registered | ||
| self.assertEqual(len(op.builder.functions), 1) | ||
| registered = next(iter(op.builder.functions.values())) |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2851 +/- ##
==========================================
+ Coverage 71.86% 71.90% +0.03%
==========================================
Files 239 239
Lines 29139 29241 +102
Branches 2875 2876 +1
==========================================
+ Hits 20942 21026 +84
- Misses 7219 7237 +18
Partials 978 978 ☔ View full report in Codecov by Sentry. |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
gramalingam
reviewed
Mar 14, 2026
| *args, | ||
| _outputs: Sequence[str] | None = None, | ||
| _prefix: str = "", | ||
| _inline: bool = True, |
Collaborator
There was a problem hiding this comment.
Should we enforce that _prefix is "" if _inline is False?
Collaborator
|
Not sure if it is better to have a separate function or an option (as in this PR). For now, this seems fine. |
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 new mechanism for function calls in the ONNXScript builder, allowing users to choose between inlining a function's body into the graph or creating a single function call node. It also adds tracking and registration of called functions, exposes the registered functions via a property, and provides comprehensive tests to ensure correctness of the new feature.
Function call mechanism enhancements:
_inlineoption to thecallmethod inGraphBuilderandOpBuilder, allowing users to either inline the function's body (default) or create a single function call node and register the function. When_inline=False, the function is added to the builder'sfunctionsdictionary and a single node is created in the graph; when_inline=True, the function is inlined and not registered. (onnxscript/_internal/builder.py[1] [2] [3] [4]functionsproperty to bothGraphBuilderandOpBuilder, exposing the dictionary of registered functions for inspection and downstream use. (onnxscript/_internal/builder.py[1] [2] [3]Testing and validation:
_inlineoption, verifying function registration, node creation, output renaming, prefix handling, and the difference in node count between inlining and non-inlining. (onnxscript/_internal/builder_test.pyonnxscript/_internal/builder_test.pyR851-R1001)API improvements:
domainproperty to theOnnxFunctionclass for easier access to the function's domain. (onnxscript/_internal/values.pyonnxscript/_internal/values.pyR225-R228)