Fix dangling special val pointers after StatementGuard rollback

Special vals (trueVal, falseVal, oneVal, etc.) can be lazily created
inside a StatementGuard scope (e.g. by simplifyExpr called from
haveDifferentShardings). When the guard rolls back, it pops vals_up_
back to the snapshot, destroying those vals while the Fusion cache
pointers still reference them. Subsequent calls return dangling pointers
causing UB — this manifested as LoopShardedSplitReshapeIds incorrectly
classifying a reshape as resharding on CI.

Fusion::removeStatementsCreatedAfter now nulls out any special val cache
pointers that are about to be destroyed, so they get re-created on next
access.

Author: mdavis36

csrc/fusion.cpp

@@ -401,6 +401,34 @@ void Fusion::removeVal(Val* val) {
   invalidateTvsAndUses();
 }
 
+void Fusion::removeStatementsCreatedAfter(
+    int64_t num_exprs_before,
+    int64_t num_vals_before) {
+  // Before removing vals from vals_up_, null out any special value caches that
+  // point to vals that are about to be destroyed. This prevents dangling
+  // pointers when special vals are lazily created inside a StatementGuard scope.
+  int64_t current_num_vals = ir_container()->numVals();
+  if (current_num_vals > num_vals_before) {
+    auto& vals_up = ir_container()->vals_up_;
+    for (int64_t i = num_vals_before; i < current_num_vals; ++i) {
+      Val* v = vals_up[i].get();
+      if (v == zero_val_) {
+        zero_val_ = nullptr;
+      } else if (v == one_val_) {
+        one_val_ = nullptr;
+      } else if (v == true_val_) {
+        true_val_ = nullptr;
+      } else if (v == false_val_) {
+        false_val_ = nullptr;
+      } else if (v == magic_zero_val_) {
+        magic_zero_val_ = nullptr;
+      }
+    }
+  }
+  ir_container()->removeStatementsCreatedAfter(
+      num_exprs_before, num_vals_before);
+}
+
 void Fusion::addInput(Val* input) {
   assertInContainer(input, "Cannot register input ");
 

csrc/fusion.h

@@ -574,10 +574,7 @@ class NVF_API Fusion : public PolymorphicBase {
   // Statement removal
   void removeStatementsCreatedAfter(
       int64_t num_exprs_before,
-      int64_t num_vals_before) {
-    ir_container()->removeStatementsCreatedAfter(
-        num_exprs_before, num_vals_before);
-  }
+      int64_t num_vals_before);
 
  protected:
   friend SegmentCandidateFinder;

tests/cpp/test_statement_guard.cpp

@@ -51,4 +51,55 @@ TEST_F(StatementGuardTest, ExecuteAfterGuard) {
       executor_cache.fusion(), {out_tensor}, {in_tensor}, __LINE__, __FILE__);
 }
 
+// Regression test: special vals lazily created inside a StatementGuard scope
+// must not become dangling pointers after the guard rolls back.
+TEST_F(StatementGuardTest, LazySpecialValsNotDangling) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  TensorView* in = makeContigTensor(1);
+  fusion->addInput(in);
+  TensorView* out = set(in);
+  fusion->addOutput(out);
+
+  // Force lazy creation of trueVal/falseVal inside a StatementGuard scope.
+  // This reproduces the bug where haveDifferentShardings calls simplifyExpr
+  // inside a StatementGuard, which can lazily create special vals that then
+  // become dangling pointers when the guard rolls back.
+  {
+    StatementGuard sg(fusion.get());
+    // Directly trigger lazy creation of trueVal and falseVal
+    fusion->trueVal();
+    fusion->falseVal();
+    fusion->oneVal();
+  }
+
+  // After the guard, the special vals should still be valid (re-created if the
+  // originals were destroyed by the guard's rollback).
+  Val* z = fusion->zeroVal();
+  Val* o = fusion->oneVal();
+  Val* t = fusion->trueVal();
+  Val* f = fusion->falseVal();
+  EXPECT_NE(z, nullptr);
+  EXPECT_NE(o, nullptr);
+  EXPECT_NE(t, nullptr);
+  EXPECT_NE(f, nullptr);
+  EXPECT_TRUE(z->isZeroInt());
+  EXPECT_TRUE(o->isOneInt());
+  EXPECT_TRUE(t->isTrue());
+  EXPECT_TRUE(f->isFalse());
+
+  // The fusion should still be executable
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor in_tensor = at::randn({8}, at::device(at::kCUDA));
+  auto out_tensors = executor_cache.runFusionWithInputs({in_tensor});
+  ASSERT_EQ(out_tensors.size(), 1);
+  testValidate(
+      executor_cache.fusion(),
+      {out_tensors[0].as<at::Tensor>()},
+      {in_tensor},
+      __LINE__,
+      __FILE__);
+}
+
 } // namespace nvfuser