Per-Fusion name counters for shared container name correspondence

Replace global IrContainer name counters with per-Fusion counters so
cloned Fusions produce matching statement names (T0=T0, T1=T1) instead
of incrementing names (T0=T10). This fixes cross-fusion name lookups
in GreedyParams and normalization_utils which use tv->name() as map keys.

Changes:
- Add per_fusion_val_name_map_ and per_fusion_expr_name_counter_ to IrContainer
- Update getValName/getExprName to use per-Fusion counter with global fallback
- Update registerVal/registerExpr to pass owning Fusion to name generators
- Handle counter lifecycle in swap, copy, clear, destroy, transferOwnership
- Use deterministic_vals() in Fusion::copy for stable clone ordering
- Add 8 new tests for name correspondence (71/71 Phase 2 tests pass)

Author: mdavis36

csrc/fusion.cpp

@@ -218,13 +218,16 @@ IrCloner Fusion::copy(const Fusion* from, Fusion* to) {
   IrCloner ir_cloner(to);
 
   // Phase 2: Clone only 'from's owned vals (not all vals in shared container)
-  // Use ownedVals() to get only vals belonging to 'from'
-  for (auto val : from->ownedVals()) {
+  // CRITICAL: Use deterministic_vals() to get vals in insertion order.
+  // Using ownedVals() (unordered_set) causes non-deterministic clone order,
+  // which assigns different name() values to cloned vals between runs.
+  // This breaks code that uses tv->name() as map keys (e.g., GreedyParams).
+  for (auto val : from->deterministic_vals()) {
     ir_cloner.clone(val);
   }
 
   // Update definition_ and uses_ on cloned vals
-  for (auto val : from->ownedVals()) {
+  for (auto val : from->deterministic_vals()) {
     ir_cloner.clone(val)->setDefinition(ir_cloner.clone(val->definition_));
     ir_cloner.clone(val)->setUses(ir_cloner.clone(val->uses_));
   }

csrc/fusion.h

@@ -744,7 +744,13 @@ T* IrBuilder::clone(const T* src, IrCloner* ir_cloner) {
 
   dest_container->registerStmt(IrBuilderPasskey(dest_container), dest_stmt);
 
-  if (src_container != dest_container) {
+  // Phase 2 Task 10: For same-container cloning (shared IrContainer),
+  // per-Fusion name counters produce matching names naturally (both start
+  // at 0), so the name override below is NOT needed and is skipped.
+  // For cross-container cloning (different IrContainers), we still need
+  // to force the source name since the destination's global counter may
+  // have diverged.
+  if (src_container->ir_container() != dest_container->ir_container()) {
     dest_stmt->setName(IrBuilderPasskey(dest_container), src_stmt->name());
   }
 

csrc/ir/container.cpp

@@ -268,6 +268,24 @@ void IrContainer::transferStatementOwnership(Fusion* from, Fusion* to) {
     to_exprs.insert(exprs_it->second.begin(), exprs_it->second.end());
     per_fusion_exprs_.erase(exprs_it);
   }
+
+  // Transfer per-Fusion name counters (Phase 2 Task 10)
+  auto val_names_it = per_fusion_val_name_map_.find(from);
+  if (val_names_it != per_fusion_val_name_map_.end()) {
+    // Merge counter maps: take max of each ValType counter
+    auto& to_map = per_fusion_val_name_map_[to];
+    for (auto& [vtype, counter] : val_names_it->second) {
+      to_map[vtype] = std::max(to_map[vtype], counter);
+    }
+    per_fusion_val_name_map_.erase(val_names_it);
+  }
+
+  auto expr_names_it = per_fusion_expr_name_counter_.find(from);
+  if (expr_names_it != per_fusion_expr_name_counter_.end()) {
+    auto& to_counter = per_fusion_expr_name_counter_[to];
+    to_counter = std::max(to_counter, expr_names_it->second);
+    per_fusion_expr_name_counter_.erase(expr_names_it);
+  }
 }
 
 void IrContainer::removeStatementsOwnedBy(Fusion* fusion) {
@@ -303,6 +321,10 @@ void IrContainer::removeStatementsOwnedByUnlocked(Fusion* fusion) {
   // Clean up per-Fusion tracking (Phase 2 Task 4)
   per_fusion_vals_.erase(fusion);
   per_fusion_exprs_.erase(fusion);
+
+  // Clean up per-Fusion name counters (Phase 2 Task 10)
+  per_fusion_val_name_map_.erase(fusion);
+  per_fusion_expr_name_counter_.erase(fusion);
 }
 
 void IrContainer::swap(IrContainer& a, IrContainer& b) noexcept {
@@ -334,6 +356,10 @@ void IrContainer::swap(IrContainer& a, IrContainer& b) noexcept {
   std::swap(a.sharing_fusions_, b.sharing_fusions_);
   std::swap(a.per_fusion_vals_, b.per_fusion_vals_);
   std::swap(a.per_fusion_exprs_, b.per_fusion_exprs_);
+
+  // Swap per-Fusion name counters (Phase 2 Task 10)
+  std::swap(a.per_fusion_val_name_map_, b.per_fusion_val_name_map_);
+  std::swap(a.per_fusion_expr_name_counter_, b.per_fusion_expr_name_counter_);
 }
 
 IrCloner IrContainer::copy(const IrContainer* from, IrContainer* to) {
@@ -354,6 +380,8 @@ IrCloner IrContainer::copy(const IrContainer* from, IrContainer* to) {
   to->expr_name_counter_ = 0;
   to->per_fusion_vals_.clear();
   to->per_fusion_exprs_.clear();
+  to->per_fusion_val_name_map_.clear();
+  to->per_fusion_expr_name_counter_.clear();
 
   // NOTE: In Phase 2, we can't use to->parent() here because parent_ might
   // not be set correctly for shared containers. Fusion::copy handles this.
@@ -468,12 +496,16 @@ void IrContainer::registerVal(Val* val) {
   // Otherwise handle registration locally
   vals_up_.emplace_back(val);
   vals_.insert(val);
-  val->setName(IrContainerPasskey(), getValName(val->vtype()));
+
+  // Phase 2 Task 10: Use per-Fusion counter if val has an owning Fusion.
+  // This ensures cloned Fusions get matching names (T0=T0, T1=T1)
+  // instead of incrementing global names (T0=T10, T1=T11).
+  Fusion* owning_fusion = val->container();
+  val->setName(IrContainerPasskey(), getValName(owning_fusion, val->vtype()));
 
   // Track per-Fusion ownership (Phase 2 Task 4)
-  // val->container() returns the owning Fusion
-  if (val->container() != nullptr) {
-    per_fusion_vals_[val->container()].insert(val);
+  if (owning_fusion != nullptr) {
+    per_fusion_vals_[owning_fusion].insert(val);
   }
 }
 
@@ -486,12 +518,14 @@ void IrContainer::registerExpr(Expr* expr) {
   // Otherwise handle registration locally
   exprs_up_.emplace_back(expr);
   exprs_.insert(expr);
-  expr->setName(IrContainerPasskey(), getExprName());
+
+  // Phase 2 Task 10: Use per-Fusion counter if expr has an owning Fusion.
+  Fusion* owning_fusion = expr->container();
+  expr->setName(IrContainerPasskey(), getExprName(owning_fusion));
 
   // Track per-Fusion ownership (Phase 2 Task 4)
-  // expr->container() returns the owning Fusion
-  if (expr->container() != nullptr) {
-    per_fusion_exprs_[expr->container()].insert(expr);
+  if (owning_fusion != nullptr) {
+    per_fusion_exprs_[owning_fusion].insert(expr);
   }
 }
 
@@ -507,6 +541,10 @@ void IrContainer::clear() noexcept {
   // Clear per-Fusion tracking (Phase 2 Task 4)
   per_fusion_vals_.clear();
   per_fusion_exprs_.clear();
+
+  // Clear per-Fusion name counters (Phase 2 Task 10)
+  per_fusion_val_name_map_.clear();
+  per_fusion_expr_name_counter_.clear();
 }
 
 bool IrContainer::inContainer(const Statement* const_stmt) const {

csrc/ir/container.h

@@ -192,14 +192,43 @@ class IrContainer {
   //! Register expr with this container.
   NVF_API void registerExpr(Expr* expr);
 
-  StmtNameType getValName(ValType vtype) {
+  //! Get next val name, using per-Fusion counter if fusion is non-null,
+  //! falling back to global counter otherwise.
+  //! Per-Fusion counters ensure cloned Fusions produce matching names.
+  StmtNameType getValName(Fusion* fusion, ValType vtype) {
+    if (fusion != nullptr) {
+      auto& name_map = per_fusion_val_name_map_[fusion];
+      if (name_map.find(vtype) == name_map.end()) {
+        name_map[vtype] = 0;
+      }
+      // Also advance global counter to keep it >= all per-Fusion counters
+      // This prevents conflicts if global counter is used later
+      auto& global = val_type_name_map_[vtype];
+      auto per_fusion_name = name_map[vtype]++;
+      if (global <= per_fusion_name) {
+        global = per_fusion_name + 1;
+      }
+      return per_fusion_name;
+    }
+    // Global fallback for non-Fusion contexts
     if (val_type_name_map_.find(vtype) == val_type_name_map_.end()) {
       val_type_name_map_[vtype] = 0;
     }
     return val_type_name_map_[vtype]++;
   }
 
-  StmtNameType getExprName() {
+  //! Get next expr name, using per-Fusion counter if fusion is non-null,
+  //! falling back to global counter otherwise.
+  StmtNameType getExprName(Fusion* fusion) {
+    if (fusion != nullptr) {
+      auto& counter = per_fusion_expr_name_counter_[fusion];
+      auto per_fusion_name = counter++;
+      // Also advance global counter
+      if (expr_name_counter_ <= per_fusion_name) {
+        expr_name_counter_ = per_fusion_name + 1;
+      }
+      return per_fusion_name;
+    }
     return expr_name_counter_++;
   }
 
@@ -237,12 +266,21 @@ class IrContainer {
   // something like check if an Expr is in this container
   std::unordered_set<Expr*> exprs_;
 
-  // Values names counters
+  // Values names counters (global fallback for non-Fusion contexts)
   std::unordered_map<ValType, StmtNameType> val_type_name_map_;
 
-  // Expression names counter
+  // Expression names counter (global fallback for non-Fusion contexts)
   StmtNameType expr_name_counter_ = 0;
 
+  // Per-Fusion name counters (Phase 2 Task 10)
+  // Each Fusion gets its own counter starting at 0, so cloned Fusions
+  // produce matching names (T0=T0, T1=T1) instead of incrementing names.
+  // This is critical for GreedyParams and normalization_utils which use
+  // tv->name() as map keys across cloned Fusions.
+  std::unordered_map<Fusion*, std::unordered_map<ValType, StmtNameType>>
+      per_fusion_val_name_map_;
+  std::unordered_map<Fusion*, StmtNameType> per_fusion_expr_name_counter_;
+
   // 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.