Skip to content

Unify BuilderBase and GraphBuilder: shared call_op with feature flags#2910

Open
gramalingam wants to merge 2 commits intomainfrom
rama/unify-builders
Open

Unify BuilderBase and GraphBuilder: shared call_op with feature flags#2910
gramalingam wants to merge 2 commits intomainfrom
rama/unify-builders

Conversation

@gramalingam
Copy link
Copy Markdown
Collaborator

Summary

Extends the BuilderBase / TapeBuilder infrastructure (from PR #2905) so that GraphBuilder inherits from BuilderBase, sharing a single call_op() implementation gated by feature flags.

Changes

Class rename and API merge

  • Renamed OpBuilderBaseBuilderBase (no backward-compat alias needed — not yet released)
  • Merged the separate op() and _make_node() methods into a single op(op_type, *args, **kwargs) ergonomic API
  • Promoted _create_op to BuilderBase.call_op() — the single core node-creation method

Feature flags (BuilderFeature enum)

  • SCHEMA_PARTITION — partition positional args into inputs vs attributes using op schema
  • CAST_INPUTS — promote scalars to tensors, match type variables
  • CAST_ATTRIBUTES — cast attributes using schema info (stub for future)
  • INFER_SHAPES — run shape inference after node creation
  • CONSTANT_PROPAGATION — run constant propagation after node creation
  • FULL = all of the above (used by GraphBuilder)
  • NONE = no extra processing (used by TapeBuilder for rewriter/optimizer/version-converter)

GraphBuilder inherits from BuilderBase

  • GraphBuilder(BuilderBase) with features=BuilderFeature.FULL
  • call_op() is inherited directly — no override needed
  • Default hook implementations (_get_schema, _partition_inputs_attributes, _cast_inputs, _cast_attributes) live in BuilderBase
  • GraphBuilder only overrides hooks needing graph-specific state:
    • _input_to_ir_value — constant promotion with caching
    • _generate_node_name / _adapt_outputs — qualified naming
    • _post_create_node — constant propagation + shape inference
    • _annotate_node — scope metadata

Exports

  • BuilderFeature exported from onnxscript package

Testing

  • All 559 rewriter/optimizer/version-converter tests pass
  • All 90 GraphBuilder tests pass
  • All 15 TapeBuilder/context tests pass

@gramalingam gramalingam changed the title Unify BuilderBase and GraphBuilder: shared call_op with feature flags [DRAFT] Unify BuilderBase and GraphBuilder: shared call_op with feature flags May 7, 2026
Comment thread onnxscript/_internal/tape_builder.py Fixed
Comment thread onnxscript/_internal/tape_builder.py Fixed
@codecov
Copy link
Copy Markdown

codecov Bot commented May 7, 2026

Codecov Report

❌ Patch coverage is 88.54167% with 22 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.63%. Comparing base (029441f) to head (d3b52dd).
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
onnxscript/_internal/tape_builder.py 86.80% 14 Missing and 5 partials ⚠️
onnxscript/_internal/builder.py 91.89% 2 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2910      +/-   ##
==========================================
+ Coverage   72.61%   72.63%   +0.01%     
==========================================
  Files         259      259              
  Lines       31597    31651      +54     
  Branches     2973     2980       +7     
==========================================
+ Hits        22945    22990      +45     
- Misses       7643     7653      +10     
+ Partials     1009     1008       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Comment thread onnxscript/rewriter/rules/common/_remove_optional_bias.py
Comment thread onnxscript/_internal/tape_builder.py Fixed
Rename and restructure:
- Rename OpBuilderBase to BuilderBase
- Merge op() and _make_node() into single op() method
- Promote _create_op to BuilderBase.call_op() as the core node-creation method

Feature flags (BuilderFeature enum):
- SCHEMA_PARTITION, CAST_INPUTS, CAST_ATTRIBUTES, INFER_SHAPES,
  CONSTANT_PROPAGATION
- FULL = all flags (used by GraphBuilder)
- NONE = no extra processing (used by TapeBuilder)

GraphBuilder inherits from BuilderBase:
- Shares call_op() with feature-gated hooks
- Default hook implementations (_get_schema, _partition_inputs_attributes,
  _cast_inputs, _cast_attributes, _constant_propagation, _infer_shapes)
  live in BuilderBase
- GraphBuilder only overrides: _promote_constant (cached initializers),
  _generate_node_name, _adapt_outputs (named outputs), _annotate_node
  (scope metadata), _qualify_value_name (scope prefixing)

Constant promotion in BuilderBase:
- _input_to_ir_value with scalar-to-tensor promotion and CastLike
- _promote_constant hook: base class uses Constant nodes (no name
  collisions), GraphBuilder uses cached initializers

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@gramalingam gramalingam force-pushed the rama/unify-builders branch from f90ed1f to 6cb4e71 Compare May 10, 2026 23:16
@gramalingam gramalingam marked this pull request as ready for review May 10, 2026 23:18
@gramalingam gramalingam changed the title [DRAFT] Unify BuilderBase and GraphBuilder: shared call_op with feature flags Unify BuilderBase and GraphBuilder: shared call_op with feature flags May 10, 2026
Comment thread onnxscript/_internal/tape_builder.py Fixed
@gramalingam gramalingam enabled auto-merge (squash) May 10, 2026 23:25
…27 lint

- Remove unused _PYTHON_TYPE_TO_DTYPE from tape_builder.py (builder.py
  has its own local copy that is actually used)
- Remove dead import of _PYTHON_TYPE_TO_DTYPE in builder.py
- Add noqa: B027 to _annotate_node (intentional no-op hook, not abstract)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@gramalingam gramalingam force-pushed the rama/unify-builders branch from d2b25f5 to d3b52dd Compare May 10, 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

Development

Successfully merging this pull request may close these issues.

3 participants