[SOT][3.11] fix pass NULL to resume function in jump breakgraph (#58187)

SigureMo · web-flow · commit 4787ac1d04e5 · 2023-10-18T23:20:59.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -245,11 +245,15 @@ def start_compile_with_name_store(self, ret_vars, to_store_vars):
         class VariableLoader:
             def __init__(self, index_for_load, pycode_gen):
                 self._index_for_load = index_for_load
-                self._pycode_gen = pycode_gen
+                self._pycode_gen: PyCodeGen = pycode_gen
 
-            def load(self, var):
+            def load(self, var, allow_push_null=True):
                 if isinstance(var, NullVariable):
-                    var.reconstruct(self._pycode_gen)
+                    if allow_push_null:
+                        var.reconstruct(self._pycode_gen)
+                    else:
+                        # Avoid passing NULL as a parameter to the resume function
+                        self._pycode_gen.gen_load_null_variable()
                     return
                 self._pycode_gen.gen_load_fast(self._index_for_load[var.id])
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -1570,8 +1570,10 @@ def _break_graph_in_jump(self, result: VariableBase, instr: Instruction):
                 if_fn, if_fn.__code__.co_name
             )
             insert_index = len(self._graph.pycode_gen._instructions) - 1
-            for stack_arg in self.stack:
-                var_loader.load(stack_arg)
+            for i, stack_arg in enumerate(self.stack):
+                var_loader.load(
+                    stack_arg, allow_push_null=i >= len(self.stack) - 1
+                )
             for name in if_inputs:
                 var_loader.load(self.get_var(name))
             self._graph.pycode_gen.gen_call_function(
@@ -1587,8 +1589,10 @@ def _break_graph_in_jump(self, result: VariableBase, instr: Instruction):
                 else_fn, else_fn.__code__.co_name
             )
             jump_to = self._graph.pycode_gen._instructions[-1]
-            for stack_arg in self.stack:
-                var_loader.load(stack_arg)
+            for i, stack_arg in enumerate(self.stack):
+                var_loader.load(
+                    stack_arg, allow_push_null=i >= len(self.stack) - 1
+                )
             for name in else_inputs:
                 var_loader.load(self.get_var(name))
             self._graph.pycode_gen.gen_call_function(
@@ -1658,16 +1662,9 @@ def _break_graph_in_call(
         pop_n = push_n - stack_effect
 
         for i, stack_arg in enumerate(self.stack):
-            # Avoid passing NULL as a parameter to the resume function
-            if (
-                isinstance(stack_arg, NullVariable)
-                and i < len(self.stack) - pop_n
-            ):
-                self._graph.pycode_gen.gen_load_object(
-                    NullVariable(), f'null_var_{i}', push_null=False
-                )
-            else:
-                var_loader.load(stack_arg)
+            var_loader.load(
+                stack_arg, allow_push_null=i >= len(self.stack) - pop_n
+            )
 
         # gen call resume fn opcode
         # NOTE(SigureMo): In Python 3.11，we need generate KW_NAMES if the call shape is not None.
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -21,6 +21,7 @@
 import random
 import sys
 import types
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 import opcode
@@ -551,6 +552,12 @@ def gen_resume_fn_at(
 
         return fn, inputs
 
+    @cached_property
+    def global_null_variable(self):
+        from .variables.basic import NullVariable
+
+        return NullVariable()
+
     def gen_disable_eval_frame(self):
         """
         Generates instructions to disable the evaluation frame.
@@ -744,6 +751,13 @@ def gen_load_object(self, obj, obj_name: str, push_null: bool = True):
             self._f_globals[obj_name] = obj
         self.gen_load_global(obj_name, push_null=push_null)
 
+    def gen_load_null_variable(self):
+        """
+        Generate the bytecode for loading a null variable.
+        """
+        null_var = self.global_null_variable
+        self.gen_load_object(null_var, "___null_var", push_null=False)
+
     def gen_load_fast(self, name):
         """
         Generate the bytecode for loading a local variable.
diff --git a/test/sot/test_break_graph.py b/test/sot/test_break_graph.py
@@ -153,5 +153,16 @@ def test_simple(self):
         self.assert_results(test_break_graph_repeat, x)
 
 
+def break_graph_resume_pass_null(x, y):
+    return paddle.add(x, y[0:50] if y is not None else None)
+
+
+class TestBreakGraphResumePassNull(TestCaseBase):
+    def test_break_graph_resume_pass_null(self):
+        x = paddle.rand([50, 50], dtype=paddle.float32)
+        y = paddle.rand([100, 50], dtype=paddle.float32)
+        self.assert_results(break_graph_resume_pass_null, x, y)
+
+
 if __name__ == "__main__":
     unittest.main()
