[IR Container] Phase 2.1 Ownership Fusion Special Vals & Axioms

[mdavis36]

Summary

Move per-Fusion state — special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_), axioms, and metadata — from IrContainer to Fusion. This eliminates the 1:1 relationship assumption between IrContainer and Fusion that would break when multiple Fusions share a container in Phase 2.

Key Changes

Special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_):

• Moved from IrContainer to Fusion as raw cache pointers (memory still owned by IrContainer's vals_up_)
• Lazy creation preserved — created on first access via IrBuilder::createInContainer
• Fusion::clear() nulls cache pointers so they re-create on next access
• StatementGuard rollback nulls any special vals created within the guard scope to prevent dangling pointers

Axioms and metadata (axioms_, metadata_):

• Moved from IrContainer to Fusion
• Methods that used parent_ to create vals (metadataOf, axioms, assumePositive, assumeNonNegative) reimplemented on Fusion
• Properly swapped in Fusion::swap and cloned in Fusion::copy

Bug fixes surfaced by the migration:

• SubstituteInExpr self-substitution: with per-Fusion special vals, zero_val_ is now the same pointer as IterDomain extents (both go through IrCloner). SubstituteInExpr(ext, zero) where ext == zero created a self-mapping that tripped a two-hop assertion. Fixed by guarding against reference == substitute.
• Dangling special vals after StatementGuard rollback: special vals lazily created inside a guard scope were destroyed on rollback while Fusion cache pointers still referenced them.

Relationship to Phase 2

This is a hard prerequisite for shared containers. Without per-Fusion special vals:

Problem: Shared container with IrContainer-owned special vals

  Fusion A ─┐
             ├──→ IrContainer ──→ { zero_val_ = 0x1111 }
  Fusion B ─┘

  ~Fusion A:  IrContainer still alive (refcount > 0)
              But A's destruction might invalidate zero_val_
              B calls zeroVal() → dangling pointer or wrong Fusion's val

With per-Fusion special vals, each Fusion's zeroVal() returns its own independently-cached val. Destroying Fusion A has zero effect on Fusion B's special values.

CI Risk

Low. Each commit is independently CI-clean. The migration preserves identical behavior for the single-Fusion case (all existing usage). The two bug fixes (SubstituteInExpr, StatementGuard rollback) address real correctness issues that were latent in the old code and only surfaced because the migration changed pointer identity relationships.

[mdavis36]

!test

[github-actions]

Review updated until commit 3693e3d

Description

• Moved special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_) from IrContainer to Fusion as raw cache pointers

• Moved axioms and metadata from IrContainer to Fusion, with methods reimplemented on Fusion

• Fixed SubstituteInExpr self-substitution bug: guard against reference == substitute to prevent assertion failure

• Fixed dangling special vals bug: StatementGuard rollback now nulls cached special val pointers to prevent use-after-free

Changes walkthrough

	Relevant files
Enhancement	fusion.cpp Implement per-Fusion special values and axioms                      csrc/fusion.cpp Added special value cache pointers (zero_val_, one_val_, etc.) and moved axioms_/metadata_ from IrContainer Implemented lazy creation for zeroVal(), oneVal(), trueVal(), falseVal(), magicZeroVal() with dtype variants Added metadataOf(), axioms(), assumePositive(), assumeNonNegative() methods to Fusion Updated Fusion::swap() and Fusion::copy() to handle new per-Fusion state Updated Fusion::clear() to reset special value caches and clear axioms/metadata Added nulling of special val caches in removeStatementsCreatedAfter() to prevent dangling pointers +211/-0  fusion.h Add per-Fusion special value members to Fusion class          csrc/fusion.h Added private members: zero_val_, one_val_, true_val_, false_val_, magic_zero_val_, axioms_, metadata_ Changed method signatures from delegating to ir_container() to direct implementations Removed include_shortcuts parameter from numVals() +27/-48  container.cpp Remove special value methods from IrContainer                        csrc/ir/container.cpp Removed all special value methods (zeroVal, oneVal, trueVal, falseVal, magicZeroVal) Removed axioms_, metadata_, and related methods (lazyInitAxioms, assumePositive, assumeNonNegative) Removed removeStatementsCreatedAfter() method (moved to Fusion) Simplified deterministic_vals/exprs using std::ranges::transform +14/-199 container.h Remove special value members from IrContainer header          csrc/ir/container.h Removed special value cache members (true_val_, false_val_, one_val_, zero_val_, magic_zero_val_) Removed axioms_ and metadata_ members Removed method declarations for special values and axioms Simplified numVals() to not have include_shortcuts parameter +2/-50    statement_guard.cpp Update StatementGuard to use simplified numVals                    csrc/statement_guard.cpp Changed prev_num_vals_ to use numVals() without include_shortcuts parameter +1/-1
Bug fix	utils.cpp Fix SubstituteInExpr self-substitution bug                              csrc/ir/utils.cpp Fixed SubstituteInExpr self-substitution bug by guarding against reference != substitute before adding to mutations map Updated integer comparisons to use std::cmp_equal for safer comparison Added NOLINT comments for non-deterministic pointer iteration order +8/-3
Tests	test_statement_guard.cpp Add regression test for dangling special vals                        tests/cpp/test_statement_guard.cpp Added regression test LazySpecialValsNotDangling to verify special vals created inside StatementGuard scope don't become dangling after rollback +51/-0    test_indexing.cpp Update test expected output for reshape                                    tests/cpp/test_indexing.cpp Updated expected string in Reshape test to reflect internal variable numbering change +2/-2

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Memory ownership clarity The special value pointers (zero_val_, one_val_, etc.) are now raw pointers in Fusion but the actual memory is still owned by IrContainer's vals_up_. This is intentional based on the PR description, but the code comments should be very clear about this ownership model to prevent future bugs. The current comment "the vals themselves are owned by ir_container" at line 312 is good but could be more explicit at the member declaration site in fusion.h. os << " " << out << '\n'; } os << "\n%kernel {\n"; IrPrinter op_exprs(os); op_exprs.handle(this); if (include_tensor_transforms) { os << "\nTransformPrinter : \n"; IrTransformPrinter t_exprs(os); Potential null pointer safety In Fusion::removeVal, special val cache pointers are checked for equality with the val being removed. However, if a special val was cloned (e.g., during Fusion::copy), the cloned version would have a different pointer and wouldn't be caught by these checks. Consider if this edge case needs handling or documentation. if (val == zero_val_ || val == one_val_ || val == true_val_ || val == false_val_ || val == magic_zero_val_) { return; } Self-substitution guard correctness The fix for SubstituteInExpr self-substitution (guarding against reference == substitute) is correct and addresses the reported bug. However, this is a behavioral change that could affect existing code that relied on the previous assertion failure. Ensure this change is properly documented and the test coverage is sufficient. explicit SubstituteInExpr(Val* reference, Val* substitute) { if (reference != substitute) { mutations_[reference] = substitute; } }

[mdavis36]

!test

[mdavis36]

!test

[mdavis36]

!test

[mdavis36]

!test

[greptile-apps]

Greptile Summary

This PR successfully moves per-Fusion state — special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_), axioms, and metadata — from IrContainer to Fusion, which is a hard prerequisite for Phase 2's shared-container design where multiple Fusion objects can co-own a single IrContainer.

Key changes verified:

• Special vals are now raw cache pointers on Fusion, memory still owned by IrContainer::vals_up_; lazy creation and Fusion::clear() / StatementGuard rollback paths updated correctly.
• Axioms and metadata moved to Fusion; Fusion::copy and Fusion::swap properly updated to handle them.
• StatementGuard constructor eagerly pre-initializes axioms via fusion->axioms() to ensure their backing vals fall below the rollback watermark.
• SubstituteInExpr self-substitution bug fixed by guarding reference != substitute.
• numVals(bool include_shortcuts) simplified to numVals() since special-val "side-channel" allocation no longer exists.
• std::transform + iterator pairs replaced with std::ranges::transform throughout IrContainer.

All existing behavior is preserved for the single-Fusion case, and the two explicit bug fixes (SubstituteInExpr, StatementGuard rollback) are correct and address real correctness issues.

Confidence Score: 5/5

• PR is safe to merge. The migration is coherent, all verified behavior is correct, and the two explicit bug fixes properly address latent correctness issues surfaced by the migration.
• No verified issues found. The PR successfully completes a significant architectural migration (per-Fusion ownership of special values, axioms, and metadata) that enables Phase 2's shared-container design. All commits are CI-clean, the special-val dangling pointer bug is properly fixed with null cache pointers on rollback, and existing behavior is preserved for the single-Fusion case.
• No files require special attention

Sequence Diagram

sequenceDiagram
    participant Fusion
    participant IrContainer
    participant IrBuilder

    Note over Fusion,IrBuilder: Lazy special-val creation (new design)
    Fusion->>Fusion: zeroVal() — zero_val_ == nullptr
    Fusion->>IrBuilder: createInContainer(this, 0L, Index)
    IrBuilder->>IrContainer: vals_up_.push_back(unique_ptr<Val>)
    IrBuilder-->>Fusion: Val* stored as zero_val_

    Note over Fusion,IrContainer: StatementGuard rollback
    Fusion->>Fusion: removeStatementsCreatedAfter(num_exprs, num_vals)
    Fusion->>Fusion: null zero_val_/one_val_/true_val_/false_val_/magic_zero_val_ if about to be destroyed
    Fusion->>IrContainer: vals_.erase + vals_up_.pop_back per destroyed val

    Note over Fusion,IrContainer: Fusion copy
    Fusion->>IrContainer: IrContainer::copy clones all vals_up_ including special vals
    Fusion->>Fusion: remap to->zero_val_ via ir_cloner cache lookup
    Fusion->>Fusion: clone axioms_ and metadata_ entries

Loading

_{Last reviewed commit: 4111350}

[greptile-apps]

_{8 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[mdavis36]

!test

[mdavis36]

!test

[wujingyue]

special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_)

Why are these scalar values per Fusion not per IrContainer?

[mdavis36]

special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_)

Why are these scalar values per Fusion not per IrContainer?

For Phase 2 Statements need to maintain a 1:1 relationship with the registered Fusion object since we frequently query Fusion* Statement::fusion(). If the special values are owned by a container with multiple fusions that invariant breaks (this is one of the core behaviors that need to be carefully addressed w/ scalar sharing in Phase 3) . Phase 2 maintains this the 1:1 relationship between statement and fusions.

Don't we need to move these back to IrContainer in Phase 3 with scalar sharing?
For current work in Phase 3 scalar sharing Fusions will actually continue to own special vals uniquely. This is in order to help with significant lock contention we would encounter with special values. They are lazily constructed, if they were owned by the IrContainer and shared across fusions every read / write access would need to be locked to check:

1. The value exists.
2. The calling fusion is / isn't registered.

Since special values have a known static value it is safe to maintain them uniquely for each Fusion.

[wujingyue]

They are lazily constructed

That's true but I suspect that has been a premature optimization. Creating them in a shared IrContainer will likely

1. be more efficient because their storage can be shared, and
2. lead to a simpler contract -- scalars are per-container and TensorViews are per-fusion.

(Non-blocking of course)

[mdavis36]

They are lazily constructed

That's true but I suspect that has been a premature optimization. Creating them in a shared IrContainer will likely

1. be more efficient because their storage can be shared, and
2. lead to a simpler contract -- scalars are per-container and TensorViews are per-fusion.

(Non-blocking of course)

This was my first instinct too. I thought about making all special values eagerly initialized during IrContainer construction however this still hits the lock contention issue since you need to check / initialize registration of the special value on it's first use to take it from an "un-owned" state to an "owned" one to ensure you have registered the Fusion object with the pre-initialized Val (this relationship is necessary for checks like inContainer or traversal).

We've already seen a regression due to lock contention on parallel builds. My current approach to combat this (eventually) is to separate the unique and shared IR, with the hopes of reducing the IR nodes that we need to live in a thread safe container. Since the nature of the shared values is that they are used extensively throughout the IR, keeping them per fusion is easiest and will stop the flip-flopping of ownership between fusion and the shared container.