[dynamo] Replace raw None with ConstantVariable on stack (#169325)

dsashidh · pytorchmergebot · commit 59b048f0bb04 · 2026-03-10T20:09:02.000Z
Fixes #168994 Fixes the issue where raw None values were being pushed onto the dynamo stack instead of wrapped ConstantVariable(None) objects. This caused crashes when code expected VariableTracker methods. - Updated push( ) signature: removed Optional[VariableTracker] - Fixed bytecode handlers (BEGIN_FINALLY, WITH_CLEANUP_START, etc.) to push ConstantVariable.create(None) instead of raw None - Updated assertions to check for ConstantVariable with value is None instead of raw None - Fixed DebuggingVariable.call_function( ) to return ConstantVariable.create(None) Pull Request resolved: #169325 Approved by: https://github.com/williamwen42
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
@@ -6353,6 +6353,29 @@ def fn(x):
         res2 = opt_fn(x)
         self.assertEqual(res, res2)
 
+    def test_function_return_none_creates_constant_variable(self):
+        """
+        Test that functions returning None properly return ConstantVariable.create(None)
+        instead of raw None, which would violate the stack's type contract.
+
+        Regression test for: Avoid using Optional[VariableTracker]
+        """
+
+        def gn(x):
+            return
+
+        torch._dynamo.config.reorderable_logging_functions.add(gn)
+
+        @torch.compile(backend="eager")
+        def fn(x):
+            x = x + 1
+            if gn(x) is None:
+                return x + 2
+            return x + 4
+
+        # If this doesn't crash, the test passes
+        fn(torch.ones(3))
+
     @patch.object(torch._dynamo.config, "capture_scalar_outputs", True)
     def test_tensor_ctor_list_of_tensor(self):
         def fn(x):
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -1758,11 +1758,11 @@ def run(self) -> None:
                     # twice is not an issue (second stop is a no op).
                     self.output.mark_bytecode_tracing_stop()
 
-    def push(self, val: VariableTracker | None) -> None:
-        assert val is None or isinstance(val, VariableTracker), (
+    def push(self, val: VariableTracker) -> None:
+        assert isinstance(val, VariableTracker), (
             f"push expects VariableTracker, got {typestr(val)}"
         )
-        self.stack.append(val)  # type: ignore[arg-type]
+        self.stack.append(val)
 
     def push_many(self, vals: list[VariableTracker]) -> None:
         for val in vals:
@@ -2111,20 +2111,6 @@ def SETUP_FINALLY(self, inst: Instruction) -> None:
         assert inst.target is not None
         self.block_stack.append(BlockStackEntry(inst, inst.target, len(self.stack)))
 
-    def BEGIN_FINALLY(self, inst: Instruction) -> None:
-        self.push(None)
-
-    def WITH_CLEANUP_START(self, inst: Instruction) -> None:
-        exit, exc = self.popn(2)
-        assert exc is None
-        self.push(exc)
-
-        self.push(exit.call_function(self, [CONSTANT_VARIABLE_NONE] * 3, {}))
-
-    def WITH_CLEANUP_FINISH(self, inst: Instruction) -> None:
-        self.popn(2)
-        self.push(None)
-
     def FOR_ITER(self, inst: Instruction) -> None:
         it = self.pop().realize()
         self.push(it)
@@ -2736,15 +2722,14 @@ def LOAD_METHOD(self, inst: Instruction) -> None:
             self.PUSH_NULL(inst)
             self.push(obj)
         else:
-            self.push(obj)
-            self.push(None)
+            raise AssertionError(
+                "LOAD_METHOD should have been rewritten to LOAD_ATTR. We should never reach here."
+            )
 
     def CALL_METHOD(self, inst: Instruction) -> None:
-        args = self.popn(inst.argval)
-        dummy = self.pop()
-        assert dummy is None
-        fn = self.pop()
-        self.call_function(fn, args, {})
+        raise AssertionError(
+            "CALL_METHOD should have been rewritten to CALL_FUNCTION. This function should never be called."
+        )
 
     def _load_attr(self, attr: Any) -> None:
         obj = self.pop()
diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
@@ -2026,16 +2026,15 @@ def is_reorderable_logging_function(
             and obj in torch._dynamo.config.reorderable_logging_functions
         )
 
-    # type: ignore[override]
     def call_function(
         self,
         tx: "InstructionTranslator",
         args: Sequence[VariableTracker],
         kwargs: dict[str, VariableTracker],
-    ) -> None:
+    ) -> VariableTracker:
         if tx.export:
             # For export cases, we can just make debugging functions no-ops
-            return
+            return CONSTANT_VARIABLE_NONE
 
         if not self.can_reorder_logs(self.value, args, kwargs):
             unimplemented(
@@ -2049,6 +2048,7 @@ def call_function(
             )
 
         tx.debug_locals.append((self, list(args)))
+        return CONSTANT_VARIABLE_NONE
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
         assert self.source is not None
