Add experimental support for externalizing graphs to a separate artifact - #77
Draft
jakesabathia2 wants to merge 3 commits into
Draft
Add experimental support for externalizing graphs to a separate artifact#77jakesabathia2 wants to merge 3 commits into
jakesabathia2 wants to merge 3 commits into
Conversation
`ExternalizeSpec` already emits a matched submodule as its own `coreai.graph` invoked via `coreai.invoke`. This adds the pieces needed to *ship those graphs separately* from the program that calls them, so their weights can be replaced without re-exporting the base: * `ExternalizeSpec._graph_externalize` marks the emitted graph with the `externalize` attribute. Set after construction via the op's own setter, so no extra flag has to be threaded through `_get_graph_op`. Rejected in combination with `composite_op_name`, since composite graphs are `private` and the op definition forbids `externalize` with `private`. * `ExternalizeSpec._namespace` takes a dotted path and places the graph in nested symbol tables, calling it through the qualified reference `@group_a::@stage_one::@name`. Nesting uses `builtin.module` because the `coreai` dialect has no module op; a builtin module is a symbol table, which is what makes the qualified reference resolve. Existing levels are reused. * `_externalize_graphs(program)` returns a new `AIProgram` holding only the marked graphs. It deep-copies first: `ISOLATE_EXTERNALIZED_GRAPHS` mutates in place and `AIProgram` holds its module by reference, so calling the pass directly would consume the base program the caller still needs to save. Callers are therefore order-independent and can inspect both artifacts. * Graph names become deterministic when `_graph_externalize` is set, derived from the module path plus the call site's input type signature. A separately-exported base program and side artifact must agree on symbol names, which the default `uuid4` suffix cannot guarantee. Left as-is otherwise: graphs that stay in one module only need uniqueness, and renaming them would churn existing IR expectations for no benefit. The whole surface is underscore-prefixed deliberately -- experimental, with no backwards-compatibility guarantee, and removable once something first-class exists. Tests cover the multi-input submodule boundary (a two-tensor side branch must not be flattened to one argument), the actionable failure for an untyped `*args` forward, one graph per call site, marking, extraction leaving the source intact, namespace nesting, name determinism, and the underlying pass semantics pinned directly so a compiler regression stays distinguishable from a regression here.
jakesabathia2
marked this pull request as draft
September 1, 2026 23:26
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.
ExternalizeSpecalready emits a matched submodule as its owncoreai.graphinvoked viacoreai.invoke. This adds the pieces needed to ship those graphs separately from the program that calls them, so their weights can be replaced without re-exporting the base:ExternalizeSpec._graph_externalizemarks the emitted graph with theexternalizeattribute. Set after construction via the op's own setter, so no extra flag has to be threaded through_get_graph_op. Rejected in combination withcomposite_op_name, since composite graphs areprivateand the op definition forbidsexternalizewithprivate.ExternalizeSpec._namespacetakes a dotted path and places the graph in nested symbol tables, calling it through the qualified reference@group_a::@stage_one::@name. Nesting usesbuiltin.modulebecause thecoreaidialect has no module op; a builtin module is a symbol table, which is what makes the qualified reference resolve. Existing levels are reused._externalize_graphs(program)returns a newAIProgramholding only the marked graphs. It deep-copies first:ISOLATE_EXTERNALIZED_GRAPHSmutates in place andAIProgramholds its module by reference, so calling the pass directly would consume the base program the caller still needs to save. Callers are therefore order-independent and can inspect both artifacts._graph_externalizeis set, derived from the module path plus the call site's input type signature. A separately-exported base program and side artifact must agree on symbol names, which the defaultuuid4suffix cannot guarantee. Left as-is otherwise: graphs that stay in one module only need uniqueness, and renaming them would churn existing IR expectations for no benefit.The whole surface is underscore-prefixed deliberately -- experimental, with no backwards-compatibility guarantee, and removable once something first-class exists.
Tests cover the multi-input submodule boundary (a two-tensor side branch must not be flattened to one argument), the actionable failure for an untyped
*argsforward, one graph per call site, marking, extraction leaving the source intact, namespace nesting, name determinism, and the underlying pass semantics pinned directly so a compiler regression stays distinguishable from a regression here.