chore: import upstream snapshot with attribution

python/paddle/jit/sot/utils/call_ast_utils.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2023 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
+#
+#     http://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 inspect
+import types
+
+import paddle
+
+from .envs import ENV_SOT_WITH_CONTROL_FLOW
+from .exceptions import InnerError
+from .utils import Singleton
+
+try_ast_codes = set()
+
+
+def try_ast_func(func):
+    def _is_wrapped(f):
+        return hasattr(f, '__wrapped__')
+
+    unwrapped_f = func
+    if hasattr(unwrapped_f, "__code__"):
+        try_ast_codes.add(func.__code__)
+
+    while _is_wrapped(unwrapped_f):
+        unwrapped_f = unwrapped_f.__wrapped__
+        if hasattr(unwrapped_f, "__code__"):
+            try_ast_codes.add(func.__code__)
+
+    return func
+
+
+class StaticFunctionManager(metaclass=Singleton):
+    def __init__(self):
+        self.code_map = {}
+
+    def ast_transform_with_frame(self, frame):
+        code = frame.f_code
+        if code not in try_ast_codes:
+            return None
+        if code not in self.code_map:
+            if code.co_name.startswith("#") or code.co_name.startswith("$"):
+                self.code_map[code] = None
+            elif len(code.co_cellvars) + len(code.co_freevars) != 0:
+                self.code_map[code] = None
+            else:
+                function = types.FunctionType(
+                    code,
+                    frame.f_globals,
+                    code.co_name,
+                    (),
+                    (),
+                )
+                function = paddle.jit.to_static(function, full_graph=True)
+                self.code_map[code] = function
+
+        return self.code_map[code]
+
+    def ast_transform_with_callable(self, fn):
+        if not inspect.isfunction(fn) or not hasattr(fn, "__code__"):
+            return None
+
+        code = fn.__code__
+        if code not in try_ast_codes:
+            return None
+        if code not in self.code_map:
+            if code.co_name.startswith("#") or code.co_name.startswith("$"):
+                self.code_map[code] = None
+            elif len(code.co_cellvars) + len(code.co_freevars) != 0:
+                self.code_map[code] = None
+            else:
+                self.code_map[code] = paddle.jit.to_static(fn, full_graph=True)
+
+        return self.code_map[code]
+
+
+def get_static_function(obj, type_):
+    if ENV_SOT_WITH_CONTROL_FLOW.get():
+        if type_ == "eval_frame":
+            return StaticFunctionManager().ast_transform_with_frame(obj)
+        elif type_ == "inline_call":
+            return StaticFunctionManager().ast_transform_with_callable(obj)
+        else:
+            raise InnerError(f"Can not get static function with type {type_}.")
+    return None