Add Basic Node Var Type Analysis (#22603)

1. Move AstNodeWrapper, StaticAnalysisVisitor to a new python file: static_analysis.py
2. Add basic node var type analysis

python/paddle/fluid/dygraph/dygraph_to_static/__init__.py

@@ -17,5 +17,9 @@ from __future__ import print_function
 from . import ast_transformer
 from .ast_transformer import *
 
+from . import static_analysis
+from .static_analysis import *
+
 __all__ = []
 __all__ += ast_transformer.__all__
+__all__ += static_analysis.__all__

python/paddle/fluid/dygraph/dygraph_to_static/ast_transformer.py

@@ -16,87 +16,9 @@ from __future__ import print_function
 
 import gast
 
-__all__ = ['AstNodeWrapper', 'DygraphToStaticAst', 'StaticAnalysisVisitor']
+from .static_analysis import AstNodeWrapper, StaticAnalysisVisitor
 
-
-class NodeVarType(object):
-    """
-    Enum class of python variable types. We have to know some variable types
-    during compile time to transfer AST. For example, a string variable and a
-    tensor variable in if clause may lead to different conversion from dygraph
-    to static graph.
-    """
-    UNKNOWN = 0  # Reserve for AST nodes have not known the type
-    STATEMENT = 1  # For nodes representing statement (non-variable type)
-    PADDLE_DYGRAPH_API = 2
-    PADDLE_CONTROL_IF = 3
-    PADDLE_CONTROL_WHILE = 4
-    PADDLE_CONTROL_FOR = 5
-
-    NONE = 100
-    INT = 101
-    FLOAT = 102
-    STRING = 103
-    TENSOR = 104
-
-
-class AstNodeWrapper(object):
-    """
-    Wrapper for python ast.node. We need a node wrapper because ast.node
-    doesn't store all required information when we are transforming AST.
-    We should collect additional information which the actual transformation
-    needs.
-    """
-
-    def __init__(self, node):
-        self.node = node
-        self.parent = None
-        self.children = []
-        self.node_var_type = NodeVarType.UNKNOWN
-
-
-class StaticAnalysisVisitor(object):
-    """
-    A class that does static analysis
-    """
-
-    def __init__(self, ast_root=None):
-        if ast_root is not None:
-            self.run(ast_root)
-
-    def run(self, ast_root):
-        self.node_wrapper_root = None
-        self.ancestor_wrappers = []
-        self.node_to_wrapper_map = {}
-        self.dfs_visit(ast_root)
-
-    def dfs_visit(self, node):
-        # AST reuses some ast.nodes, such as Param node of expr_context
-        if node not in self.node_to_wrapper_map:
-            cur_wrapper = AstNodeWrapper(node)
-            self.node_to_wrapper_map[node] = cur_wrapper
-        else:
-            cur_wrapper = self.node_to_wrapper_map[node]
-
-        if self.node_wrapper_root is None:
-            self.node_wrapper_root = cur_wrapper
-
-        if len(self.ancestor_wrappers) != 0:
-            last_wrapper = self.ancestor_wrappers[-1]
-            last_wrapper.children.append(cur_wrapper)
-            cur_wrapper.parent = last_wrapper
-
-        self.ancestor_wrappers.append(cur_wrapper)
-        for child in gast.iter_child_nodes(node):
-            self.dfs_visit(child)
-        self.ancestor_wrappers.pop()
-        return cur_wrapper.node_var_type
-
-    def get_node_wrapper_root(self):
-        return self.node_wrapper_root
-
-    def get_node_to_wrapper_map(self):
-        return self.node_to_wrapper_map
+__all__ = ['DygraphToStaticAst']
 
 
 class DygraphToStaticAst(gast.NodeTransformer):

python/paddle/fluid/tests/unittests/test_ast_transformer_static_analysis.py

@@ -14,18 +14,23 @@
 
 from __future__ import print_function
 
-import ast
+import gast
 import inspect
+import numpy as np
+import paddle.fluid as fluid
 import unittest
 
-from paddle.fluid.dygraph.dygraph_to_static import AstNodeWrapper, StaticAnalysisVisitor
+from paddle.fluid.dygraph.dygraph_to_static import AstNodeWrapper, NodeVarType, StaticAnalysisVisitor
 
 
-def func_to_test_1(a, b):
+def func_to_test1(a, b):
     return a + b
 
 
-def func_to_test_2(x):
+result_var_type1 = {}
+
+
+def func_to_test2(x):
     for i in range(10):
         x += i
     m = 3
@@ -37,6 +42,57 @@ def func_to_test_2(x):
         return x
 
 
+result_var_type2 = {'m': NodeVarType.INT}
+
+
+def func_to_test3():
+    a = 1
+    b = 3.0
+    c = a * b
+    d = True + c
+    e = a < b
+    f = 9 * (a * 4)
+    g = "dddy"
+    h = None
+    i = False
+    j = None + 1
+
+
+result_var_type3 = {
+    'a': NodeVarType.INT,
+    'b': NodeVarType.FLOAT,
+    'c': NodeVarType.FLOAT,
+    'd': NodeVarType.FLOAT,
+    'e': NodeVarType.BOOLEAN,
+    'f': NodeVarType.INT,
+    'g': NodeVarType.STRING,
+    'h': NodeVarType.NONE,
+    'i': NodeVarType.BOOLEAN,
+    'j': NodeVarType.UNKNOWN
+}
+
+
+def func_to_test4():
+    with fluid.dygraph.guard():
+        a = np.random.uniform(0.1, 1, [1, 2])
+        b = 1 + a
+        c = fluid.dygraph.to_variable(b)
+        d = (c + 1) * 0.3
+
+
+result_var_type4 = {
+    'a': NodeVarType.NUMPY_NDARRAY,
+    'b': NodeVarType.NUMPY_NDARRAY,
+    'c': NodeVarType.TENSOR,
+    'd': NodeVarType.TENSOR
+}
+
+test_funcs = [func_to_test1, func_to_test2, func_to_test3, func_to_test4]
+result_var_type = [
+    result_var_type1, result_var_type2, result_var_type3, result_var_type4
+]
+
+
 class TestStaticAnalysis(unittest.TestCase):
     def _check_wrapper(self, wrapper, node_to_wrapper_map):
         self.assertEqual(node_to_wrapper_map[wrapper.node], wrapper)
@@ -44,7 +100,7 @@ class TestStaticAnalysis(unittest.TestCase):
             self.assertTrue(wrapper in wrapper.parent.children)
 
         children_ast_nodes = [
-            child for child in ast.iter_child_nodes(wrapper.node)
+            child for child in gast.iter_child_nodes(wrapper.node)
         ]
         self.assertEqual(len(wrapper.children), len(children_ast_nodes))
         for child in wrapper.children:
@@ -52,15 +108,30 @@ class TestStaticAnalysis(unittest.TestCase):
             self._check_wrapper(child, node_to_wrapper_map)
 
     def test_construct_node_wrapper(self):
-        for func in [func_to_test_1, func_to_test_2]:
+        for func in test_funcs:
             test_source_code = inspect.getsource(func)
-            ast_root = ast.parse(test_source_code)
-
+            ast_root = gast.parse(test_source_code)
             visitor = StaticAnalysisVisitor(ast_root)
             wrapper_root = visitor.get_node_wrapper_root()
             node_to_wrapper_map = visitor.get_node_to_wrapper_map()
             self._check_wrapper(wrapper_root, node_to_wrapper_map)
 
+    def test_var_env(self):
+        for i in range(4):
+            func = test_funcs[i]
+            var_type = result_var_type[i]
+            test_source_code = inspect.getsource(func)
+            ast_root = gast.parse(test_source_code)
+            print(gast.dump(ast_root))
+            visitor = StaticAnalysisVisitor(ast_root)
+            var_env = visitor.get_var_env()
+            scope_var_type = var_env.get_scope_var_type()
+            self.assertEqual(len(scope_var_type), len(var_type))
+            for name in scope_var_type:
+                print("Test var name %s" % (name))
+                self.assertTrue(name in var_type)
+                self.assertEqual(scope_var_type[name], var_type[name])
+
 
 if __name__ == '__main__':
     unittest.main()

python/requirements.txt

@@ -18,3 +18,4 @@ decorator
 prettytable
 objgraph
 gast
+astor