[SOT] Add inline call codeobj to global guard (#69803)

SigureMo · web-flow · commit 97bacd154e3a · 2024-11-29T14:07:33.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py
@@ -163,7 +163,8 @@ def __init__(
         self._fn_var = fn_variable
         self.return_value: VariableBase | None = None
         self._fn_value = fn_variable.value
-        super().__init__(fn_variable.get_code(), fn_variable.graph)
+        self._code_var = fn_variable.get_code()
+        super().__init__(self._code_var.value, fn_variable.graph)
         self._name = "Inline"
         self._prepare_locals(*args, **kwargs)
         self._prepare_closure()
@@ -273,6 +274,7 @@ def inline_call(self) -> VariableBase:
         """
         Execute the inline call of the function.
         """
+        self._graph.add_global_guarded_variable(self._code_var)
         self.run()
         assert self.return_value is not None
         return self.return_value
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -125,8 +125,11 @@ def __init__(
     def get_py_value(self, allow_tensor=False):
         return self.value
 
-    def get_code(self) -> types.CodeType:
-        return self.value.__code__
+    def get_code(self) -> VariableBase:
+        code_obj_var = VariableFactory.from_value(
+            self.value.__code__, self.graph, GetAttrTracker(self, "__code__")
+        )
+        return code_obj_var
 
     def bind(self, instance: VariableBase, name: str):
         method_var = MethodVariable(
diff --git a/test/sot/test_06_call_function.py b/test/sot/test_06_call_function.py
@@ -14,7 +14,10 @@
 
 import unittest
 
-from test_case_base import TestCaseBase
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
 
 import paddle
 
@@ -150,5 +153,30 @@ def test_call8(self):
         self.assert_results(foo_8, paddle.to_tensor(9))
 
 
+def apply_fn(fn, x):
+    return fn(x)
+
+
+def fn1(x):
+    return x + 1
+
+
+def fn2(x):
+    return x - 1
+
+
+class TestApplyDifferentFunctions(TestCaseBase):
+    def test_apply_fn(self):
+        x = 1
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            self.assert_results(apply_fn, fn1, x)
+            self.assertEqual(ctx.translate_count, 1)
+            self.assert_results(apply_fn, fn2, x)
+            self.assertEqual(ctx.translate_count, 2)
+            self.assert_results(apply_fn, fn1, x)
+            self.assertEqual(ctx.translate_count, 2)
+
+
 if __name__ == "__main__":
     unittest.main()
