[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 251ae68

Description

• Move statement registration (registerVal, registerExpr) and removal (removeVal, removeExpr) logic from IrContainer directly into Fusion

• Inline the container manipulation code in Fusion methods instead of delegating to IrContainer

• Remove the corresponding declarations and implementations from container.cpp and container.h

• Add Fusion as friend to IrContainerPasskey for setName() calls, and remove vestigial friend class StatementGuard

Changes walkthrough

Relevant files

Enhancement

fusion.cpp Inline registration/removal logic in Fusion                            csrc/fusion.cpp Inline removeExpr logic to directly erase from container's exprs_up_ and exprs_ Inline removeVal logic to directly erase from container's vals_up_ and vals_ Inline registerVal logic to add to vals_up_, vals_, and call setName Inline registerExpr logic to add to exprs_up_, exprs_, and call setName +28/-4    container.cpp Remove registration methods from IrContainer                          csrc/ir/container.cpp Remove implementations of removeExpr, removeVal, registerVal, registerExpr +0/-58    container.h Remove registration declarations, add Fusion friend            csrc/ir/container.h Remove declarations of removeExpr, removeVal, registerVal, registerExpr Add friend class Fusion to IrContainerPasskey for passkey construction Remove vestigial friend class StatementGuard +1/-14

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Non-deterministic iteration A NOLINTNEXTLINE comment was added at line 632 to suppress a warning about non-deterministic pointer iteration order. This suggests the iteration over all_inputs may not have a defined order, which could affect determinism. Please verify this is intentional and acceptable for the project's determinism requirements. // NOLINTNEXTLINE(bugprone-nondeterministic-pointer-iteration-order) Code duplication The registration and removal logic is now duplicated in Fusion (for vals and exprs). While this is necessary for the refactoring, ensure both paths are kept in sync if future changes are made. Consider adding helper functions if this pattern grows. auto* c = ir_container(); auto expr_in_deque = std::ranges::find_if( c->exprs_up_, [expr](std::unique_ptr<Expr>& expr_up) { return expr_up.get() == expr; }); NVF_ERROR( expr_in_deque != c->exprs_up_.end(), "Wanted to remove an expression but its unique ptr is missing."); c->exprs_.erase(expr); c->exprs_up_.erase(expr_in_deque); } Code duplication The registration and removal logic is now duplicated in Fusion (for vals and exprs). While this is necessary for the refactoring, ensure both paths are kept in sync if future changes are made. Consider adding helper functions if this pattern grows. auto* c = ir_container(); auto val_in_deque = std::ranges::find_if( c->vals_up_, [val](std::unique_ptr<Val>& val_up) { return val_up.get() == val; }); NVF_ERROR( val_in_deque != c->vals_up_.end(), "Wanted to remove a value but its unique ptr is missing."); c->vals_.erase(val); c->vals_up_.erase(val_in_deque);

[mdavis36]

!test

[mdavis36]

!test

[greptile-apps]

Greptile Summary

This PR consolidates statement registration (registerVal, registerExpr, removeVal, removeExpr) directly into Fusion by inlining what was previously delegated to IrContainer. It is the final foundation change before the Phase 2 shared-container transition, ensuring Fusion (which knows the owning fusion identity) is the single authority for all write-path mutations.

Key changes:

• Fusion::registerVal/registerExpr now directly populate IrContainer's vals_up_/exprs_up_ and the corresponding sets, constructing IrContainerPasskey directly (made possible by adding friend class Fusion to IrContainerPasskey).
• Fusion::removeVal/removeExpr now directly search the deques and erase from both the deque and the set; the pre-existing assertInContainer call at the start of each method preserves the existence validation that was previously duplicated inside IrContainer::removeExpr/removeVal.
• friend class StatementGuard is correctly removed from IrContainer since StatementGuard only ever calls public Fusion APIs (numExprs(), numVals(), removeStatementsCreatedAfter()).
• Kernel::registerVal/registerExpr correctly continue to delegate to Fusion::registerVal/registerExpr, so the subclass chain is intact.
• A NOLINTNEXTLINE(bugprone-nondeterministic-pointer-iteration-order) suppression is added for the existing iteration over an std::unordered_set<Val*> in validateInputs.

Confidence Score: 4/5

• This PR is safe to merge—it is pure code motion with equivalent behavior and no new logic introduced.
• All four inlined methods produce identical runtime behavior: the existence-validation previously done by a set-lookup in IrContainer is fully covered by the pre-existing assertInContainer call at the top of each Fusion method. The IrContainerPasskey friend extension is correct and minimal. The StatementGuard friend removal is verified safe. Kernel's delegation chain is intact. Score is 4 rather than 5 only because the PR directly reaches into IrContainer's protected storage from Fusion (a deliberate architectural choice for the Phase 2 transition) which slightly increases the surface for future maintenance errors.
• No files require special attention; all changes are straightforward code motion.

Sequence Diagram

sequenceDiagram
    participant C as Client Code
    participant F as Fusion
    participant IC as IrContainer

    note over C,IC: Before this PR
    C->>F: registerVal(val)
    F->>IC: ir_container()->registerVal(val)
    IC->>IC: vals_up_.emplace_back(val)
    IC->>IC: vals_.insert(val)
    IC->>IC: val->setName(IrContainerPasskey(), ...)

    note over C,IC: After this PR
    C->>F: registerVal(val)
    F->>IC: ir_container() → get pointer
    F->>IC: c->vals_up_.emplace_back(val)
    F->>IC: c->vals_.insert(val)
    F->>F: val->setName(IrContainerPasskey(), ...) [Fusion is now friend of passkey]

Loading

_{Last reviewed commit: 0716dc5}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

csrc/fusion.cpp, line 358
Dropped set-membership pre-check in removeExpr

The old IrContainer::removeExpr had an explicit guard:

NVF_ERROR(
    exprs_.find(expr) != exprs_.end(),
    "Wanted to remove an expression but it doesn't exist in this container.");

before searching exprs_up_. In the new code this check is absent between the input/output cleanup loops and the deque search. In practice this is safe — assertInContainer at the top of Fusion::removeExpr (line 329) already verifies the expr is in exprs_ before any mutation happens. Worth noting explicitly that the guard has moved to the function's entry point rather than immediately preceding the erase, so a future reader doesn't think it's missing.

Same pattern applies to removeVal (line 410–417 below).

[mdavis36]

!test

[mdavis36]

!test