Per-Fusion statement tracking and ownership-filtered accessors

Add per_fusion_vals_ / per_fusion_exprs_ maps to IrContainer so each
Fusion can efficiently query only its own statements in a shared
container. Fusion forwarding methods (vals(), unordered_exprs(),
deterministic_vals(), etc.) now return per-Fusion filtered results.
Fusion::clear() uses removeStatementsOwnedBy(this) instead of
ir_container()->clear().

Author: mdavis36

csrc/fusion.cpp

@@ -111,10 +111,11 @@ void Fusion::swap(Fusion& a, Fusion& b) noexcept {
   // will only swap the ptrs NOT the contents.
   IrContainer::swap(*(a.ir_container()), *(b.ir_container()));
 
-  // Fix parent pointers after swapping containers
-  // After swap, each Fusion owns a different IrContainer, so we must
-  // update the parent backpointers in those containers to point to their new
-  // owners
+  // After swapping container contents, per-Fusion tracking keys point to the
+  // wrong Fusions. Rename: a's container had b's entries, b's had a's.
+  a.ir_container()->transferStatementOwnership(&b, &a);
+  b.ir_container()->transferStatementOwnership(&a, &b);
+
   if (a.ir_container_) {
     for (auto val : a.vals()) {
       val->ir_container_ = &a;
@@ -295,10 +296,9 @@ void Fusion::clear() noexcept {
   // constructor of Trace, which could throw an exception.
   // FUSER_PERF_SCOPE("Fusion clear");
 
-  // Clear container contents instead of destroying it
-  // This preserves the container object so Statement pointers don't become
-  // dangling
-  ir_container()->clear();
+  if (ir_container_) {
+    ir_container_->removeStatementsOwnedBy(this);
+  }
 
   inputs_.clear();
   outputs_.clear();
@@ -308,8 +308,6 @@ void Fusion::clear() noexcept {
   managed_data_.clear();
   managed_named_data_.clear();
 
-  // Reset per-Fusion special value caches (the vals themselves are owned by
-  // ir_container and were already destroyed by ir_container()->clear() above).
   zero_val_ = nullptr;
   one_val_ = nullptr;
   true_val_ = nullptr;
@@ -354,6 +352,7 @@ void Fusion::removeExpr(Expr* expr) {
   NVF_ERROR(
       expr_in_deque != c->exprs_up_.end(),
       "Wanted to remove an expression but its unique ptr is missing.");
+  c->per_fusion_exprs_[this].erase(expr);
   c->exprs_.erase(expr);
   c->exprs_up_.erase(expr_in_deque);
 }
@@ -414,6 +413,7 @@ void Fusion::removeVal(Val* val) {
   NVF_ERROR(
       val_in_deque != c->vals_up_.end(),
       "Wanted to remove a value but its unique ptr is missing.");
+  c->per_fusion_vals_[this].erase(val);
   c->vals_.erase(val);
   c->vals_up_.erase(val_in_deque);
 
@@ -446,13 +446,11 @@ void Fusion::removeStatementsCreatedAfter(
     for (Val* in : e->inputs()) {
       in->removeUse(e);
     }
+    c->per_fusion_exprs_[this].erase(e);
     c->exprs_.erase(e);
     c->exprs_up_.pop_back();
   }
 
-  // Null out any special value caches that point to vals about to be destroyed.
-  // This prevents dangling pointers when special vals are lazily created inside
-  // a StatementGuard scope.
   while (std::ssize(c->vals_up_) > num_vals_before) {
     Val* v = c->vals_up_.back().get();
     if (v == zero_val_) {
@@ -466,6 +464,7 @@ void Fusion::removeStatementsCreatedAfter(
     } else if (v == magic_zero_val_) {
       magic_zero_val_ = nullptr;
     }
+    c->per_fusion_vals_[this].erase(v);
     c->vals_.erase(v);
     c->vals_up_.pop_back();
   }
@@ -932,6 +931,7 @@ void Fusion::registerVal(Val* val) {
   auto* c = ir_container();
   c->vals_up_.emplace_back(val);
   c->vals_.insert(val);
+  c->per_fusion_vals_[this].insert(val);
   val->setName(IrContainerPasskey(), c->getValName(val->vtype()));
 }
 
@@ -948,6 +948,7 @@ void Fusion::registerExpr(Expr* expr) {
   auto* c = ir_container();
   c->exprs_up_.emplace_back(expr);
   c->exprs_.insert(expr);
+  c->per_fusion_exprs_[this].insert(expr);
   expr->setName(IrContainerPasskey(), c->getExprName());
 
   for (Val* input : expr->inputs()) {

csrc/fusion.h

@@ -521,31 +521,29 @@ class NVF_API Fusion : public PolymorphicBase {
   }
 
   // Collections access (return values in insertion order)
-  const std::deque<Val*> deterministic_vals() const noexcept {
-    return ir_container()->deterministic_vals();
+  std::deque<Val*> deterministic_vals() const noexcept {
+    return ir_container()->deterministicValsOwnedBy(this);
   }
 
-  const std::deque<Expr*> deterministic_exprs() const noexcept {
-    return ir_container()->deterministic_exprs();
+  std::deque<Expr*> deterministic_exprs() const noexcept {
+    return ir_container()->deterministicExprsOwnedBy(this);
   }
 
-  const std::unordered_map<Val*, int64_t> deterministic_vals_map()
-      const noexcept {
-    return ir_container()->deterministic_vals_map();
+  std::unordered_map<Val*, int64_t> deterministic_vals_map() const noexcept {
+    return ir_container()->deterministicValsMapOwnedBy(this);
   }
 
-  const std::unordered_map<Expr*, int64_t> deterministic_exprs_map()
-      const noexcept {
-    return ir_container()->deterministic_exprs_map();
+  std::unordered_map<Expr*, int64_t> deterministic_exprs_map() const noexcept {
+    return ir_container()->deterministicExprsMapOwnedBy(this);
   }
 
   // Collections access (unordered sets)
   const std::unordered_set<Expr*>& unordered_exprs() const noexcept {
-    return ir_container()->unordered_exprs();
+    return ir_container()->exprsOwnedBy(this);
   }
 
   const std::unordered_set<Val*>& vals() const noexcept {
-    return ir_container()->vals();
+    return ir_container()->valsOwnedBy(this);
   }
 
   // Count queries

csrc/ir/container.cpp

@@ -80,6 +80,9 @@ 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.per_fusion_vals_, b.per_fusion_vals_);
+  std::swap(a.per_fusion_exprs_, b.per_fusion_exprs_);
 }
 
 IrCloner IrContainer::copy(
@@ -124,6 +127,8 @@ void IrContainer::clear() noexcept {
   exprs_up_.clear();
   val_type_name_map_.clear();
   expr_name_counter_ = 0;
+  per_fusion_vals_.clear();
+  per_fusion_exprs_.clear();
 }
 
 bool IrContainer::inContainer(const Statement* const_stmt) const {
@@ -183,4 +188,130 @@ const std::unordered_set<Fusion*>& IrContainer::sharingFusions() const {
   return sharing_fusions_;
 }
 
+const std::unordered_set<Val*>& IrContainer::valsOwnedBy(
+    const Fusion* fusion) const {
+  static const std::unordered_set<Val*> empty;
+  auto it = per_fusion_vals_.find(fusion);
+  return it != per_fusion_vals_.end() ? it->second : empty;
+}
+
+const std::unordered_set<Expr*>& IrContainer::exprsOwnedBy(
+    const Fusion* fusion) const {
+  static const std::unordered_set<Expr*> empty;
+  auto it = per_fusion_exprs_.find(fusion);
+  return it != per_fusion_exprs_.end() ? it->second : empty;
+}
+
+void IrContainer::transferStatementOwnership(
+    const Fusion* from,
+    const Fusion* to) {
+  auto vals_it = per_fusion_vals_.find(from);
+  if (vals_it != per_fusion_vals_.end()) {
+    auto& to_vals = per_fusion_vals_[to];
+    to_vals.insert(vals_it->second.begin(), vals_it->second.end());
+    per_fusion_vals_.erase(vals_it);
+  }
+
+  auto exprs_it = per_fusion_exprs_.find(from);
+  if (exprs_it != per_fusion_exprs_.end()) {
+    auto& to_exprs = per_fusion_exprs_[to];
+    to_exprs.insert(exprs_it->second.begin(), exprs_it->second.end());
+    per_fusion_exprs_.erase(exprs_it);
+  }
+}
+
+void IrContainer::removeStatementsOwnedBy(const Fusion* fusion) {
+  auto vals_it = per_fusion_vals_.find(fusion);
+  if (vals_it != per_fusion_vals_.end()) {
+    for (auto it = vals_up_.begin(); it != vals_up_.end();) {
+      if (vals_it->second.count(it->get()) > 0) {
+        vals_.erase(it->get());
+        it = vals_up_.erase(it);
+      } else {
+        ++it;
+      }
+    }
+    per_fusion_vals_.erase(vals_it);
+  }
+
+  auto exprs_it = per_fusion_exprs_.find(fusion);
+  if (exprs_it != per_fusion_exprs_.end()) {
+    for (auto it = exprs_up_.begin(); it != exprs_up_.end();) {
+      if (exprs_it->second.count(it->get()) > 0) {
+        exprs_.erase(it->get());
+        it = exprs_up_.erase(it);
+      } else {
+        ++it;
+      }
+    }
+    per_fusion_exprs_.erase(exprs_it);
+  }
+}
+
+std::deque<Val*> IrContainer::deterministicValsOwnedBy(
+    const Fusion* fusion) const noexcept {
+  std::deque<Val*> result;
+  auto it = per_fusion_vals_.find(fusion);
+  if (it == per_fusion_vals_.end()) {
+    return result;
+  }
+  const auto& owned = it->second;
+  for (const auto& val_up : vals_up_) {
+    if (owned.count(val_up.get()) > 0) {
+      result.push_back(val_up.get());
+    }
+  }
+  return result;
+}
+
+std::deque<Expr*> IrContainer::deterministicExprsOwnedBy(
+    const Fusion* fusion) const noexcept {
+  std::deque<Expr*> result;
+  auto it = per_fusion_exprs_.find(fusion);
+  if (it == per_fusion_exprs_.end()) {
+    return result;
+  }
+  const auto& owned = it->second;
+  for (const auto& expr_up : exprs_up_) {
+    if (owned.count(expr_up.get()) > 0) {
+      result.push_back(expr_up.get());
+    }
+  }
+  return result;
+}
+
+std::unordered_map<Val*, int64_t> IrContainer::deterministicValsMapOwnedBy(
+    const Fusion* fusion) const noexcept {
+  std::unordered_map<Val*, int64_t> result;
+  auto it = per_fusion_vals_.find(fusion);
+  if (it == per_fusion_vals_.end()) {
+    return result;
+  }
+  const auto& owned = it->second;
+  int64_t count = 0;
+  for (const auto& val_up : vals_up_) {
+    if (owned.count(val_up.get()) > 0) {
+      result[val_up.get()] = count++;
+    }
+  }
+  return result;
+}
+
+std::unordered_map<Expr*, int64_t> IrContainer::deterministicExprsMapOwnedBy(
+    const Fusion* fusion) const noexcept {
+  std::unordered_map<Expr*, int64_t> result;
+  auto it = per_fusion_exprs_.find(fusion);
+  if (it == per_fusion_exprs_.end()) {
+    return result;
+  }
+  const auto& owned = it->second;
+  int64_t count = 0;
+  for (const auto& expr_up : exprs_up_) {
+    if (owned.count(expr_up.get()) > 0) {
+      result[expr_up.get()] = count++;
+    }
+  }
+  return result;
+}
+
 } // namespace nvfuser

csrc/ir/container.h

@@ -137,8 +137,26 @@ class IrContainer {
   bool hasMultipleFusions() const;
   const std::unordered_set<Fusion*>& sharingFusions() const;
 
+  NVF_API const std::unordered_set<Val*>& valsOwnedBy(
+      const Fusion* fusion) const;
+  const std::unordered_set<Expr*>& exprsOwnedBy(const Fusion* fusion) const;
+  void transferStatementOwnership(const Fusion* from, const Fusion* to);
+  void removeStatementsOwnedBy(const Fusion* fusion);
+
+  std::deque<Val*> deterministicValsOwnedBy(
+      const Fusion* fusion) const noexcept;
+  std::deque<Expr*> deterministicExprsOwnedBy(
+      const Fusion* fusion) const noexcept;
+  std::unordered_map<Val*, int64_t> deterministicValsMapOwnedBy(
+      const Fusion* fusion) const noexcept;
+  std::unordered_map<Expr*, int64_t> deterministicExprsMapOwnedBy(
+      const Fusion* fusion) const noexcept;
+
  private:
   std::unordered_set<Fusion*> sharing_fusions_;
+  std::unordered_map<const Fusion*, std::unordered_set<Val*>> per_fusion_vals_;
+  std::unordered_map<const Fusion*, std::unordered_set<Expr*>>
+      per_fusion_exprs_;
 };
 
 } // namespace nvfuser