mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/rewrite/ast_transformers/flatten_recursive_stmt.py

@@ -14,14 +14,20 @@
 # ============================================================================
 """Ast optimizer for flatten recursive call."""
 
-from typing import Any, Tuple
+import sys
+from typing import Any, Tuple, List
+import keyword
 import ast
-from ast import FunctionDef
-import astunparse
 
 from mindspore import log as logger
 from ..common import error_str
 
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
+
+FLATTEN_BLACK_LIST = ["set_vertex_attr",]
 
 class FlattenRecursiveStmt(ast.NodeTransformer):
     """Ast optimizer for flatten recursive call."""
@@ -40,17 +46,35 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
             ast.BoolOp: ["values"],
             ast.UnaryOp: ["operand"],
             ast.Compare: ["left", "comparators"],
+            ast.If: ["test"]
         }
+        self._transform_functions = []
+        self._transform_if = False
+        self._symbol_tree = None # Used to get unique name
 
     @staticmethod
-    def _generate_target_name(node: ast.AST, target_names):
+    def _check_flatten_black_list(node: ast.AST):
+        """Check whether node in flatten black list"""
+        func_name = ""
+        # Get func name of node
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name):
+                func_name = node.func.id
+            elif isinstance(node.func, ast.Attribute):
+                func_name = node.func.attr
+        # Check func name of node
+        if func_name and func_name in FLATTEN_BLACK_LIST:
+            return True
+        return False
+
+    def _generate_target_name(self, node: ast.AST, target_names):
         """Generate unique target name."""
         if isinstance(node, ast.Call):
             func = node.func
             if isinstance(func, ast.Name):
-                target_name = func.id
+                target_name = func.id + "_var"
             elif isinstance(func, ast.Attribute):
-                target_name = func.attr
+                target_name = func.attr + "_var"
             else:
                 logger.info("unhandled type of func of ast.Call while generating new target name: %s ", type(func))
                 target_name = "function"
@@ -67,30 +91,33 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
         else:
             logger.info("unhandled type of node while generating new target name: %s ", type(node))
             target_name = type(node).__name__.lower() + "_var"
+        # avoid python keyword
+        if keyword.iskeyword(target_name):
+            target_name = target_name + "_var"
         suffix = 0
         result = target_name
         while result in target_names:
             suffix += 1
             result = f"{target_name}_{suffix}"
+        if self._symbol_tree:
+            result = self._symbol_tree.unique_name(result)
         target_names.append(result)
         return result
 
-    @staticmethod
-    def _create_new_assign_node(node: ast.AST, target_names) -> Tuple[str, ast.AST]:
+    def _create_new_assign_node(self, node: ast.AST, target_names) -> Tuple[str, ast.AST]:
         """Create new assign node to be inserted into ast.FunctionDef."""
         if isinstance(node, (ast.Name, ast.Constant, ast.Num, ast.Str, ast.NameConstant, ast.Bytes, ast.Ellipsis)):
             return "", node
-        new_target_name = FlattenRecursiveStmt._generate_target_name(node, target_names)
+        new_target_name = self._generate_target_name(node, target_names)
         return new_target_name, ast.Assign(targets=[ast.Name(id=new_target_name, ctx=ast.Store())], value=node)
 
-    @staticmethod
-    def _flatten_list(node_list, target_names):
+    def _flatten_list(self, node_list, target_names):
         """Flatten a list of node."""
         results = list()
         new_list = list()
         for node in node_list:
             if isinstance(node, ast.Call):
-                new_target, new_node = FlattenRecursiveStmt._create_new_assign_node(node, target_names)
+                new_target, new_node = self._create_new_assign_node(node, target_names)
                 results.append(new_node)
                 new_list.append(ast.Name(id=new_target, ctx=ast.Load()))
             else:
@@ -99,6 +126,8 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
 
     def _flatten_statement(self, node: ast.AST, target_names) -> [ast.AST]:
         """Flatten recursive statement according to different node type."""
+        if FlattenRecursiveStmt._check_flatten_black_list(node):
+            return []
         flatten_config = self._flatten_table.get(type(node))
         if flatten_config is None:
             return []
@@ -108,21 +137,25 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
             if isinstance(todos, list):
                 new_list = []
                 for todo in todos:
-                    new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(todo, target_names)
+                    # Starred expression(e.g. *args) cannot be flatten.
+                    if isinstance(todo, ast.Starred):
+                        new_list.append(todo)
+                        continue
+                    new_target_name, new_node = self._create_new_assign_node(todo, target_names)
                     if id(new_node) == id(todo):
                         new_list.append(todo)
                     else:
                         new_list.append(ast.Name(id=new_target_name, ctx=ast.Load()))
                         results.append(new_node)
                     if isinstance(todo, (ast.Tuple, tuple)):
-                        _res, _new_list = FlattenRecursiveStmt._flatten_list(new_node.value.elts, [new_target_name])
+                        _res, _new_list = self._flatten_list(new_node.value.elts, [new_target_name])
                         new_node.value.elts = _new_list
                         results.extend(_res)
                 setattr(node, todo_name, new_list)
             elif isinstance(todos, dict):
                 new_dict = []
                 for key, value in todos:
-                    new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(value, target_names)
+                    new_target_name, new_node = self._create_new_assign_node(value, target_names)
                     if id(new_node) == id(value):
                         new_dict[key] = value
                     else:
@@ -130,16 +163,15 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                         results.append(new_node)
                 setattr(node, todo_name, new_dict)
             else:
-                new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(todos, target_names)
+                new_target_name, new_node = self._create_new_assign_node(todos, target_names)
                 if id(new_node) != id(todos):
                     setattr(node, todo_name, ast.Name(id=new_target_name, ctx=ast.Load()))
                     results.append(new_node)
         return results
 
-    def _fill_in_original_target_names(self, target_names, node):
-        """Fill in original target names before getting unique names."""
-        for function_index in range(len(node.body)):
-            child = node.body[function_index]
+    def _save_target_names(self, target_names, ast_body: List[ast.AST]):
+        """Saving target names in ast_body before getting unique names."""
+        for child in ast_body:
             if isinstance(child, (ast.Assign, ast.Expr)):
                 child_value = child.value
             else:
@@ -151,7 +183,7 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                 continue
             targets = child.targets
             for target in targets:
-                if not isinstance(target, (ast.Name, ast.Tuple)):
+                if not isinstance(target, (ast.Name, ast.Tuple, ast.List)):
                     raise RuntimeError(
                         error_str(f"currently only support ast.Name targets, but got ast type "
                                   f"'{type(target).__name__}'", child_node=target, father_node=child))
@@ -159,7 +191,7 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                     target_name = target.id
                     if target_name not in target_names:
                         target_names.append(target_name)
-                elif isinstance(target, ast.Tuple):
+                elif isinstance(target, (ast.Tuple, ast.List)):
                     for elt in target.elts:
                         if not isinstance(elt, ast.Name):
                             raise RuntimeError(
@@ -170,47 +202,66 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                         if target_name not in target_names:
                             target_names.append(target_name)
 
-    def visit_FunctionDef(self, node: FunctionDef) -> Any:
-        """Traverse construct node and flatten recursive nodes."""
-        if node.name != "construct":
-            return node
-
+    def _visit_ast_bodies(self, ast_body: List[ast.AST]):
+        """Traverse nodes in ast_body and flatten nodes recursive."""
         target_names = []
-        self._fill_in_original_target_names(target_names, node)
-        index = len(node.body) - 1
+        self._save_target_names(target_names, ast_body)
+        index = len(ast_body) - 1
         while index >= 0:
-            child = node.body[index]
+            child = ast_body[index]
             if isinstance(child, ast.Assign):
                 stmt = child.value
             elif isinstance(child, ast.If):
                 if isinstance(child.body[0], ast.Return) and not isinstance(child.test, ast.UnaryOp):
-                    if isinstance(child.body[0].value, ast.Call):
-                        if_body = child.body
-                        if_func = if_body[0].value
-                        expr = "x = " + astunparse.unparse(if_func)
-                        if_body = ast.parse(expr)
-                        if_body = if_body.body+ast.parse("return x").body
-                        child.body = if_body
-                        stmt = child
-                    else:
-                        stmt = child
-                else:
-                    stmt = child
+                    if not isinstance(child.body[0].value, (ast.Name, ast.Constant)):
+                        return_val_ast = child.body[0].value
+                        return_name = self._generate_target_name(return_val_ast, target_names)
+                        new_assign_code = f"{return_name} = {astunparse.unparse(return_val_ast)}"
+                        new_assign_ast = ast.parse(new_assign_code).body[0]
+                        new_return_ast = ast.parse(f"return {return_name}").body[0]
+                        child.body = [new_assign_ast, new_return_ast]
+                stmt = child
             elif isinstance(child, ast.Expr):
                 stmt = child.value
             else:
                 stmt = child
             results = self._flatten_statement(stmt, target_names)
             if results:
-                results.reverse()
-                for result in results:
-                    node.body.insert(index, result)
+                for result in reversed(results):
+                    ast_body.insert(index, result)
                     index += 1
             index -= 1
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: # pylint: disable=invalid-name
+        """Traverse nodes in _transform_functions and flatten recursive nodes."""
+        if node.name not in self._transform_functions:
+            return node
+        self._visit_ast_bodies(node.body)
         return node
 
-    def transform(self, ast_root):
+    def visit_If(self, node: ast.If) -> Any: # pylint: disable=invalid-name
+        """Traverse nodes in if node and flatten recursive nodes."""
+        if not self._transform_if:
+            return node
+        self._visit_ast_bodies(node.body)
+        if node.orelse:
+            self._visit_ast_bodies(node.orelse)
+        return node
+
+    def transform(self, ast_root, transform_functions=None, stree=None):
         """Interface of FlattenRecursiveStmt."""
+        self._transform_functions = transform_functions if transform_functions else ["construct"]
+        self._transform_if = False
+        self._symbol_tree = stree
         ast_root = self.visit(ast_root)
         ast_root = ast.fix_missing_locations(ast_root)
         return ast_root
+
+    def transform_if(self, ast_if, stree=None):
+        """Interface of FlattenRecursiveStmt."""
+        self._transform_functions = []
+        self._transform_if = True
+        self._symbol_tree = stree
+        ast_if = self.visit(ast_if)
+        ast_if = ast.fix_missing_locations(ast_if)
+        return ast_if

mindspore/rewrite/common/rewrite_elog.py

@@ -14,8 +14,12 @@
 # ============================================================================
 """Error Log for Rewrite."""
 
+import sys
 import ast
-import astunparse
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
 
 
 def error_str(reason: str, child_node: ast.expr = None, father_node: ast.expr = None) -> str:

mindspore/rewrite/namer.py

@@ -14,9 +14,9 @@
 # ============================================================================
 """Unique name producer for target, name of node, class name, etc."""
 
-from typing import Union
+from typing import Union, Tuple
 
-from .node import Node
+from .node.node import Node
 from .api.node_type import NodeType
 
 
@@ -33,7 +33,7 @@ class Namer:
         self._names: {str: int} = {}
 
     @staticmethod
-    def _real_name(name: str) -> str:
+    def _real_name(name: str) -> Tuple[str, int]:
         """
         Find real name. For example, "name1" is the real name of "name1_10", "name1" is the real name of "name1_10_3".
         If not find real name before find unique name, unique name may be not unique. For example:
@@ -47,21 +47,21 @@ class Namer:
             name (str): Origin name which may have digit prefix.
 
         Returns:
-            A string represents real-name.
+            A string represents real-name and a int represents suffix.
         """
         if name == '_':
-            return name
+            return name, None
         pos = name.rfind("_")
-        if pos == -1:
-            return name
+        if pos == -1 or pos == len(name) - 1:
+            return name, None
         digit = True
         for i in range(pos + 1, len(name)):
             if not name[i].isdigit():
                 digit = False
                 break
         if digit:
-            return Namer._real_name(name[:pos])
-        return name
+            return name[:pos], int(name[pos + 1:])
+        return name, None
 
     def get_name(self, origin_name: str) -> str:
         """
@@ -75,13 +75,31 @@ class Namer:
         """
         if origin_name == '_':
             return origin_name
-        origin_name = Namer._real_name(origin_name)
-        number = self._names.get(origin_name)
+        real_name, suffix_idx = Namer._real_name(origin_name)
+        name = origin_name
+        number = self._names.get(name)
         if number is None:
-            self._names[origin_name] = 1
-            return origin_name
-        self._names[origin_name] = number + 1
-        return f"{origin_name}_{number}"
+            self._names[name] = 1
+            if not suffix_idx:
+                # When _names is {x:2} and origin_name is y,
+                # origin_name is not in _names and can be returned.
+                return name
+            if suffix_idx and not self._names.get(real_name, -1) >= suffix_idx:
+                # When _names is {x:2} and origin_name is x_3,
+                # return x_3 and update _names to {x:2, x_3:1}
+                return name
+            # When _names is {x:2} and origin_name is x_1,
+            # set new_name to x_1_1 by set number to 1, and continue to update name.
+            number = 1
+        while True:
+            new_name = f"{name}_{number}"
+            number += 1
+            self._names[name] = number
+            # When _names is {x:2, x_3:1}, origin_name is x and number is update to 3,
+            # new_name x_3 is conflict with key x_3, so this new_name need to be skipped.
+            if new_name in self._names.keys():
+                continue
+            return new_name
 
     def add_name(self, name: str):
         """
@@ -93,44 +111,25 @@ class Namer:
         Raises:
             RuntimeError: If name is not unique in current namer.
         """
-        real_name = Namer._real_name(name)
-        number = self._names.get(real_name)
-        if number is not None:
-            raise RuntimeError("name duplicated: ", name)
-        self._names[name] = 1
+        if self._names.get(name) is None:
+            self._names[name] = 1
 
 
 class TargetNamer(Namer):
     """
     Used for unique-ing targets of node.
     """
-    def __init__(self):
-        super().__init__()
-        self._origin_name_map = {}
-
-    def get_name(self, origin_name: str) -> str:
-        ret = super(TargetNamer, self).get_name(origin_name)
-        self._origin_name_map[origin_name] = ret
-        return ret
-
-    def add_name(self, name: str):
-        super(TargetNamer, self).add_name(name)
-        self._origin_name_map[name] = name
-
-    def get_real_arg(self, origin_arg: str) -> str:
+    def get_unique_name(self, origin_name: str) -> str:
         """
-        Get real argument from 'origin_arg' because target of node which produces 'origin_arg' may be change for unique.
+        Get unique name from 'origin_name'.
 
         Args:
-            origin_arg (str): Origin argument string which may be undefined cause to target unique-lize.
+            origin_name (str): Origin name which may be duplicated.
 
         Returns:
-            A string represents real argument name.
+            A string represents unique-name.
         """
-        real_arg = self._origin_name_map.get(origin_arg)
-        if real_arg:
-            return real_arg
-        return origin_arg
+        return super(TargetNamer, self).get_name(origin_name)
 
 
 class NodeNamer(Namer):
@@ -153,19 +152,20 @@ class NodeNamer(Namer):
         if isinstance(node_or_name, Node):
             origin_name = node_or_name.get_name()
             if origin_name is None or not origin_name:
-                if node_or_name.get_node_type() in (NodeType.CallCell, NodeType.CallPrimitive, NodeType.CallFunction):
-                    if not isinstance(node_or_name, Node):
-                        raise TypeError("node_or_name should be Node, got: ", type(node_or_name))
-                    targets = node_or_name.get_targets()
-                    # return node and head node will not call this method
-                    if not targets:
-                        raise RuntimeError("node should has at lease one target except return-node and head-node: ",
-                                           node_or_name)
-                    origin_name = str(targets[0].value)
+                if node_or_name.get_node_type() in (NodeType.CallCell, NodeType.CallPrimitive, NodeType.CallFunction,
+                                                    NodeType.Tree):
+                    origin_name = type(node_or_name.get_instance()).__name__
                 elif node_or_name.get_node_type() == NodeType.Python:
-                    origin_name = node_or_name.get_instance().__name__
+                    if node_or_name.get_instance():
+                        origin_name = type(node_or_name.get_instance()).__name__
+                    else:
+                        origin_name = "python_node"
                 elif node_or_name.get_node_type() == NodeType.Input:
                     origin_name = "parameter"
+                elif node_or_name.get_node_type() == NodeType.Output:
+                    origin_name = "return"
+                elif node_or_name.get_node_type() == NodeType.MathOps:
+                    origin_name = "math_ops"
                 else:
                     raise RuntimeError("Node type unsupported:", node_or_name.get_node_type())
         elif isinstance(node_or_name, str):

mindspore/rewrite/namespace.py

@@ -21,12 +21,21 @@ _ms_nn_ns = CellNamespace('mindspore.nn')
 _ms_ops_ns = CellNamespace('mindspore.ops.operations')
 _ms_functional_ns = CellNamespace('mindspore.ops.functional')
 
-
-def is_subtree(cls_name):
-    """Determine whether 'cls_name' is a subtree."""
-    if cls_name == "QuantizeWrapperCell":
+# Elements in _subtree_black_list will not be converted to symbol tree.
+# Only str and types are stored in _subtree_black_list.
+_subtree_black_list = ["QuantizeWrapperCell",]
+
+def is_subtree(cls_inst):
+    """Determine whether 'cls_inst' is a subtree."""
+    cls_name = type(cls_inst).__name__
+    black_list_types = tuple([elem for elem in _subtree_black_list if not isinstance(elem, str)])
+    if cls_name in _subtree_black_list or isinstance(cls_inst, black_list_types):
+        return False
+    if cls_name in _ms_common_ns and isinstance(cls_inst, _ms_common_ns[cls_name]):
+        return False
+    if cls_name in _ms_nn_ns and isinstance(cls_inst, _ms_nn_ns[cls_name]):
         return False
-    if cls_name in _ms_common_ns or cls_name in _ms_nn_ns or cls_name in _ms_ops_ns:
+    if cls_name in _ms_ops_ns and isinstance(cls_inst, _ms_ops_ns[cls_name]):
         return False
 
     return True

mindspore/{ops/bprop_mindir → rewrite/node}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
+# Copyright 2022 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-
-"""bprop mindir."""
-from __future__ import absolute_import
+"""
+SymbolTree node
+"""
+from mindspore.rewrite.node.node import Node, TreeNode
+from mindspore.rewrite.node.node_manager import NodeManager
+from mindspore.rewrite.node.call_function import CallFunction
+from mindspore.rewrite.node.cell_container import CellContainer
+from mindspore.rewrite.node.control_flow import ControlFlow

mindspore/rewrite/node/call_function.py

@@ -0,0 +1,79 @@
+# Copyright 2022 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+"""CallFunction Node."""
+import ast
+from .node import Node
+from .node_manager import NodeManager
+from ..api.scoped_value import ScopedValue
+from ..api.node_type import NodeType
+from ..ast_helpers import AstModifier
+
+
+class CallFunction(Node, NodeManager):
+    """CallFunction is used for class internal function."""
+    def __init__(self, targets: [ScopedValue], func_name: ScopedValue, args: [ScopedValue],
+                 kwargs: {str: ScopedValue}, node_name: str, ast_node: ast.AST, ast_functiondef: ast.FunctionDef,
+                 stree, instance):
+        """
+        Constructor of CallFunction.
+
+        Args:
+            targets (list[ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class.
+            args (list[ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class.
+            kwargs (dict{str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class.
+            func_name ([ScopedValue, optional]): An instance of ScopedValue. See detail in docstring of Node class.
+            node_name (str): A string represents name of node. Name of node will be unique when inserted into
+                SymbolTree. Name of node also used as field name in network class.
+            ast_node (ast.AST): An instance of ast.AST represents corresponding node in ast.
+            ast_functiondef (ast.FunctionDef): An instance of ast.FunctionDef represents corresponding function
+                definition in ast.
+            stree (SymbolTree): Symbol tree used to get node_namer.
+            instance: Object in network corresponding to this node.
+        """
+        if isinstance(func_name, str):
+            func_name = ScopedValue.create_naming_value(func_name)
+        Node.__init__(self, NodeType.CallFunction, ast_node, targets, func_name, args, kwargs, node_name, instance)
+        NodeManager.__init__(self, stree.get_node_namer())
+        NodeManager.set_ast_functiondef(self, ast_functiondef)
+        NodeManager.set_manager_name(self, func_name.value)
+
+    def erase_node(self, node):
+        """Erase node from CallFunction."""
+        NodeManager.erase_node(self, node)
+        # erase asts
+        ret = AstModifier.erase_ast_from_function(self.get_ast_functiondef(), node.get_ast())
+        if not ret:
+            raise ValueError(f"erase node failed, node {node.get_name()} not in function ast tree.")
+
+    def insert_node(self, new_node: Node, base_node: Node, before_node: bool, insert_to_ast: bool = True):
+        """
+        Insert a node before or after base_node.
+
+        Args:
+            new_node (Node): Node to be inserted.
+            base_node (Node): New node will be inserted before or after base_node.
+            before_node (bool): Indicate whether new node is inserted before base_node.
+            insert_to_ast (bool): Indicate whether ast nodes need to be updated.
+        """
+        NodeManager.insert_node(self, new_node, base_node, before_node)
+        if insert_to_ast:
+            stree = self.get_belong_symbol_tree()
+            stree.insert_to_ast_while_insert_node(new_node, base_node, before_node, self)
+
+    def set_belong_symbol_tree(self, symbol_tree):
+        """Set the symbol tree to which node belongs."""
+        self._belong_tree = symbol_tree
+        for node in self.nodes():
+            node.set_belong_symbol_tree(symbol_tree)