Review updated until commit f3e849f

Description

• Changed Fusion from inheritance to composition pattern with IrContainer

• Added back-reference from IrContainer to owning Fusion via fusion() method

• Fixed virtual dispatch for registerExpr by adding Fusion::registerStmt()

• Updated IrBuilder::registerWithContainer() to route through owning Fusion

• All IrContainer methods now forwarded through Fusion composition

Fusion* Statement::fusion() const {
  Fusion* fusion = ir_container_->fusion();
  NVF_ERROR(fusion != nullptr, "Statement does not belong to a fusion.");
  return fusion;
}

kir::Kernel* Statement::kernel() const {
  Fusion* fusion = ir_container_->fusion();
  NVF_ERROR(
      fusion != nullptr && fusion->isA<kir::Kernel>(),
      "Statement does not belong to a kernel.");
  return fusion->as<kir::Kernel>();
}

// the Fusion destructor and the scope initializer could call the
// constructor of Trace, which could throw an exception.
// FUSER_PERF_SCOPE("Fusion clear");

if (container_) {
  container_->clear();
}

inputs_.clear();
outputs_.clear();

io_alias_.clear();

Fusion* fusion() const {
  return owning_fusion_;
}

// Set the owning fusion - should only be called by Fusion
void setOwningFusion(Fusion* fusion) {
  owning_fusion_ = fusion;
}

Greptile Overview

Greptile Summary

This PR refactors Fusion from an inheritance-based design to a composition pattern. Previously Fusion inherited from IrContainer; now it contains an IrContainer* member and forwards container methods. Similarly, IrContainer no longer inherits from PolymorphicBase, with Fusion maintaining that inheritance instead.

Key changes:

• Fusion now owns an IrContainer via std::unique_ptr<IrContainer> container_ with forwarding methods
• IrContainer adds owning_fusion_ back-reference to support bidirectional navigation
• Type checking pattern changed: container()->isA<Kernel>() becomes container()->fusion()->isA<Kernel>()
• Added updateAllStatementContainerPointers() to fix dangling pointers after container swaps
• Registration now routes through Fusion::registerStmt() when container is owned by a Fusion

Issue identified:
The type checking logic in csrc/ir/base_nodes.cpp:82 appears incorrect. The code checks fusion->isA<kir::Kernel>() on a Fusion*, but since Kernel inherits from Fusion (not from IrContainer), this check should work correctly. However, the previous review comment suggests this is problematic - needs verification of the inheritance hierarchy.

Confidence Score: 3/5

• This PR requires careful testing due to fundamental architectural changes in type checking and container ownership
• Score reflects the complexity of the refactoring and potential issues with type checking logic. The composition pattern is well-implemented with proper back-references and pointer updates, but the type checking changes (especially fusion()->isA<Kernel>() pattern) need thorough validation across all edge cases
• csrc/ir/base_nodes.cpp requires close attention for the type checking logic changes, and csrc/fusion.cpp needs validation of the swap and copy operations with container pointer updates

Important Files Changed

File Analysis

Sequence Diagram

sequenceDiagram
    participant User
    participant Fusion
    participant IrContainer
    participant Statement
    participant IrBuilder
    
    Note over Fusion,IrContainer: Composition Pattern
    
    User->>Fusion: new Fusion()
    Fusion->>IrContainer: create unique_ptr<IrContainer>
    IrContainer-->>Fusion: container_
    Fusion->>IrContainer: setOwningFusion(this)
    IrContainer->>IrContainer: owning_fusion_ = this
    
    Note over User,Statement: Creating IR Nodes
    
    User->>IrBuilder: create<Val>(args...)
    IrBuilder->>Fusion: getActiveContainer()
    Fusion-->>IrBuilder: container()
    IrBuilder->>Val: new Val(passkey, args)
    Val->>Statement: ir_container_ = container
    IrBuilder->>IrBuilder: registerWithContainer(container, val)
    IrBuilder->>IrContainer: fusion()
    IrContainer-->>IrBuilder: owning_fusion_
    alt Container owned by Fusion
        IrBuilder->>Fusion: registerStmt(val)
        Fusion->>IrContainer: registerVal(passkey, val)
    else Standalone container
        IrBuilder->>IrContainer: registerStmt(passkey, val)
    end
    
    Note over User,Statement: Type Checking Pattern
    
    User->>Statement: kernel()
    Statement->>IrContainer: fusion()
    IrContainer-->>Statement: owning_fusion_
    Statement->>Fusion: isA<kir::Kernel>()
    Fusion-->>Statement: true/false
    Statement->>Fusion: as<kir::Kernel>()
    Fusion-->>Statement: kernel ptr
    
    Note over Fusion,IrContainer: Swap Operations
    
    User->>Fusion: swap(fusion_a, fusion_b)
    Fusion->>Fusion: swap(container_, container_)
    Fusion->>IrContainer: setOwningFusion(&fusion_a)
    Fusion->>IrContainer: updateAllStatementContainerPointers()
    loop For each Val/Expr
        IrContainer->>Statement: ir_container_ = this
    end