Skip to content

Add experimental support for externalizing graphs to a separate artifact - #77

Draft
jakesabathia2 wants to merge 3 commits into
apple:mainfrom
jakesabathia2:dev/henry/externalize-graphs-to-artifact
Draft

Add experimental support for externalizing graphs to a separate artifact#77
jakesabathia2 wants to merge 3 commits into
apple:mainfrom
jakesabathia2:dev/henry/externalize-graphs-to-artifact

Conversation

@jakesabathia2

Copy link
Copy Markdown
Contributor

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.

`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
jakesabathia2 marked this pull request as draft September 1, 2026 23:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant