Per-Fusion Axioms and Metadata

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.

Author: mdavis36

csrc/fusion.cpp

@@ -20,7 +20,9 @@
 #include <host_ir/container.h>
 #include <instrumentation.h>
 #include <ir/all_nodes.h>
+#include <ir/builder.h>
 #include <ir/cloner.h>
+#include <ir/internal_nodes.h>
 #include <ir/printer.h>
 #include <ir/utils.h>
 #include <iter_visitor.h>
@@ -280,6 +282,10 @@ void Fusion::clear() noexcept {
   false_val_ = nullptr;
   magic_zero_val_ = nullptr;
 
+  // Reset per-Fusion axioms and metadata (Phase 2)
+  axioms_.reset();
+  metadata_.clear();
+
   invalidateTvsAndUses();
 
   is_during_update_uses_ = false;
@@ -768,6 +774,54 @@ Val* Fusion::oneVal(DataType dtype) {
   }
 }
 
+// =========================================================================
+// Per-Fusion Metadata and Axioms (Phase 2)
+// These are per-Fusion to avoid ownership issues with shared containers.
+// =========================================================================
+
+Val* Fusion::metadataOf(Val* v) {
+  if (metadata_.count(v) == 0) {
+    // Create metadata val owned by the same Fusion as v
+    Fusion* owner = v->container();
+    auto metadata_val =
+        IrBuilder::createInContainer<Val>(owner, metaDataTypeOf(v));
+    auto metadata_expr =
+        IrBuilder::createInContainer<GetMetaData>(owner, metadata_val, v);
+    metadata_[v] = std::make_pair(metadata_val, metadata_expr);
+  }
+  return metadata_.at(v).first;
+}
+
+const std::vector<Val*>& Fusion::axioms() {
+  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_;
+}
+
+void Fusion::assumePositive(Val* val) {
+  NVF_ERROR(inContainer(val));
+  // Lazy init axioms, then add the assumption
+  axioms();
+  axioms_->emplace_back(IrBuilder::gtExpr(val, zeroVal()));
+}
+
+void Fusion::assumeNonNegative(Val* val) {
+  NVF_ERROR(inContainer(val));
+  // Lazy init axioms, then add the assumption
+  axioms();
+  axioms_->emplace_back(IrBuilder::geExpr(val, zeroVal()));
+}
+
 void Fusion::registerVal(Val* val) {
   if (inContainer(val)) {
     return;

csrc/fusion.h

@@ -566,22 +566,15 @@ class NVF_API Fusion : public PolymorphicBase {
   Val* zeroVal(DataType dtype);
   Val* oneVal(DataType dtype);
 
-  Val* metadataOf(Val* val) {
-    return ir_container()->metadataOf(val);
-  }
+  // Phase 2: Per-Fusion metadata and axioms
+  // These are now per-Fusion to avoid ownership issues with shared containers.
+  Val* metadataOf(Val* val);
 
   // Axioms (CUDA programming assumptions)
-  const std::vector<Val*>& axioms() {
-    return ir_container()->axioms();
-  }
+  const std::vector<Val*>& axioms();
 
-  void assumePositive(Val* val) {
-    ir_container()->assumePositive(val);
-  }
-
-  void assumeNonNegative(Val* val) {
-    ir_container()->assumeNonNegative(val);
-  }
+  void assumePositive(Val* val);
+  void assumeNonNegative(Val* val);
 
   // Statement removal
   void removeStatementsCreatedAfter(
@@ -661,6 +654,14 @@ class NVF_API Fusion : public PolymorphicBase {
   Val* true_val_ = nullptr;
   Val* false_val_ = nullptr;
   NamedScalar* magic_zero_val_ = nullptr;
+
+  // Phase 2: Per-Fusion axioms (CUDA programming assumptions)
+  // These are per-Fusion to avoid ownership issues with shared containers.
+  std::unique_ptr<std::vector<Val*>> axioms_;
+
+  // Phase 2: Per-Fusion metadata cache
+  // Maps Val* to (metadata_val, metadata_expr) pairs
+  std::unordered_map<Val*, std::pair<Val*, Expr*>> metadata_;
 };
 
 // Template implementations for Fusion::manage<T>() that use IrCloner

csrc/ir/container.cpp

@@ -81,17 +81,15 @@ void IrContainer::swap(IrContainer& a, IrContainer& b) noexcept {
   std::swap(a.val_type_name_map_, b.val_type_name_map_);
   std::swap(a.expr_name_counter_, b.expr_name_counter_);
 
-  std::swap(a.metadata_, b.metadata_);
-
   std::swap(a.parent_, b.parent_);
 
-  // Note: Special values (zero_val_, one_val_, etc.) are now per-Fusion,
-  // not per-IrContainer. They are swapped as part of the Fusion-level swap.
-  std::swap(a.axioms_, b.axioms_);
+  // Note: Special values, axioms, and metadata are now per-Fusion,
+  // not per-IrContainer. They are handled by Fusion::swap.
 }
 
 IrCloner IrContainer::copy(const IrContainer* from, IrContainer* to) {
   to->clear();
+
   IrCloner ir_cloner(to->parent());
 
   // Copy values in deterministic order
@@ -113,14 +111,7 @@ IrCloner IrContainer::copy(const IrContainer* from, IrContainer* to) {
   to->val_type_name_map_ = from->val_type_name_map_;
   to->expr_name_counter_ = from->expr_name_counter_;
 
-  if (from->axioms_ != nullptr) {
-    to->axioms_ = std::make_unique<std::vector<Val*>>();
-    for (auto pred : *from->axioms_) {
-      to->axioms_->push_back(ir_cloner.clone(pred));
-    }
-  }
-
-  to->metadata_ = ir_cloner.clone(from->metadata_);
+  // Note: axioms and metadata are now per-Fusion, handled by Fusion::copy
 
   return ir_cloner;
 }
@@ -201,9 +192,7 @@ void IrContainer::clear() noexcept {
   vals_up_.clear();
   exprs_.clear();
   exprs_up_.clear();
-  axioms_.reset();
   val_type_name_map_.clear();
-  metadata_.clear();
   expr_name_counter_ = 0;
 }
 
@@ -239,51 +228,11 @@ bool IrContainer::inContainer(const Statement* const_stmt) const {
   return true;
 }
 
-// Note: Shortcut values (zeroVal, oneVal, trueVal, falseVal, magicZeroVal)
-// are now per-Fusion. Use Fusion::zeroVal() etc. instead.
+// Note: Shortcut values (zeroVal, oneVal, trueVal, falseVal, magicZeroVal),
+// metadata, and axioms are now per-Fusion. Use Fusion::zeroVal(),
+// Fusion::metadataOf(), Fusion::axioms(), etc. instead.
 // This avoids ownership conflicts when multiple Fusions share an IrContainer.
 
-Val* IrContainer::metadataOf(Val* v) {
-  if (metadata_.count(v) == 0) {
-    auto metadata_val =
-        IrBuilder::createInContainer<Val>(this->parent(), metaDataTypeOf(v));
-    auto metadata_expr = IrBuilder::createInContainer<GetMetaData>(
-        this->parent(), metadata_val, v);
-    metadata_[v] = std::make_pair(metadata_val, metadata_expr);
-  }
-  return metadata_.at(v).first;
-}
-
-void IrContainer::lazyInitAxioms() {
-  if (!axioms_) {
-    axioms_ = std::make_unique<std::vector<Val*>>();
-    axioms_->reserve(kParallelTypeThreads.size() * 3);
-    // Use parent()->zeroVal() since special values are now per-Fusion
-    auto zero = parent()->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));
-    }
-  }
-}
-
-void IrContainer::assumePositive(Val* val) {
-  NVF_ERROR(val->container() == this->parent());
-  lazyInitAxioms();
-  // Use parent()->zeroVal() since special values are now per-Fusion
-  axioms_->emplace_back(IrBuilder::gtExpr(val, parent()->zeroVal()));
-}
-
-void IrContainer::assumeNonNegative(Val* val) {
-  NVF_ERROR(val->container() == this->parent());
-  lazyInitAxioms();
-  // Use parent()->zeroVal() since special values are now per-Fusion
-  axioms_->emplace_back(IrBuilder::geExpr(val, parent()->zeroVal()));
-}
-
 void IrContainer::removeStatementsCreatedAfter(
     int64_t prev_num_exprs,
     int64_t prev_num_vals) {

csrc/ir/container.h

@@ -88,21 +88,11 @@ class IrContainer {
     return include_shortcuts ? std::ssize(vals_) : std::ssize(vals_up_);
   }
 
-  // Note: Shortcut values (zeroVal, oneVal, trueVal, falseVal, magicZeroVal)
-  // are now per-Fusion. Use Fusion::zeroVal() etc. instead.
+  // Note: Shortcut values (zeroVal, oneVal, trueVal, falseVal, magicZeroVal),
+  // metadata, and axioms are now per-Fusion. Use Fusion::zeroVal(),
+  // Fusion::metadataOf(), Fusion::axioms(), etc. instead.
   // This avoids ownership conflicts when multiple Fusions share an IrContainer.
 
-  Val* metadataOf(Val*);
-
-  // Axioms about CUDA programming, for example: threadIdx.x < blockDim.x
-  const std::vector<Val*>& axioms() {
-    lazyInitAxioms();
-    return *axioms_;
-  }
-
-  void assumePositive(Val* val);
-  void assumeNonNegative(Val* val);
-
  protected:
   static IrCloner copy(const IrContainer* from, IrContainer* to);
 
@@ -136,8 +126,6 @@ class IrContainer {
 
   void clear() noexcept;
 
-  void lazyInitAxioms();
-
   friend class StatementGuard;
 
   // A simple garbage collection mechanism to remove all Exprs and Vals that
@@ -173,10 +161,8 @@ class IrContainer {
   // 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.
-
-  std::unique_ptr<std::vector<Val*>> axioms_;
-  std::unordered_map<Val*, std::pair<Val*, Expr*>> metadata_;
+  // See Fusion::zeroVal(), Fusion::axioms(), Fusion::metadataOf(), etc.
+  // for the per-Fusion implementations.
 
  public:
   Fusion* parent() const {