[Dy2St] Cleanup unused IsControlFlowVisitor and remove shim of unwrap between py2 and py3 (#60238)

SigureMo · web-flow · commit d65cbe898f40 · 2023-12-22T02:02:11.000+08:00
diff --git a/python/paddle/jit/dy2static/convert_call_func.py b/python/paddle/jit/dy2static/convert_call_func.py
@@ -40,7 +40,7 @@
     convert_to_static,
     unwrap_decorators,
 )
-from .utils import is_builtin, is_paddle_func, unwrap
+from .utils import is_builtin, is_paddle_func
 
 __all__ = []
 
@@ -250,7 +250,7 @@ def convert_call(func):
             # `foo` will be converted into a wrapper class, suppose as `StaticFunction`.
             # And `foo.__globals__['foo']` will still return this `StaticFunction` instead of
             # `foo` function. So `isinstance(fn, StaticFunction)` is added here.
-            _origfunc = unwrap(func)
+            _origfunc = inspect.unwrap(func)
             global_functions = set()
             for fn in _origfunc.__globals__.values():
                 if inspect.isfunction(fn):
diff --git a/python/paddle/jit/dy2static/origin_info.py b/python/paddle/jit/dy2static/origin_info.py
@@ -19,7 +19,7 @@
 from paddle.base.framework import Program
 from paddle.utils import gast
 
-from .utils import ORIGI_INFO, unwrap
+from .utils import ORIGI_INFO
 
 __all__ = []
 
@@ -95,7 +95,7 @@ class OriginInfoAttacher(gast.NodeTransformer):
 
     def __init__(self, root, func):
         self.root = root
-        self.func = unwrap(func)
+        self.func = inspect.unwrap(func)
         self.filepath = inspect.getsourcefile(self.func)
         self.source_code = inspect.getsource(self.func)
         self.current_func = []
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -65,7 +65,6 @@
     prim_is_enabled,
     prim_or_cinn_is_enabled,
     type_name,
-    unwrap,
 )
 
 if TYPE_CHECKING:
@@ -134,9 +133,7 @@ def foo(x, y):
         If the conversion of A.foo happens after B.foo, it will reuse the transformed ast node of B.foo
         to speed up the conversion.
         """
-        # Note: In Python2, it will raise OSError when inspect function
-        # with decorator directly and function.__wrapped__ holds the actual function.
-        func = unwrap(func)
+        func = inspect.unwrap(func)
         source_code = func_to_source_code(func)
 
         # TODO(liym27):
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -39,7 +39,6 @@
 from paddle.utils import gast
 
 from .ast_utils import ast_to_source_code
-from .static_analysis import StaticAnalysisVisitor
 from .utils_helper import (  # noqa: F401
     DYGRAPH_MODULE_PREFIX,
     DYGRAPH_TO_STATIC_MODULE_PREFIX,
@@ -657,249 +656,6 @@ def func_to_source_code(function, dedent=True):
     return source_code
 
 
-def is_candidate_node(node):
-    """
-    Nodes with specified type will be dependent on tensor.
-    """
-    is_compare_node = isinstance(
-        node,
-        (
-            gast.Compare,
-            gast.BoolOp,
-            gast.UnaryOp,
-            gast.For,
-            gast.If,
-            gast.While,
-        ),
-    )
-    # TODO(Aurelius84): `.numpy()` may be an customized function,
-    # and should consider a more elegant way to solve this problem.
-    has_numpy_attr = ".numpy()" in ast_to_source_code(node)
-    return is_compare_node or has_numpy_attr
-
-
-def compare_with_none(node):
-    """
-    Whether the comparator of `gast.Compare` node is `None`.
-    """
-    if isinstance(node, gast.Compare):
-        for child in [node.left, node.comparators]:
-            # node.comparators is a list.
-            if isinstance(child, list):
-                child = child[0]
-            if (isinstance(child, gast.Constant) and child.value is None) or (
-                isinstance(child, gast.Name) and child.id == 'None'
-            ):
-                return True
-    return False
-
-
-class IsControlFlowVisitor(gast.NodeVisitor):
-    """
-    Judge whether the ast_node of control flow from Dygraph code dependent on paddle Tensor.
-    `ast_node` can be gast.If, gast.For, gast.While, gast.If.test(gast.Compare, gast.BoolOp, gast.UnaryOp).
-
-    If returns True,
-    gast.If.test must meet at least one of the following requirements:
-        1. involves at least one var whose type is Tensor.
-        2. the Tensor var calls `.numpy()[]` interface or Tensor.shape is [1].
-        3. involves Tensor.shape[i] and the shape[i] is unknown in compile time.
-    gast.While must meet at least one of the requirements 1 to 5:
-        4. has `break` statement.
-        5. has `continue` statement.
-    gast.For must meet at least one of the requirements 4 to 8:
-        6. calls `range` function in `for` statement and the argument of range is Tensor.
-        7. calls `enumerate` function in `for` statement and the argument of enumerate is Tensor.
-        8. the iterable varaible in `for` statement is Tensor.
-        TODO: Support non-range case
-
-    The following examples should not be considered as control_flow_if:
-        1. `if Tensor_var` or `if Tensor_var is None`
-        2. if Tensor.shape[i] is determined with fixed value (not -1 or None)
-
-    Note: pred in ConditionalBlock require variable, which means all vars should be Tensor
-          or transformed into Tensor, like fill_constant(shape=[1], dtype='int32', value=Tensor.shape[i]).
-
-    TODO: 1. need to deal with `tensor.shape[i]` which need to eval the data of shape[i],
-             because reshape_op may be called before this statement.
-    """
-
-    def __init__(
-        self, ast_node, static_analysis_visitor=None, node_var_type_map=None
-    ):
-        assert isinstance(
-            ast_node, gast.AST
-        ), "Type of input node should be gast.AST, but received %s." % type(
-            ast_node
-        )
-        self.ast_root = ast_node
-        if static_analysis_visitor is None:
-            static_analysis_visitor = StaticAnalysisVisitor(ast_node)
-        self.static_analysis_visitor = static_analysis_visitor
-        self.node_to_wrapper_map = (
-            self.static_analysis_visitor.get_node_to_wrapper_map()
-        )
-        self.node_var_type_map = node_var_type_map
-
-        self.is_control_flow_num = 0
-        self._compare_node_tenor_set = set()
-
-    def transform(self):
-        node = self.ast_root
-        if isinstance(node, gast.If):
-            self._visit_If(node)
-        elif isinstance(node, gast.For):
-            self._visit_For(node)
-        elif isinstance(node, gast.While):
-            self._visit_While(node)
-        else:
-            self.visit(node)
-        return self.is_control_flow_num > 0
-
-    def _visit_If(self, node):
-        assert isinstance(node, gast.If)
-        self.visit(node.test)
-
-    def _visit_For(self, node):
-        assert isinstance(node, gast.For)
-        if isinstance(node.iter, gast.Call):
-            # for in range(var[0]|var.numpy()[0]) or for in enumerate(var|var.numpy())
-            if isinstance(node.iter.func, gast.Name):
-                if (
-                    node.iter.func.id == "range"
-                    or node.iter.func.id == "enumerate"
-                ):
-                    for arg in node.iter.args:
-                        self.visit(arg)
-                else:
-                    return
-            # for in var.numpy()
-            elif isinstance(node.iter.func, gast.Attribute):
-                if node.iter.func.attr == 'numpy':
-                    self._visit_Call(node.iter)
-                else:
-                    return
-            else:
-                return
-        elif isinstance(node.iter, gast.Name):
-            # for in var
-            self.visit(node.iter)
-        else:
-            return
-
-        for child_node in gast.walk(node):
-            if isinstance(child_node, (gast.Continue, gast.Break)):
-                self._visit_break_continue(child_node)
-        return
-
-    def _visit_While(self, node):
-        assert isinstance(node, gast.While)
-        test = node.test
-        self.generic_visit(test)
-        for child_node in gast.walk(node):
-            if isinstance(child_node, (gast.Continue, gast.Break)):
-                self._visit_break_continue(child_node)
-
-    def _visit_break_continue(self, node):
-        assert isinstance(node, (gast.Break, gast.Continue))
-        wrapper_node = self.node_to_wrapper_map.get(node)
-        if not wrapper_node:
-            # Transformed node is not in node_to_wrapper_map
-            return
-
-        while wrapper_node.parent:
-            parent_node = wrapper_node.parent.node
-            if isinstance(parent_node, (gast.For, gast.While)):
-                if parent_node is self.ast_root:
-                    self.is_control_flow_num += 1
-                    return
-                else:
-                    return
-
-            wrapper_node = wrapper_node.parent
-
-        return
-
-    def visit_BoolOp(self, node):
-        for i, child in enumerate(node.values):
-            self.visit(child)
-        return node
-
-    def visit_Compare(self, node):
-        pre_control_flow_num = self.is_control_flow_num
-        if not compare_with_none(node):
-            self.generic_visit(node)
-            for child in gast.walk(node):
-                if isinstance(child, gast.Subscript):
-                    self._visit_Subscript(child)
-        if self.is_control_flow_num > pre_control_flow_num:
-            self._compare_node_tenor_set.add(node)
-        return node
-
-    def _visit_Subscript(self, node):
-        self.generic_visit(node)
-        if hasattr(node, 'value') and isinstance(node.value, gast.Call):
-            self._visit_Call(node.value)
-        return node
-
-    def _visit_Call(self, node):
-        assert isinstance(node, gast.Call)
-        if isinstance(node.func, gast.Attribute):
-            attr_node = node.func
-            if attr_node.attr == 'numpy':
-                self.is_control_flow_num += 1
-
-    def visit_Call(self, node):
-        self._visit_Call(node)
-        if is_paddle_api(node):
-            self.is_control_flow_num += 1
-        return node
-
-    def visit_Name(self, node):
-        if self._is_node_with_tensor(node, node.id):
-            self.is_control_flow_num += 1
-        return node
-
-    def visit_Constant(self, node):
-        if self._is_node_with_tensor(node, node.value):
-            self.is_control_flow_num += 1
-        return node
-
-    def _is_node_with_tensor(self, node, name_id):
-        # Look up the node_var_type_map by name_id.
-        if self.node_var_type_map:
-            if name_id and isinstance(name_id, str):
-                var_type = self.node_var_type_map.get(name_id, None)
-                if var_type and var_type & NodeVarType.TENSOR_TYPES:
-                    return True
-        # if not found, look up the node_to_wrapper_map by node.
-        wrapper_node = self.node_to_wrapper_map.get(node, None)
-        if wrapper_node is not None:
-            if wrapper_node.node_var_type & NodeVarType.TENSOR_TYPES:
-                return True
-
-        return False
-
-    def get_compare_nodes_with_tensor(self):
-        return self._compare_node_tenor_set
-
-
-# NOTE: inspect.unwrap() exits in PY3 but not in PY2.
-def unwrap(func):
-    """
-    Returns the object wrapped by decorators.
-    """
-
-    def _is_wrapped(f):
-        return hasattr(f, '__wrapped__')
-
-    unwrapped_f = func
-    while _is_wrapped(unwrapped_f):
-        unwrapped_f = unwrapped_f.__wrapped__
-
-    return unwrapped_f
-
-
 def input_specs_compatible(src_input_specs, desired_input_specs):
     """
     Returns True if the two input specs are compatible, otherwise False.
diff --git a/test/dygraph_to_static/test_legacy_error.py b/test/dygraph_to_static/test_legacy_error.py
@@ -21,7 +21,6 @@
 import paddle
 from paddle import base
 from paddle.jit.dy2static import error
-from paddle.jit.dy2static.origin_info import unwrap
 
 
 def inner_func():
@@ -55,14 +54,14 @@ def func_error_in_runtime(x):
     return x
 
 
-@unwrap
+@inspect.unwrap
 @paddle.jit.to_static(full_graph=True)
 def func_decorated_by_other_1():
     return 1
 
 
 @paddle.jit.to_static(full_graph=True)
-@unwrap
+@inspect.unwrap
 def func_decorated_by_other_2():
     return 1
 
@@ -172,7 +171,7 @@ def setUp(self):
         self.set_input()
         self.set_func()
         self.set_func_call()
-        self.filepath = inspect.getfile(unwrap(self.func_call))
+        self.filepath = inspect.getfile(inspect.unwrap(self.func_call))
         self.set_exception_type()
         self.set_message()
 
@@ -297,7 +296,7 @@ class TestErrorStaticLayerCallInCompiletime_3(
     def setUp(self):
         self.reset_flags_to_default()
         self.set_func_call()
-        self.filepath = inspect.getfile(unwrap(self.func_call))
+        self.filepath = inspect.getfile(inspect.unwrap(self.func_call))
         self.set_exception_type()
         self.set_message()
 
@@ -354,7 +353,7 @@ class TestJitSaveInCompiletime(TestErrorBase):
     def setUp(self):
         self.reset_flags_to_default()
         self.set_func_call()
-        self.filepath = inspect.getfile(unwrap(self.func_call))
+        self.filepath = inspect.getfile(inspect.unwrap(self.func_call))
         self.set_exception_type()
         self.set_message()
 
diff --git a/test/dygraph_to_static/test_origin_info.py b/test/dygraph_to_static/test_origin_info.py
@@ -23,7 +23,6 @@
     OriginInfo,
     attach_origin_info,
     create_and_update_origin_info_map,
-    unwrap,
 )
 from paddle.jit.dy2static.utils import ast_to_func
 from paddle.utils import gast
@@ -56,7 +55,7 @@ def decorated_func2(x):
 class TestOriginInfo(unittest.TestCase):
     def setUp(self):
         self.set_test_func()
-        self.dygraph_func = unwrap(self.func)
+        self.dygraph_func = inspect.unwrap(self.func)
         self.dygraph_filepath = inspect.getfile(self.dygraph_func)
         self.source_code = inspect.getsource(self.dygraph_func)
         lines, self.start_lineno = inspect.getsourcelines(self.dygraph_func)
