[SOT] Cache resume fn code to avoid repeated simulation (#69300)

SigureMo · web-flow · commit f471051a9737 · 2024-11-12T01:56:30.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -66,7 +66,10 @@
     MAKE_FUNCTION_FLAG as MF,
     IntrinsicsUnaryFunctions,
 )
-from .pycode_generator import PyCodeGen
+from .pycode_generator import (
+    ResumeFunctionCreator,
+    ResumeFunctionType,
+)
 from .tracker import (
     CellTracker,
     ConstTracker,
@@ -1825,15 +1828,21 @@ def create_if_branch_fn(start_idx, input_var_names, is_pop_jump_branch):
                 and not is_pop_jump_branch
             ):
                 return None
-            pycode_gen = PyCodeGen(self._frame)
+            cache_key = (ResumeFunctionType.IF_RESUME, self._code, start_idx)
+            resume_fn_creator = ResumeFunctionCreator(self._frame)
+            if (
+                maybe_resume_fn := resume_fn_creator.lookup(cache_key)
+            ) is not None:
+                return maybe_resume_fn
+            pycode_gen = resume_fn_creator.codegen
             origin_instrs = get_instructions(pycode_gen._origin_code)
-            pycode_gen.set_function_inputs(
+            resume_fn_creator.set_inputs(
                 input_var_names, stack_size=stack_size_after_if
             )
             pycode_gen.extend_instrs(origin_instrs[start_idx:])
             # the resume_fn contains return code, so we don't need set output here
             # global vars are updated correctly, and need local vars will return
-            resume_fn = pycode_gen.create_function()
+            resume_fn = resume_fn_creator.generate(cache_key=cache_key)
             return resume_fn
 
         true_fn_read_names, _ = analysis_used_names(
@@ -1960,15 +1969,21 @@ def _break_graph_when_call(
         def create_resume_fn():
             if self._instructions[next_index].opname == "RETURN_VALUE":
                 return None
-            pycode_gen = PyCodeGen(self._frame)
+            cache_key = (ResumeFunctionType.CALL_RESUME, self._code, next_index)
+            resume_fn_creator = ResumeFunctionCreator(self._frame)
+            if (
+                maybe_resume_fn := resume_fn_creator.lookup(cache_key)
+            ) is not None:
+                return maybe_resume_fn
+            pycode_gen = resume_fn_creator.codegen
             origin_instrs = get_instructions(pycode_gen._origin_code)
-            pycode_gen.set_function_inputs(
+            resume_fn_creator.set_inputs(
                 input_var_names, stack_size=stack_size_after_call
             )
             pycode_gen.extend_instrs(origin_instrs[next_index:])
             # the resume_fn contains return code, so we don't need set output here
             # global vars are updated correctly, and need local vars will return
-            resume_fn = pycode_gen.create_function()
+            resume_fn = resume_fn_creator.generate(cache_key=cache_key)
             return resume_fn
 
         resume_fn = create_resume_fn()
@@ -2048,9 +2063,20 @@ def _break_graph_when_for_loop(
         loop_body_outputs = [*list(loop_body_write_names), "_break_flag"]
 
         def create_loop_body():
-            pycode_gen = PyCodeGen(self._frame)
+            cache_key = (
+                ResumeFunctionType.LOOP_BODY_RESUME,
+                self._code,
+                loop_body_start_idx,
+                loop_body_end_idx,
+            )
+            resume_fn_creator = ResumeFunctionCreator(self._frame)
+            if (
+                maybe_resume_fn := resume_fn_creator.lookup(cache_key)
+            ) is not None:
+                return maybe_resume_fn
+            pycode_gen = resume_fn_creator.codegen
 
-            pycode_gen.set_function_inputs(loop_body_inputs, stack_size=0)
+            resume_fn_creator.set_inputs(loop_body_inputs, stack_size=0)
 
             origin_instrs = get_instructions(pycode_gen._origin_code)
             for_iter = origin_instrs[for_iter_idx]
@@ -2085,8 +2111,8 @@ def create_loop_body():
                     instr.jump_to = nop_for_break
 
             # outputs is the same as inputs
-            pycode_gen.set_function_outputs(loop_body_outputs)
-            loop_body_fn = pycode_gen.create_function()
+            resume_fn_creator.set_outputs(loop_body_outputs)
+            loop_body_fn = resume_fn_creator.generate(cache_key=cache_key)
 
             log(
                 3,
@@ -2109,7 +2135,17 @@ def create_loop_body():
         def create_after_loop_fn():
             if self._instructions[loop_body_end_idx].opname == "RETURN_VALUE":
                 return None
-            pycode_gen = PyCodeGen(self._frame)
+            cache_key = (
+                ResumeFunctionType.AFTER_LOOP_RESUME,
+                self._code,
+                loop_body_end_idx,
+            )
+            resume_fn_creator = ResumeFunctionCreator(self._frame)
+            if (
+                maybe_resume_fn := resume_fn_creator.lookup(cache_key)
+            ) is not None:
+                return maybe_resume_fn
+            pycode_gen = resume_fn_creator.codegen
             origin_instrs = get_instructions(pycode_gen._origin_code)
             resume_fn_end_idx = loop_body_end_idx
 
@@ -2118,13 +2154,13 @@ def create_after_loop_fn():
                 assert origin_instrs[loop_body_end_idx].opname == "END_FOR"
                 resume_fn_end_idx += 1
 
-            pycode_gen.set_function_inputs(
+            resume_fn_creator.set_inputs(
                 after_loop_fn_inputs, stack_size=len(self.stack) - 1
             )
             pycode_gen.extend_instrs(origin_instrs[resume_fn_end_idx:])
             # the resume_fn contains return code, so we don't need set output here
             # global vars are updated correctly, and need local vars will return
-            after_loop_fn = pycode_gen.create_function()
+            after_loop_fn = resume_fn_creator.generate(cache_key=cache_key)
             return after_loop_fn
 
         after_loop_fn = create_after_loop_fn()
@@ -2250,10 +2286,21 @@ def _inline_call_for_loop(
 
         # 2. create inline call loop fn
         def create_inline_call_fn():
-            pycode_gen = PyCodeGen(self._frame)
+            cache_key = (
+                ResumeFunctionType.LOOP_BODY_INLINE_CALL,
+                self._code,
+                start_idx,
+                end_idx,
+            )
+            resume_fn_creator = ResumeFunctionCreator(self._frame)
+            if (
+                maybe_resume_fn := resume_fn_creator.lookup(cache_key)
+            ) is not None:
+                return maybe_resume_fn
+            pycode_gen = resume_fn_creator.codegen
             origin_instrs = get_instructions(pycode_gen._origin_code)
 
-            pycode_gen.set_function_inputs(input_var_names, stack_size=0)
+            resume_fn_creator.set_inputs(input_var_names, stack_size=0)
 
             # 2.1. load iter, it is a input of loop fn
             pycode_gen.gen_load_fast(iterator.id)
@@ -2292,8 +2339,8 @@ def create_inline_call_fn():
             if sys.version_info >= (3, 12):
                 for_iter_instr.jump_to = end_for
 
-            pycode_gen.set_function_outputs(output_var_names)
-            inline_call_fn = pycode_gen.create_function()
+            resume_fn_creator.set_outputs(output_var_names)
+            inline_call_fn = resume_fn_creator.generate(cache_key=cache_key)
 
             log(
                 3,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -23,6 +23,7 @@
 import random
 import sys
 import types
+from enum import Enum
 from functools import cached_property
 from typing import TYPE_CHECKING
 
@@ -435,45 +436,6 @@ def __init__(
         self.hooks = []
         if self.disable_eval_frame:
             self.gen_disable_eval_frame()
-        self.fn_name = ResumeFnNameFactory().next()
-
-    def set_function_inputs(self, inputs: list[str], stack_size: int):
-        stack_arg_str = self.fn_name + '_stack_{}'
-
-        self._code_options['co_argcount'] = len(inputs) + stack_size
-        self._code_options['co_varnames'] = list(
-            [stack_arg_str.format(i) for i in range(stack_size)]
-            + inputs
-            + [
-                var_name
-                for var_name in self._origin_code.co_varnames
-                if var_name not in inputs
-            ]
-        )
-
-        self._instructions.extend(
-            [
-                gen_instr('LOAD_FAST', argval=stack_arg_str.format(i))
-                for i in range(stack_size)
-            ]
-        )
-
-    def set_function_outputs(self, outputs: list[str]):
-        for name in outputs:
-            self.gen_load(name)
-        self.gen_build_tuple(len(outputs))
-        self.gen_return()
-
-    def create_function(self) -> types.FunctionType:
-        self.update_code_name(self.fn_name, is_resumed_fn=True)
-        self._code_options['co_flags'] &= ~(
-            inspect.CO_VARARGS | inspect.CO_VARKEYWORDS
-        )
-        new_code = self.gen_pycode()
-        if len(new_code.co_freevars) + len(new_code.co_cellvars) > 0:
-            raise FallbackError("Break graph in closure is not support.")
-        fn = types.FunctionType(new_code, self._f_globals, new_code.co_name)
-        return fn
 
     def insert_prefix_instructions(self):
         """
@@ -1019,3 +981,81 @@ def extend_instrs(self, instrs):
 
     def pop_instr(self):
         self._instructions.pop()
+
+
+class ResumeFunctionType(Enum):
+    # If breakgraph
+    IF_RESUME = 0
+    # Call breakgraph
+    CALL_RESUME = 1
+    # Loop breakgraph
+    LOOP_BODY_RESUME = 2
+    AFTER_LOOP_RESUME = 3
+    # Loop inline call
+    LOOP_BODY_INLINE_CALL = 4
+
+
+class ResumeFunctionCreator:
+    CODE_CACHE = {}
+
+    def __init__(
+        self, frame: types.FrameType, disable_eval_frame: bool = False
+    ):
+        self.codegen = PyCodeGen(frame, disable_eval_frame)
+        self.name = ResumeFnNameFactory().next()
+
+    def set_inputs(self, inputs: list[str], stack_size: int):
+        stack_arg_str = self.name + '_stack_{}'
+
+        self.codegen._code_options['co_argcount'] = len(inputs) + stack_size
+        self.codegen._code_options['co_varnames'] = list(
+            [stack_arg_str.format(i) for i in range(stack_size)]
+            + inputs
+            + [
+                var_name
+                for var_name in self.codegen._origin_code.co_varnames
+                if var_name not in inputs
+            ]
+        )
+
+        self.codegen._instructions.extend(
+            [
+                gen_instr('LOAD_FAST', argval=stack_arg_str.format(i))
+                for i in range(stack_size)
+            ]
+        )
+
+    def set_outputs(self, outputs: list[str]):
+        for name in outputs:
+            self.codegen.gen_load(name)
+        self.codegen.gen_build_tuple(len(outputs))
+        self.codegen.gen_return()
+
+    @staticmethod
+    def validate_code(code):
+        if len(code.co_freevars) + len(code.co_cellvars) > 0:
+            raise FallbackError("Break graph in closure is not support.")
+
+    def lookup(self, cache_key):
+        if cache_key in self.CODE_CACHE:
+            cached_code = self.CODE_CACHE[cache_key]
+            ResumeFunctionCreator.validate_code(cached_code)
+            return types.FunctionType(
+                cached_code, self.codegen._f_globals, cached_code.co_name
+            )
+        return None
+
+    def generate(self, cache_key=None) -> types.FunctionType:
+        self.codegen.update_code_name(self.name, is_resumed_fn=True)
+        self.codegen._code_options['co_flags'] &= ~(
+            inspect.CO_VARARGS | inspect.CO_VARKEYWORDS
+        )
+        new_code = self.codegen.gen_pycode()
+        # TODO(SigureMo): cache_key should not be None
+        if cache_key is not None:
+            self.CODE_CACHE[cache_key] = new_code
+        ResumeFunctionCreator.validate_code(new_code)
+        fn = types.FunctionType(
+            new_code, self.codegen._f_globals, new_code.co_name
+        )
+        return fn
diff --git a/test/sot/test_resume_cache.py b/test/sot/test_resume_cache.py
@@ -0,0 +1,46 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
+
+import paddle
+
+
+def if_breakgraph(x, y):
+    z = x + y
+    if x:
+        z = z + 1
+    else:
+        z = z - 1
+    return z
+
+
+class TestResumeCache(TestCaseBase):
+    def test_resume_cache_in_if_breakgraph(self):
+        x = paddle.to_tensor(0)
+        with test_instruction_translator_cache_context() as cache:
+            self.assertEqual(cache.translate_count, 0)
+            self.assert_results(if_breakgraph, x, 1)
+            self.assertEqual(cache.translate_count, 2)
+            self.assert_results(if_breakgraph, x, 2)
+            self.assertEqual(cache.translate_count, 3)
+
+
+if __name__ == "__main__":
+    unittest.main()
