[IR Container] Phase 2.2 Fusion Statement Registration

[mdavis36]

Summary

Inline registerVal, registerExpr, removeVal, and removeExpr logic directly into Fusion, eliminating the delegation to IrContainer. This consolidates the statement registration path after per-Fusion special values, axioms, and metadata were moved to Fusion in PR #5954 .

Key Changes

• fusion.cpp: registerVal, registerExpr, removeVal, removeExpr implemented directly on Fusion, no longer delegating to IrContainer
• container.cpp: Remove registration/removal implementations
• container.h: Remove registerVal/registerExpr/removeVal/removeExpr declarations, remove vestigial friend class StatementGuard (it only uses public Fusion API)
• container.h: Add Fusion as friend of IrContainerPasskey so it can construct passkeys for setName() calls

Why This Matters

Statement registration is the write path for container mutation. In Phase 2, when multiple Fusions share a container, registration must populate per-Fusion tracking maps (per_fusion_vals_, per_fusion_exprs_). Having registration live in Fusion (which knows this — the owning Fusion) rather than in IrContainer (which doesn't know which Fusion is registering) is essential for correct per-Fusion tracking.

Relationship to Phase 2

This is the final foundation change before the shared_ptr transition begins. Now Fusion is the single authority for:

Fusion owns and manages:
  ✓ Special values (zero_val_, one_val_, etc.)     ← PR A
  ✓ Axioms and metadata                            ← PR A
  ✓ Statement registration/removal                 ← PR B (this PR)
  ✓ Copy/move/swap semantics                       ← existing

IrContainer is now pure storage:
  - vals_up_, exprs_up_ (owning unique_ptrs)
  - vals_, exprs_ (lookup sets)
  - deterministic_vals_, deterministic_exprs_ (insertion-ordered)
  - Name counters (val_type_name_map_, expr_name_counter_) <- Fusion will manage this in #5964.

CI Risk

Low. Pure code motion — identical behavior, just different call location. The registration logic is unchanged; only the method boundaries moved.

Note : the lint error here will is resolved in later PRs in this chain

[github-actions]

Review updated until commit 54cd0fe

Description

• Inlines registerVal/registerExpr/removeVal/removeExpr logic directly into Fusion class

• Eliminates delegation pattern to IrContainer for statement registration

• Consolidates registration path after per-Fusion special values migration

• Removes StatementGuard friend class and adds Fusion as IrContainerPasskey friend

Changes walkthrough

Relevant files

Enhancement

fusion.cpp Inline registration/removal logic into Fusion                        csrc/fusion.cpp Inlines removeExpr logic directly instead of delegating to IrContainer Inlines removeVal logic directly instead of delegating to IrContainer Inlines registerVal logic directly instead of delegating to IrContainer Inlines registerExpr logic directly instead of delegating to IrContainer +31/-4    container.cpp Remove delegated registration methods                                        csrc/ir/container.cpp Removes removeExpr method implementation Removes removeVal method implementation Removes registerVal method implementation Removes registerExpr method implementation +0/-60    container.h Update friend relationships and method declarations            csrc/ir/container.h Removes method declarations for removeExpr, removeVal, registerVal, registerExpr Adds Fusion as friend class to IrContainerPasskey Removes StatementGuard friend class declaration +1/-14

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ No major issues detected

Test failures

• (Medium, 3) CUDA OOM in AGMatmul and StreamParallelBackend tests on A100 runner

  Test Name	A100	Source
  AGMatmulTest.CollectiveBasedPipeline/AG_before_Matmul	❌	Link
  StreamParallelBackendTest.AG_matmul_P2p/Broadcast_Nccl	❌	Link
  StreamParallelBackendTest.AG_matmul_P2p/p2p_Cuda	❌	Link

[mdavis36]

!test

[mdavis36]

!test

[greptile-apps]

Greptile Summary

This PR moves statement registration (registerVal, registerExpr) and removal (removeVal, removeExpr) logic from IrContainer into Fusion, completing Phase 2.2 of the IR container refactor. The logic is unchanged — only the method boundaries moved. Fusion now directly manipulates IrContainer's internal data structures (vals_up_, vals_, exprs_up_, exprs_) via its existing friend access, and constructs IrContainerPasskey objects for setName() calls (enabled by adding Fusion as a friend of IrContainerPasskey).

• csrc/fusion.cpp: Inlines the container storage manipulation (deque find/erase, set insert/erase, name assignment) that was previously delegated to IrContainer
• csrc/ir/container.cpp: Removes the four now-unused IrContainer::registerVal, registerExpr, removeVal, removeExpr implementations (~60 lines)
• csrc/ir/container.h: Removes the four method declarations, adds Fusion as friend of IrContainerPasskey, and removes the vestigial friend class StatementGuard declaration (it only uses public Fusion API)
• No remaining callers of the removed IrContainer methods exist in the codebase; all external call sites already go through Fusion or its subclasses (Kernel, HostIrContainer)

Confidence Score: 5/5

• This PR is safe to merge — it is pure code motion with no behavioral changes.
• The registration and removal logic is byte-for-byte identical after the move. The only structural differences are: (1) Fusion's existing assertInContainer checks subsume the old IrContainer existence checks, (2) Fusion now constructs IrContainerPasskey directly (enabled by the new friend declaration). All external callers already target Fusion (not IrContainer), and the Kernel subclass overrides correctly call Fusion::registerVal/registerExpr as their base.
• No files require special attention.

Important Files Changed

Filename	Overview
csrc/fusion.cpp	Inlines registerVal, registerExpr, removeVal, removeExpr logic that was previously delegated to IrContainer. The code is functionally identical to the old IrContainer implementations, with Fusion's existing assertInContainer checks covering the removed container-level existence checks.
csrc/ir/container.cpp	Removes ~60 lines of registerVal, registerExpr, removeVal, removeExpr implementations that have been moved to Fusion. No remaining callers exist in the codebase.
csrc/ir/container.h	Removes four method declarations, adds Fusion as friend of IrContainerPasskey (needed for setName calls), and removes vestigial friend class StatementGuard (it only uses public Fusion API).

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    subgraph Before["Before (delegated)"]
        A1["Fusion::registerVal(val)"] --> B1["IrContainer::registerVal(val)"]
        B1 --> C1["vals_up_.emplace_back(val)"]
        B1 --> D1["vals_.insert(val)"]
        B1 --> E1["val->setName(...)"]
        
        A2["Fusion::removeVal(val)"] --> B2["IrContainer::removeVal(val)"]
        B2 --> C2["vals_up_.erase(...)"]
        B2 --> D2["vals_.erase(val)"]
    end

    subgraph After["After (inlined in Fusion)"]
        F1["Fusion::registerVal(val)"] --> G1["c->vals_up_.emplace_back(val)"]
        F1 --> H1["c->vals_.insert(val)"]
        F1 --> I1["val->setName(IrContainerPasskey(), ...)"]
        
        F2["Fusion::removeVal(val)"] --> G2["c->vals_up_.erase(...)"]
        F2 --> H2["c->vals_.erase(val)"]
    end

    Before -.->|"code motion"| After

Loading

_{Last reviewed commit: 54cd0fe}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}