[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 cf4ae4a

Description

• Moved special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_) from IrContainer to Fusion class for better ownership semantics

• Implemented lazy creation pattern for all special value accessors with caching to prevent recreation overhead

• Updated Fusion copy, swap, clear, and remove operations to properly handle per-Fusion special values

• Added regression test for StatementGuard scope to ensure special vals don't become dangling pointers after rollback

• Moved metadata, axioms, and assumption methods from IrContainer to Fusion class

Changes walkthrough

	Relevant files
Enhancement	fusion.cpp Implement per-Fusion special values with lazy creation      csrc/fusion.cpp Added lazy creation methods for special values (zeroVal, oneVal, falseVal, trueVal, magicZeroVal) Updated Fusion copy operation to remap cached special val pointers through IrCloner Modified Fusion swap operation to exchange per-Fusion special values Enhanced Fusion clear operation to reset special value pointers Updated Fusion removeVal to skip cached special vals during removal Implemented removeStatementsCreatedAfter with special value cache nullification Moved metadataOf, axioms, assumePositive, and assumeNonNegative methods from IrContainer +211/-0  container.cpp Remove special value management from IrContainer                  csrc/ir/container.cpp Removed special value shortcuts (zeroVal, oneVal, falseVal, trueVal, magicZeroVal) from IrContainer Eliminated metadata, axioms, and assumption methods from IrContainer Removed removeStatementsCreatedAfter method from IrContainer Updated IrContainer copy and swap operations to exclude special value handling Simplified IrContainer clear operation by removing special value cleanup +2/-180  statement_guard.cpp Update StatementGuard to use simplified numVals call          csrc/statement_guard.cpp Updated StatementGuard to call fusion_->numVals() instead of fusion_->numVals(/*include_shortcuts=*/false) Simplified numVals call by removing include_shortcuts parameter +1/-1      fusion.h Add per-Fusion special value members and accessors              csrc/fusion.h Added private members for per-Fusion special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_) Declared public accessor methods for special values with lazy creation Added axioms_ unique_ptr member and metadata_ unordered_map Updated method declarations to move from IrContainer to Fusion class Removed include_shortcuts parameter from numVals method +25/-46  container.h Remove special value management from IrContainer interface csrc/ir/container.h Removed special value shortcut method declarations from IrContainer Eliminated metadata, axioms, and assumption method declarations Removed removeStatementsCreatedAfter method declaration Simplified numVals method to remove include_shortcuts parameter Removed special value member variables and metadata structures +2/-50
Bug fix	utils.cpp Add null check for reference-substitute comparison              csrc/ir/utils.cpp Added null check in SubstituteInExpr constructor to avoid unnecessary mutations when reference equals substitute +3/-1
Tests	test_indexing.cpp Update test expectations for reshape indexing                        tests/cpp/test_indexing.cpp Updated expected string in reshape test case to match new indexing variable names (i113, i114, i115 instead of i114, i115, i116) +2/-2      test_statement_guard.cpp Add regression test for special vals in StatementGuard scope tests/cpp/test_statement_guard.cpp Added regression test "LazySpecialValsNotDangling" to verify special vals created inside StatementGuard scope remain valid after rollback Test validates that zeroVal, oneVal, trueVal, and falseVal are properly recreated after StatementGuard destruction +51/-0

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Raw Pointer Ownership The special value pointers (zero_val_, one_val_, etc.) are declared as raw pointers in the Fusion class, but the actual objects are owned by IrContainer. This creates a potential dangling pointer issue if the IrContainer is cleared or destroyed while the Fusion still holds references to these pointers. The current nullification logic in removeStatementsCreatedAfter helps mitigate this, but the ownership model could be clearer. Val* zero_val_ = nullptr; Val* one_val_ = nullptr; Val* true_val_ = nullptr; Val* false_val_ = nullptr; NamedScalar* magic_zero_val_ = nullptr; Lazy Initialization Thread Safety The lazy initialization pattern in methods like zeroVal(), oneVal(), etc. uses simple null checks without any synchronization. In a multi-threaded context, this could lead to race conditions where multiple threads simultaneously create the same special value, potentially causing memory leaks or inconsistent state. Val* Fusion::zeroVal() { if (!zero_val_) { zero_val_ = IrBuilder::createInContainer<Val>(this, 0L, DataType::Index); } return zero_val_; } Val* Fusion::oneVal() { if (!one_val_) { one_val_ = IrBuilder::createInContainer<Val>(this, 1L, DataType::Index); } return one_val_; } Val* Fusion::falseVal() { if (!false_val_) { false_val_ = IrBuilder::createInContainer<Val>(this, false, DataType::Bool); } return false_val_; } Val* Fusion::trueVal() { if (!true_val_) { true_val_ = IrBuilder::createInContainer<Val>(this, true, DataType::Bool); } return true_val_; } NamedScalar* Fusion::magicZeroVal() { if (!magic_zero_val_) { magic_zero_val_ = IrBuilder::createInContainer<NamedScalar>( this, kMagicZeroName, DataType::Index); } return magic_zero_val_; } Missing Error Handling The metadataOf() method and axioms() method don't have comprehensive error handling for edge cases. For example, what happens if metadataOf() is called with a null pointer, or if axioms() fails to create the parallel index/dim NamedScalars? Adding proper error checking would make the code more robust. Val* Fusion::metadataOf(Val* v) { if (metadata_.count(v) == 0) { auto metadata_val = IrBuilder::createInContainer<Val>(this, metaDataTypeOf(v)); auto metadata_expr = IrBuilder::createInContainer<GetMetaData>(this, metadata_val, v); metadata_[v] = std::make_pair(metadata_val, metadata_expr); } return metadata_.at(v).first; } const std::vector<Val*>& Fusion::axioms() { if (!axioms_) { axioms_ = std::make_unique<std::vector<Val*>>(); axioms_->reserve(kParallelTypeThreads.size() * 3); auto zero = zeroVal(); for (auto p : kParallelTypeThreads) { auto pidx = NamedScalar::getParallelIndex(p); auto pdim = NamedScalar::getParallelDim(p); axioms_->push_back(SimplifyingIrBuilder::geExpr(pidx, zero)); axioms_->push_back(SimplifyingIrBuilder::gtExpr(pdim, zero)); axioms_->push_back(SimplifyingIrBuilder::ltExpr(pidx, pdim)); } } return *axioms_; }

[mdavis36]

!test

[mdavis36]

!test

[mdavis36]

!test

[mdavis36]

!test

[greptile-apps]

Greptile Summary

This PR migrates per-Fusion state — special values (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_), axioms, and metadata — from IrContainer to Fusion, eliminating the 1:1 coupling between them. This is a prerequisite for Phase 2's shared-container model where multiple Fusions share a single IrContainer.

• Ownership model change: Special vals move from unique_ptr members in IrContainer (stored outside vals_up_) to raw cache pointers in Fusion (with memory owned by IrContainer::vals_up_). This simplifies the numVals API by removing the include_shortcuts parameter.
• Bug fix — SubstituteInExpr self-substitution: With per-Fusion special vals, cloned zero_val_ can be pointer-identical to IterDomain extents, causing a self-mapping that tripped a two-hop assertion. Fixed by guarding against reference == substitute.
• Bug fix — 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. removeStatementsCreatedAfter now nulls any special val caches pointing to vals about to be destroyed.
• Proper clone/swap/clear: Fusion::copy remaps special val cache pointers through IrCloner, Fusion::swap swaps all new fields, and Fusion::clear resets them. Axioms and metadata are also properly cloned and swapped.
• Regression test added for the StatementGuard dangling pointer scenario.

Confidence Score: 4/5

• This PR is safe to merge — it's a well-structured ownership migration with bug fixes, regression tests, and no identified logic issues.
• The migration is clean and methodical: all call sites already went through Fusion (via Statement::container() which returns Fusion*), so the API surface is unchanged for callers. The two bug fixes address real correctness issues with proper guards and a regression test. The only reason for not giving a 5 is the inherent complexity of ownership changes in a large IR system — there could be edge cases in rarely-exercised paths that are hard to verify statically.
• csrc/fusion.cpp deserves the most attention as it contains the core migration logic, especially removeStatementsCreatedAfter (dangling pointer prevention) and Fusion::copy (clone correctness).

Important Files Changed

Filename	Overview
csrc/fusion.cpp	Core of the migration: moves special vals, axioms, metadata from IrContainer to Fusion. Adds removeStatementsCreatedAfter with dangling-pointer prevention, removeVal guard for cached special vals, lazy creation methods, and proper clone/swap/clear handling. Well-structured with clear separation of concerns.
csrc/fusion.h	Declares the new per-Fusion special val caches (raw pointers), axioms, and metadata members. Removes delegating wrappers that forwarded to IrContainer. Simplifies numVals API by removing the include_shortcuts parameter.
csrc/ir/container.cpp	Removes all special val, axiom, metadata, and removeStatementsCreatedAfter implementations from IrContainer. Removes the special-val guard from removeVal. Clean deletion of ~150 lines of now-dead code.
csrc/ir/container.h	Removes special val unique_ptrs, axioms, metadata members, public API methods, and removeStatementsCreatedAfter declaration. Simplifies numVals by removing the include_shortcuts parameter. IrContainer is now purely a storage layer.
csrc/ir/utils.cpp	Bug fix: guards SubstituteInExpr against self-substitution (reference == substitute) which would create a self-mapping that triggers a two-hop assertion. Surfaced by per-Fusion special vals where cloned pointers can be identical.
csrc/statement_guard.cpp	Trivial change: updates numVals() call to match simplified API (removes include_shortcuts parameter). Behavior is unchanged since special vals are now in vals_up_ anyway.
tests/cpp/test_indexing.cpp	Updates hardcoded val name indices (i114→i113, etc.) that shifted because special vals are no longer stored outside vals_up_, changing the name counter offset by 1. Expected consequence of the migration.
tests/cpp/test_statement_guard.cpp	Adds regression test for the dangling special val pointer bug: lazily creates special vals inside a StatementGuard, rolls back, then verifies vals are still valid (re-created on next access) and fusion remains executable.

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
    class Fusion {
        -Val* zero_val_
        -Val* one_val_
        -Val* true_val_
        -Val* false_val_
        -NamedScalar* magic_zero_val_
        -unique_ptr~vector~Val*~~ axioms_
        -unordered_map~Val*, pair~ metadata_
        -unique_ptr~IrContainer~ ir_container_
        +zeroVal() Val*
        +oneVal() Val*
        +trueVal() Val*
        +falseVal() Val*
        +magicZeroVal() NamedScalar*
        +axioms() vector~Val*~
        +metadataOf(Val*) Val*
        +assumePositive(Val*)
        +assumeNonNegative(Val*)
        +removeStatementsCreatedAfter(int64_t, int64_t)
        +removeVal(Val*)
        +copy(Fusion*, Fusion*) IrCloner
        +swap(Fusion, Fusion)
        +clear()
    }

    class IrContainer {
        +deque~unique_ptr~Val~~ vals_up_
        +unordered_set~Val*~ vals_
        +deque~unique_ptr~Expr~~ exprs_up_
        +unordered_set~Expr*~ exprs_
        -Fusion* parent_
        +registerVal(Val*)
        +registerExpr(Expr*)
        +removeVal(Val*)
        +removeExpr(Expr*)
        +clear()
        +numVals() int64_t
    }

    class StatementGuard {
        -Fusion* fusion_
        -int64_t prev_num_exprs_
        -int64_t prev_num_vals_
        +~StatementGuard()
    }

    Fusion "1" *-- "1" IrContainer : owns via unique_ptr
    Fusion ..> IrContainer : "special vals memory\nowned by vals_up_"
    StatementGuard ..> Fusion : "calls removeStatementsCreatedAfter\non destruction"

Loading

_{Last reviewed commit: cf4ae4a}

[greptile-apps]

_{8 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}