[IR Refactor] Fusion Base Type

[mdavis36]

No description provided.

[mdavis36]

!test

[github-actions]

Review updated until commit 603d6a8

Description

• Refactor Fusion to use composition over inheritance with IrContainer

• Move IrContainer functionality into storage.h and rename to IrContainer

• Add API forwarding methods in Fusion to delegate to contained IrContainer

• Update constructors and copy/move operations for new design

• Move template implementations to avoid circular dependencies

Changes walkthrough

	Relevant files
Enhancement	15 files fusion.cpp Update Fusion implementation for composition-based design +45/-10  cloner.cpp Update IrCloner to work with Fusion instead of IrContainer +1/-1      container.cpp Remove IrContainer implementation file (functionality moved) +0/-57    storage.cpp Move IrContainer functionality into storage.cpp                    +27/-27  utils.cpp Update removeExpr signature for new design                              +1/-1      mutator.cpp Update removeExpr signature for new design                              +1/-1      dispatch.h Update removeExpr signature for new design                              +1/-1      fusion.h Major refactor: Fusion now contains IrContainer instead of inheriting +176/-16 base_nodes.h Update container pointers and method signatures                    +6/-12    builder.h Update createInContainer template signature                            +4/-9      builder_passkey.h Update passkey to use Fusion instead of IrContainer            +4/-4      cloner.h Update IrCloner to work with Fusion and move template implementation +6/-32    container.h Remove IrContainer implementation (now empty)                        +1/-226  storage.h Move IrContainer functionality into storage.h                        +21/-16  kernel.h Remove using declarations for inherited methods                    +0/-3
Configuration changes	1 files CMakeLists.txt Comment out container.cpp build (functionality moved)        +1/-1

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Complex swap logic The swap function now has extensive parent pointer update logic (lines 115-137) that manually updates ir_container_ pointers for all statements after swapping. This is complex and error-prone. The comment mentions casting safety assumptions that may not be enforced by the type system. if (a.ir_container_) { // Also update all Statement ir_container_ pointers to point to new owner // Note: IrContainer is now in impl namespace, but Statement::ir_container_ // is Fusion*. Since only Fusion (and its derived classes) inherit from // impl::IrContainer, this cast is safe. a.ir_container()->parent_ = &a; for (auto val : a.vals()) { val->ir_container_ = &a; } for (auto expr : a.deterministic_exprs()) { expr->ir_container_ = &a; } } if (b.ir_container_) { // Also update all Statement ir_container_ pointers to point to new owner b.ir_container()->parent_ = &b; for (auto val : b.vals()) { val->ir_container_ = &b; } for (auto expr : b.deterministic_exprs()) { expr->ir_container_ = &b; } } TODO comment indicates uncertainty There's a "TODO : CHECK THIS" comment in removeExpr function, suggesting the implementer was uncertain about the correctness of the change from IrContainer::removeExpr(expr) to ir_container()->removeExpr(expr). // TODO : CHECK THIS vvv ir_container()->removeExpr(expr); } Design complexity The refactor changes from inheritance (Fusion : public IrContainer) to composition (Fusion contains IrContainer*). This adds complexity with forwarding methods and parent pointer management. The benefits of this change should be clearly documented as it significantly increases code complexity. class IrContainer { public: NVF_API IrContainer(); // Copy/Move Constructors and Operators are deleted. IrContainer is managed // through a smart pointer in IrContainer. Semantic operations for Fusion // types are handled directly through copy and swap functions. IrContainer(const IrContainer& other) = delete; IrContainer(IrContainer&& other) noexcept = delete; IrContainer& operator=(const IrContainer& other) = delete; IrContainer& operator=(IrContainer&& other) noexcept = delete; ~IrContainer();

Test failures

• (Medium, 2) Thunder nvFuser NanoGPT output is zero (scalar mismatch) in test_networks on CUDA (A100 & H100)

  Test Name	A100	H100	Source
  thunder.tests.test_networks.test_nanogpt_complete_autograd_nvfuser_cuda_thunder.dtypes.float32	❌	❌

[mdavis36]

!test