Fusion val ownership by mdavis36 · Pull Request #5953 · NVIDIA/Fuser

mdavis36 · 2026-02-11T03:39:46Z

Summary

Move per-Fusion vals, axioms, and metadata from IrContainer to Fusion

Motivation : Related to Phase 2 of the IR refactor for shared containers these "special values" need to be owned and maintained at the fusion level. After shared scalars are implemented for containers we can revisit the ownership of special vals to allow them to be shared across containers. For now this maintains the current semantic relationship between special vals and Fusions for the remainder of this work.

Changes

Migrate special cached vals (zeroVal, oneVal, trueVal, falseVal, magicZeroVal), axioms, and metadata ownership from IrContainer to Fusion
IrContainer is simplified to only manage raw IR storage (vals, exprs, naming counters)
Fusion-level semantics (constants, CUDA axioms, tensor metadata) now live where they logically belong
Copy, swap, and clear operations updated to maintain correct lifecycle for the moved state

github-actions · 2026-02-11T03:44:52Z

Review updated until commit 946afc2

Description

Moved special cached values (zeroVal, oneVal, trueVal, falseVal, magicZeroVal) from IrContainer to Fusion class
Transferred ownership of axioms and tensor metadata from IrContainer to Fusion
Implemented lazy creation pattern for all special values with unique_ptr ownership
Updated copy, swap, and clear operations to maintain correct lifecycle for moved state
Simplified IrContainer to only manage raw IR storage (vals, exprs, naming counters)

Changes walkthrough

Relevant files

Enhancement

fusion.cpp `Implement per-Fusion special values and metadata management` csrc/fusion.cpp Added includes for type.h, ir/builder.h, and ir/internal_nodes.h Updated Fusion::swap() to handle per-Fusion special values, axioms, and metadata Enhanced Fusion::copy() to properly clone axioms and metadata with IrCloner Modified Fusion::clear() to reset special value pointers and clear axioms/metadata Implemented lazy creation methods for all special values (zeroVal, oneVal, falseVal, trueVal, magicZeroVal) Added dtype-specific special value methods (zeroVal(DataType), oneVal(DataType)) Implemented metadataOf() method for tensor metadata management Added axioms() getter with lazy initialization and assumePositive/assumeNonNegative methods Updated removeVal() to protect cached special vals from removal	+161/-0
fusion.h `Add per-Fusion special value and metadata members` csrc/fusion.h Added forward declaration for NamedScalar class Updated shortcut value method declarations to be implemented in Fusion class Changed numVals() method to have default include_shortcuts=true parameter Declared new special value methods: zeroVal(), oneVal(), falseVal(), trueVal(), magicZeroVal() Added dtype-specific special value method declarations Declared metadataOf() method for tensor metadata Added axioms() getter and assumePositive/assumeNonNegative method declarations Added private member variables for per-Fusion special values using unique_ptr Added axioms_ member as unique_ptr to vector of Val* Added metadata_ member as unordered_map for Val* to pair	+25/-41

Refactoring

container.cpp `Remove special values and metadata from IrContainer` csrc/ir/container.cpp Removed special value member variables (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_) Eliminated axioms_ and metadata_ member variables and their handling Updated swap() to remove special values, axioms, and metadata swapping Modified copy() to remove axioms and metadata copying logic Simplified removeVal() by removing special value protection checks Updated clear() to remove axioms reset and metadata clearing Removed all shortcut methods (zeroVal, oneVal, falseVal, trueVal, magicZeroVal, metadataOf) Eliminated axioms management methods (lazyInitAxioms, assumePositive, assumeNonNegative)	+2/-144
container.h `Remove special value and metadata declarations from IrContainer` csrc/ir/container.h Removed all shortcut method declarations (zeroVal, oneVal, falseVal, trueVal, magicZeroVal, metadataOf) Eliminated axioms management method declarations (axioms, assumePositive, assumeNonNegative, lazyInitAxioms) Updated numVals() comment to clarify special vals now owned by Fusion Removed all special value member variables (true_val_, false_val_, one_val_, zero_val_, magic_zero_val_) Removed axioms_ and metadata_ member variables	+2/-38

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 No relevant tests
⚡ Recommended focus areas for review
Memory ownership transfer In the lazy initialization methods (zeroVal, oneVal, etc.), there's a unique_ptr ownership transfer pattern where the Val is moved from ir_container()->vals_up_ to the Fusion member variable. This pattern should be carefully reviewed to ensure proper memory management and that the transferred Val is not accidentally double-deleted or leaked. auto val = IrBuilder::createInContainer<Val>(this, 0L, DataType::Index); NVF_ERROR(ir_container()->vals_up_.back().get() == val); zero_val_ = std::unique_ptr<Val>(ir_container()->vals_up_.back().release()); ir_container()->vals_up_.pop_back(); } Lazy initialization thread safety The lazy initialization methods for special values (zeroVal, oneVal, etc.) and axioms() are not thread-safe. If multiple threads could access the same Fusion concurrently, there could be race conditions in creating these cached values. This should be reviewed for thread safety requirements. 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_; Clear operation completeness* In the clear() method, special values are reset but there's a comment suggesting they were already absent from vals_up_. The logic for ensuring these values are properly cleaned up from both vals_ and vals_up_ should be verified to prevent memory leaks or dangling pointers. // These unique_ptrs own the Val objects; ir_container()->clear() above only // removed them from vals_ (they were already absent from vals_up_). zero_val_.reset(); one_val_.reset(); true_val_.reset(); false_val_.reset(); magic_zero_val_.reset(); axioms_.reset(); metadata_.clear();

mdavis36 · 2026-02-11T03:53:31Z

!test

mdavis36 · 2026-02-11T04:57:06Z

!test

greptile-apps · 2026-02-11T05:45:30Z

Greptile Overview

Greptile Summary

Migrated ownership of special cached vals (zeroVal, oneVal, trueVal, falseVal, magicZeroVal), CUDA axioms, and tensor metadata from IrContainer to Fusion, aligning with Phase 2 of the IR refactor for shared containers.

Key changes:

IrContainer simplified to pure IR storage (vals, exprs, naming counters)
Fusion now owns fusion-level semantics (constants, axioms, metadata)
Special vals remain in IrContainer::vals_ but ownership transferred to Fusion unique_ptrs
Special vals excluded from vals_up_, so not cloned during copy (recreated lazily instead)
Lifecycle operations (swap, copy, clear, removeVal) updated for new ownership model

Architecture:

Special vals created via IrBuilder::createInContainer(), then ownership transferred by .release() from vals_up_
Axioms and metadata properly cloned in Fusion::copy() using IrCloner
Lazy initialization pattern maintained for all special vals and axioms

Confidence Score: 4/5

This PR is safe to merge with careful testing of copy/clone operations
The ownership migration is architecturally sound and follows a clear refactoring pattern. The implementation correctly handles the lifecycle of special vals through swap, copy, clear, and removeVal operations. One minor issue: an outdated comment in container.cpp should be updated to reflect the new ownership model. The main risk area is ensuring all copy/clone paths properly handle the migrated state, particularly around axioms and metadata cloning.
Pay close attention to csrc/fusion.cpp and csrc/ir/container.cpp - verify copy/clone operations handle special vals correctly in integration tests

Important Files Changed

Filename	Overview
csrc/ir/container.h	Removed special vals, axioms, and metadata ownership - now pure IR storage container
csrc/ir/container.cpp	Removed special val creation/management logic and swap/copy/clear operations for migrated state
csrc/fusion.h	Added special vals, axioms, and metadata as member variables with updated method signatures
csrc/fusion.cpp	Implemented special val creation, axioms, metadata management, and lifecycle operations (swap/copy/clear)

Sequence Diagram

sequenceDiagram
    participant User
    participant Fusion
    participant IrContainer
    participant SpecialVals as Special Vals<br/>(zero_val_, etc.)
    
    Note over Fusion,SpecialVals: Before: IrContainer owned special vals
    Note over Fusion,SpecialVals: After: Fusion owns special vals
    
    User->>Fusion: zeroVal()
    alt Special val not cached
        Fusion->>IrContainer: IrBuilder::createInContainer()
        IrContainer->>IrContainer: Add to vals_up_ & vals_
        IrContainer-->>Fusion: Return new Val*
        Fusion->>IrContainer: vals_up_.back().release()
        Fusion->>IrContainer: vals_up_.pop_back()
        Fusion->>SpecialVals: Store in zero_val_ unique_ptr
        Note over Fusion,IrContainer: Val remains in vals_ but<br/>ownership moved to Fusion
    end
    Fusion-->>User: Return zero_val_.get()
    
    User->>Fusion: copy(from, to)
    Fusion->>IrContainer: IrContainer::copy()
    Note over IrContainer: deterministic_vals() only<br/>returns vals from vals_up_<br/>(excludes special vals)
    IrContainer-->>Fusion: IrCloner
    Fusion->>Fusion: Clone axioms_ vector
    Fusion->>Fusion: Clone metadata_ map
    Note over Fusion,SpecialVals: Special vals NOT copied,<br/>recreated lazily on access

greptile-apps

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

Moved special values (`zero_val_`, `one_val_`, `true_val_`, `false_val_`, `magic_zero_val_`) from `IrContainer` to the `Fusion` class. This ensures that with shared containers, each Fusion has its own special values, preventing ownership conflicts when one Fusion is destroyed. **Option Implemented:** Option A (Move Special Values to Fusion) as recommended in the prompt. Added private members and public accessors to Fusion class: ```cpp // Phase 2: Per-Fusion special values // With shared containers, each Fusion needs its own special values. // These are raw pointers - memory is owned by IrContainer's vals_up_. // Destroying this Fusion removes these vals via removeStatementsOwnedBy(). Val* zero_val_ = nullptr; Val* one_val_ = nullptr; Val* true_val_ = nullptr; Val* false_val_ = nullptr; NamedScalar* magic_zero_val_ = nullptr; ``` Public accessors: - `Val* zeroVal()` - Returns Index 0 - `Val* oneVal()` - Returns Index 1 - `Val* falseVal()` - Returns Bool false - `Val* trueVal()` - Returns Bool true - `NamedScalar* magicZeroVal()` - Returns magic zero named scalar - `Val* zeroVal(DataType dtype)` - Returns 0 for specified dtype - `Val* oneVal(DataType dtype)` - Returns 1 for specified dtype Implemented lazy creation pattern for all special value accessors: ```cpp Val* Fusion::zeroVal() { if (!zero_val_) { zero_val_ = IrBuilder::createInContainer<Val>(this, 0L, DataType::Index); } return zero_val_; } // Similar implementations for oneVal(), falseVal(), trueVal(), magicZeroVal() ``` Updated `Fusion::clear()` to reset special value pointers: ```cpp // Reset per-Fusion special values (they'll be recreated lazily if needed) // The actual Val objects were removed by removeStatementsOwnedBy above. zero_val_ = nullptr; one_val_ = nullptr; true_val_ = nullptr; false_val_ = nullptr; magic_zero_val_ = nullptr; ``` Removed special value members and added documentation comment: ```cpp // Note: Special values (zero_val_, one_val_, true_val_, false_val_, // magic_zero_val_) are now per-Fusion, stored in Fusion class. // This avoids ownership conflicts when multiple Fusions share an IrContainer. // See Fusion::zeroVal(), etc. for the per-Fusion implementation. ``` Removed special value accessor implementations (they're now in Fusion). All call sites were already updated to use `fusion->zeroVal()` instead of `ir_container()->zeroVal()`. Verified with grep that no call sites remain using the old pattern. Added 8 new unit tests for Task 7: 1. **PerFusionSpecialValuesBasic** - Tests that special values are created and owned by the Fusion 2. **SpecialValuesOwnedByFusion** - Tests that special values are tracked in `ownedVals()` 3. **SeparateFusionsHaveOwnSpecialValues** - Tests that two Fusions have different special value objects 4. **DestroyFusionDoesNotAffectOther** - Tests that destroying one Fusion doesn't affect another's special values 5. **SpecialValuesLazyCreation** - Tests that same value is returned on repeated calls 6. **AllSpecialValuesPerFusion** - Tests all five special value accessors 7. **SpecialValuesClearedOnFusionClear** - Tests that `clear()` resets special values 8. **SpecialValuesWithDtype** - Tests `zeroVal(dtype)` and `oneVal(dtype)` accessors ``` [==========] Running 34 tests from 3 test suites. [ PASSED ] 34 tests. ``` ``` [==========] Running 26 tests from 1 test suite. [ PASSED ] 26 tests. ``` Including 8 new Task 7 tests: - `Phase2ContainerTest.PerFusionSpecialValuesBasic` - PASSED - `Phase2ContainerTest.SpecialValuesOwnedByFusion` - PASSED - `Phase2ContainerTest.SeparateFusionsHaveOwnSpecialValues` - PASSED - `Phase2ContainerTest.DestroyFusionDoesNotAffectOther` - PASSED - `Phase2ContainerTest.SpecialValuesLazyCreation` - PASSED - `Phase2ContainerTest.AllSpecialValuesPerFusion` - PASSED - `Phase2ContainerTest.SpecialValuesClearedOnFusionClear` - PASSED - `Phase2ContainerTest.SpecialValuesWithDtype` - PASSED - `csrc/fusion.h` - Added special value members and accessors - `csrc/fusion.cpp` - Added accessor implementations, updated `clear()` - `csrc/ir/container.h` - Removed special values, added comment - `csrc/ir/container.cpp` - Removed accessor implementations - `tests/cpp/test_phase2_container_sharing.cpp` - Added 8 unit tests - [x] Each Fusion has its own special values - [x] Destroying Fusion A doesn't affect Fusion B's special values - [x] Special value accessors (`zeroVal()`, `oneVal()`, etc.) return this Fusion's values - [x] Lazy creation still works (create on first access) - [x] Smoke tests pass (34/34) - [x] Unit tests added (8 tests) - [x] Unit tests pass (26/26 Phase 2 tests) - [x] Code compiles without errors - [x] REPORT.md delivered 1. **Memory ownership:** Special values are raw pointers stored in Fusion, but the actual memory is owned by IrContainer's `vals_up_`. When a Fusion is destroyed, `removeStatementsOwnedBy()` cleans up these vals. 2. **Lazy creation pattern:** Special values are created on first access. This matches the original IrContainer behavior and avoids creating values that aren't needed. 3. **Clear handling:** `Fusion::clear()` now resets special value pointers to nullptr after `removeStatementsOwnedBy()` removes the actual Val objects. This ensures lazy recreation works correctly after clear. 4. **Copy/move handling:** Will be addressed in Tasks 5 and 6. This task just moves the members and accessors.

Moved `axioms_` and `metadata_` from `IrContainer` to the `Fusion` class. This completes the deprecation of `parent_` usage for val-creating methods, which was necessary because `parent_` implies a 1-1 relationship (container → Fusion), but Phase 2 has 1-many (shared containers). Methods that used `parent_` to create vals were moved to Fusion: - `metadataOf(Val*)` - Now uses `v->container()` to get owning Fusion - `axioms()` - Now creates axiom vals owned by `this` Fusion - `assumePositive/assumeNonNegative` - Now adds to `this` Fusion's axioms - Added `axioms_` and `metadata_` private members - Changed method declarations from forwarding to actual implementations - Added includes for `ir/builder.h` and `ir/internal_nodes.h` - Implemented `metadataOf()`, `axioms()`, `assumePositive()`, `assumeNonNegative()` methods - Updated `clear()` to reset `axioms_` and `metadata_` - Removed `metadataOf()`, `axioms()`, `assumePositive()`, `assumeNonNegative()` declarations - Removed `lazyInitAxioms()` declaration - Removed `axioms_` and `metadata_` members - Removed implementations of above methods - Updated `IrContainer::swap` to remove axioms_/metadata_ swapping - Updated `IrContainer::copy` to remove axioms_/metadata_ handling - Updated `IrContainer::clear` to remove axioms_/metadata_ clearing Each Fusion now has its own axioms and metadata cache. This ensures: 1. No ownership conflicts when multiple Fusions share an IrContainer 2. Correct behavior when one Fusion is destroyed (doesn't affect others) 3. Lazy creation pattern preserved (create on first access) This is a prerequisite for the copy/move semantics changes which will swap/transfer these per-Fusion members.

- Add missing swap of axioms_ and metadata_ in Fusion::swap to prevent dangling pointers after move/assignment - Add missing cloning of axioms_ and metadata_ in Fusion::copy to preserve custom assumptions and metadata cache across copies - Guard Fusion::removeVal against removing cached special vals - Use std::unique_ptr for special vals and steal from vals_up_ to preserve the original invariant (shortcuts in vals_ but not vals_up_) - Fix metadataOf to use 'this' instead of v->container()

greptile-apps

_{4 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-11T20:05:38Z

Additional Comments (1)

csrc/ir/container.cpp
outdated comment after ownership migration to Fusion

  // Copy values in deterministic order.
  // Special vals (zero_val_, one_val_, etc.) are owned by Fusion and not in vals_up_,
  // so they won't appear in deterministic_vals() and will be lazily recreated.

mdavis36 force-pushed the md/fusion-vals-axioms branch from f273b16 to 6cd9dc1 Compare February 11, 2026 03:47

mdavis36 force-pushed the md/fusion-vals-axioms branch from f982c6e to 3502a22 Compare February 11, 2026 04:56

mdavis36 marked this pull request as ready for review February 11, 2026 05:36

greptile-apps bot reviewed Feb 11, 2026

View reviewed changes

mdavis36 added 4 commits February 11, 2026 12:01

Cleanup comments

ab16a88

mdavis36 force-pushed the md/fusion-vals-axioms branch from 3502a22 to 946afc2 Compare February 11, 2026 20:01

greptile-apps bot reviewed Feb 11, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fusion val ownership#5953

Fusion val ownership#5953
mdavis36 wants to merge 4 commits intomainfrom
md/fusion-vals-axioms

mdavis36 commented Feb 11, 2026 •

edited

Loading

Uh oh!

github-actions bot commented Feb 11, 2026 •

edited

Loading

Changes walkthrough

PR Reviewer Guide

Uh oh!

mdavis36 commented Feb 11, 2026

Uh oh!

mdavis36 commented Feb 11, 2026

Uh oh!

greptile-apps bot commented Feb 11, 2026 •

edited

Loading

Uh oh!

greptile-apps bot left a comment

Uh oh!

greptile-apps bot left a comment

Uh oh!

greptile-apps bot commented Feb 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

mdavis36 commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Uh oh!

github-actions bot commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes walkthrough

PR Reviewer Guide

Uh oh!

mdavis36 commented Feb 11, 2026

Uh oh!

mdavis36 commented Feb 11, 2026

Uh oh!

greptile-apps bot commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot commented Feb 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

mdavis36 commented Feb 11, 2026 •

edited

Loading

github-actions bot commented Feb 11, 2026 •

edited

Loading

greptile-apps bot commented Feb 11, 2026 •

edited

Loading