[Dy2St] Support super call in forward without arguments (#67021)

SigureMo · web-flow · commit 77a6200a49c7 · 2024-08-08T11:13:25.000+08:00
diff --git a/paddle/scripts/paddle_build.sh b/paddle/scripts/paddle_build.sh
@@ -2706,14 +2706,14 @@ set -x
         if [ -a "$PADDLE_ROOT/added_ut" ];then
             added_uts=^$(awk BEGIN{RS=EOF}'{gsub(/\n/,"$|^");print}' $PADDLE_ROOT/added_ut)$
             if [ "$WITH_ROCM" == "ON" ];then
-                env HIP_VISIBLE_DEVICES=0 ctest -R "(${added_uts})" -LE "RUN_TYPE=DIST|RUN_TYPE=EXCLUSIVE|RUN_TYPE=HYBRID" --output-on-failure --repeat-until-fail 3 --timeout 15;added_ut_error=$?
+                env HIP_VISIBLE_DEVICES=0 ctest -R "(${added_uts})" -LE "RUN_TYPE=DIST|RUN_TYPE=EXCLUSIVE|RUN_TYPE=HYBRID" --output-on-failure --repeat-until-fail 3 --timeout 20;added_ut_error=$?
             else
-                env CUDA_VISIBLE_DEVICES=0 ctest -R "(${added_uts})" -LE "RUN_TYPE=DIST|RUN_TYPE=EXCLUSIVE|RUN_TYPE=HYBRID" --output-on-failure --repeat-until-fail 3 --timeout 15;added_ut_error=$?
+                env CUDA_VISIBLE_DEVICES=0 ctest -R "(${added_uts})" -LE "RUN_TYPE=DIST|RUN_TYPE=EXCLUSIVE|RUN_TYPE=HYBRID" --output-on-failure --repeat-until-fail 3 --timeout 20;added_ut_error=$?
             fi
-            ctest -R "(${added_uts})" -L "RUN_TYPE=DIST|RUN_TYPE=EXCLUSIVE" --output-on-failure --repeat-until-fail 3 --timeout 15;added_ut_error_1=$?
+            ctest -R "(${added_uts})" -L "RUN_TYPE=DIST|RUN_TYPE=EXCLUSIVE" --output-on-failure --repeat-until-fail 3 --timeout 20;added_ut_error_1=$?
             if [ "$added_ut_error" != 0 ] && [ "$added_ut_error_1" != 0 ];then
                 echo "========================================"
-                echo "Added UT should not exceed 15 seconds"
+                echo "Added UT should not exceed 20 seconds"
                 echo "========================================"
                 exit 8;
             fi
diff --git a/python/paddle/jit/dy2static/__init__.py b/python/paddle/jit/dy2static/__init__.py
@@ -23,6 +23,7 @@
     convert_logical_not as Not,
     convert_logical_or as Or,
     convert_shape as Shape,
+    convert_super as WrapSuper,
     convert_var_dtype as AsDtype,
     convert_while_loop as While,
     create_bool_as_type,
diff --git a/python/paddle/jit/dy2static/convert_operators.py b/python/paddle/jit/dy2static/convert_operators.py
@@ -649,6 +649,12 @@ def convert_zip(*args):
     return zip(*args)
 
 
+def convert_super(super_fn):
+    if super_fn is super:
+        return super_fn
+    return lambda cls, instance: super_fn()
+
+
 # TODO(xiongkun): delete when list<variable> is ready.
 class VariableTuple:
     """
diff --git a/python/paddle/jit/dy2static/origin_info.py b/python/paddle/jit/dy2static/origin_info.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import inspect
+import textwrap
 from collections.abc import Sequence
 
 from paddle.base import core
@@ -149,7 +150,7 @@ def create_and_update_origin_info_map(
     """
 
     origin_info_map = {}
-    static_source = inspect.getsource(static_func)
+    static_source = textwrap.dedent(inspect.getsource(static_func))
     static_node = gast.parse(static_source)
     static_node = attach_origin_info(static_node, static_func)
 
diff --git a/python/paddle/jit/dy2static/transformers/super_transformer.py b/python/paddle/jit/dy2static/transformers/super_transformer.py
@@ -0,0 +1,60 @@
+#   Copyright (c) 2020 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.
+
+from paddle.utils import gast
+
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class SuperTransformer(BaseTransformer):
+    """
+    This class transforms super() into super(__class__, <first argument>).
+    """
+
+    def __init__(self, root):
+        self.root = root
+        self.first_arg = None
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        if self.first_arg is not None:
+            return self.generic_visit(node)
+
+        positional_args = node.args.posonlyargs + node.args.args
+        if not positional_args:
+            return self.generic_visit(node)
+
+        self.first_arg = positional_args[0].id
+
+        return self.generic_visit(node)
+
+    def visit_Call(self, node):
+        # super() -> _jst.WrapSuper(super)(x.__class__, x)
+        self.generic_visit(node)
+        if self.first_arg is None:
+            return node
+        if not isinstance(node.func, gast.Name):
+            return node
+        if node.func.id != "super":
+            return node
+        if node.args:
+            return node
+
+        new_fn_call_str = f"_jst.WrapSuper(super)({self.first_arg}.__class__, {self.first_arg})"
+        new_fn_call_ast = gast.parse(new_fn_call_str).body[0]
+        return new_fn_call_ast
diff --git a/python/paddle/jit/dy2static/transformers/transform.py b/python/paddle/jit/dy2static/transformers/transform.py
@@ -42,6 +42,7 @@
     NameloadJstTransformer,
 )
 from .return_transformer import ReturnTransformer
+from .super_transformer import SuperTransformer
 from .tensor_shape_transformer import TensorShapeTransformer
 from .tensorhook_transformer import RegisterHookTransformer
 from .typehint_transformer import TypeHintTransformer
@@ -95,6 +96,7 @@ def transfer_from_node_type(self, node):
 
         transformers = [
             TypeHintTransformer,  # remove all typehint
+            SuperTransformer,  # super() -> super(__class__, <first argument>)
             RegisterHookTransformer,
             EarlyReturnTransformer,
             AttributeJstTransformer,  # Tensor.size -> Tensor.size(), it's unnecessary in PIR mode
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -26,6 +26,7 @@
 import textwrap
 import types
 from importlib.machinery import SourceFileLoader
+from typing import Any
 
 import numpy as np
 
@@ -36,7 +37,7 @@
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 from paddle.framework import CUDAPinnedPlace
 from paddle.jit.utils import OrderedSet
-from paddle.utils import flatten
+from paddle.utils import flatten, gast
 
 from .ast_utils import ast_to_source_code
 
@@ -272,6 +273,105 @@ def get_temp_dir():
     return temp_dir
 
 
+def wrap_as_closure(tree: gast.AST, closure_vars: list[str]) -> gast.AST:
+    """
+    Wrap a function to a closure function.
+
+    Before:
+
+        >>> def fn(x):
+        ...     ...
+
+    After:
+
+        >>> def create_fn():
+        ...     closure_var_1 = None
+        ...     def fn(x):
+        ...         ...
+        ...     return fn
+        ... fn = create_fn()
+    """
+
+    def create_assign_node(name, value) -> gast.Assign:
+        return gast.Assign(
+            targets=[
+                gast.Name(
+                    id=name,
+                    ctx=gast.Store(),
+                    annotation=[],
+                    type_comment=[],
+                )
+            ],
+            value=value,
+        )
+
+    def create_wrppper_fn_def_node(name, body) -> gast.FunctionDef:
+        return gast.FunctionDef(
+            name=name,
+            args=gast.arguments(
+                args=[],
+                posonlyargs=[],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=[],
+                kwarg=None,
+                defaults=[],
+            ),
+            body=body,
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+        )
+
+    if not isinstance(tree, gast.Module):
+        return tree
+    if len(tree.body) != 1:
+        return tree
+    if not isinstance(tree.body[0], gast.FunctionDef):
+        return tree
+    fn_node = tree.body[0]
+    fn_name = fn_node.name
+    wrapper_fn_name = f"create_{fn_name}"
+    wrapper_fn_def_node = create_wrppper_fn_def_node(
+        wrapper_fn_name,
+        [
+            *[
+                create_assign_node(var, gast.Constant(value=None, kind=None))
+                for var in closure_vars
+            ],
+            fn_node,
+            gast.Return(
+                value=gast.Name(
+                    id=fn_name, ctx=gast.Load(), annotation=[], type_comment=[]
+                )
+            ),
+        ],
+    )
+
+    assign_node = create_assign_node(
+        fn_name,
+        gast.Call(
+            func=gast.Name(
+                id=wrapper_fn_name,
+                ctx=gast.Load(),
+                annotation=[],
+                type_comment=[],
+            ),
+            args=[],
+            keywords=[],
+        ),
+    )
+    return gast.Module(body=[wrapper_fn_def_node, assign_node], type_ignores=[])
+
+
+def wrap_cell(var: Any) -> types.CellType:
+    def closure_fn():
+        return var
+
+    assert closure_fn.__closure__ is not None
+    return closure_fn.__closure__[0]
+
+
 def ast_to_func(ast_root, dyfunc, delete_on_exit=True):
     """
     Transform modified AST of decorated function into python callable object.
@@ -284,16 +384,43 @@ def remove_if_exit(dir_path):
             shutil.rmtree(dir_path)
 
     def func_prefix(func):
-        pre_fix = func.__name__
+        prefix = func.__name__
         if hasattr(func, '__self__'):
             try:
-                pre_fix = func.__self__.__class__.__name__ + '_' + func.__name__
+                prefix = f"{func.__self__.__class__.__name__}_{func.__name__}"
             except:
                 pass
-        return pre_fix
+        return prefix
+
+    def get_new_closure(original_fn, generated_fn):
+        if generated_fn.__closure__ is None:
+            return None
+        assert original_fn.__closure__ is not None
+        original_closure_vars = inspect.getclosurevars(original_fn).nonlocals
+        generated_free_vars = generated_fn.__code__.co_freevars
+        return tuple(
+            wrap_cell(original_closure_vars[freevar_name])
+            for freevar_name in generated_free_vars
+        )
+
+    def get_new_globals(original_fn, generated_fn):
+        globals_attr_name = "__globals__"
+        original_fn_globals = getattr(original_fn, globals_attr_name, {})
+        generated_fn_globals = getattr(generated_fn, globals_attr_name, {})
+
+        original_fn_globals_exclude_builtin = {
+            k: v
+            for k, v in original_fn_globals.items()
+            if not (k.startswith('__') and k.endswith('__'))
+        }
+        return {**generated_fn_globals, **original_fn_globals_exclude_builtin}
+
+    dyfunc_closures = inspect.getclosurevars(dyfunc).nonlocals
+    ast_root = wrap_as_closure(ast_root, list(dyfunc_closures.keys()))
 
     source = ast_to_source_code(ast_root)
     source = _inject_import_statements() + source
+
     temp_dir = get_temp_dir()
     f = tempfile.NamedTemporaryFile(
         mode='w',
@@ -330,11 +457,17 @@ def func_prefix(func):
             f'Function: {func_name} doesn\'t exist in the Module transformed from AST.'
         )
     # After transform dygraph function into callable_func saved in tmp file,
-    # it lost the global variables from imported statements or defined in source file.
-    # Recovers the necessary variables by `__globals__`.
-    recover_globals_attribute(dyfunc, callable_func)
+    # it lost the global and closure variables from imported statements or defined
+    # in source file. Recovers the necessary variables by `__globals__` and `__closure__`
+    new_fn = types.FunctionType(
+        code=callable_func.__code__,
+        globals=get_new_globals(dyfunc, callable_func),
+        name=func_name,
+        argdefs=callable_func.__defaults__,
+        closure=get_new_closure(dyfunc, callable_func),
+    )
 
-    return callable_func, f.name
+    return new_fn, f.name
 
 
 def _inject_import_statements():
@@ -351,26 +484,6 @@ def _inject_import_statements():
     return '\n'.join(import_statements) + '\n'
 
 
-def recover_globals_attribute(src_obj, dst_obj):
-    attr_name = '__globals__'
-
-    src_globals = getattr(src_obj, attr_name, {})
-    dst_globals = getattr(dst_obj, attr_name, {})
-
-    for k, v in src_globals.items():
-        # ignore builtin attribute.
-        if not (k.startswith('__') and k.endswith('__')):
-            dst_globals[k] = v
-
-    # Inject source function closure into destination function globals
-    # Because the destination function is a standalone function, the original
-    # closure of the source function is compiled as LOAD_GLOBAL in the
-    # destination function.
-    src_closure = inspect.getclosurevars(src_obj)
-    for k, v in src_closure.nonlocals.items():
-        dst_globals[k] = v
-
-
 def func_to_source_code(function, dedent=True):
     """
     Transforms function into raw string of source code.
diff --git a/test/dygraph_to_static/test_ast_util.py b/test/dygraph_to_static/test_ast_util.py
@@ -92,9 +92,7 @@ def func(x):
     def test_ast2func_error(self):
         with self.assertRaises(Exception) as e:
             self.assertRaises(TypeError, ast_to_func("x = a + b", 'foo'))
-        self.assertTrue(
-            "Type of ast_root should be gast.AST or ast.AST" in str(e.exception)
-        )
+        self.assertIn("'foo' is not a Python function", str(e.exception))
 
 
 if __name__ == '__main__':
diff --git a/test/dygraph_to_static/test_capture_free_vars.py b/test/dygraph_to_static/test_capture_free_vars.py
diff --git a/test/dygraph_to_static/test_origin_info.py b/test/dygraph_to_static/test_origin_info.py
