mindspore 2.0.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/rewrite/node/node_manager.py

@@ -0,0 +1,254 @@
+# 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.
+# ============================================================================
+"""SymbolTree nodes manager."""
+import sys
+from typing import Optional
+import ast
+from .node import Node
+from .node_topological_manager import TopoManager
+from ..api.node_type import NodeType
+
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
+
+class NodeManager:
+    """
+    NodeManager saves nodes and manager nodes' topological relationship.
+    """
+    def __init__(self, node_namer):
+        """Initializer of NodeManager"""
+        self._topo_mgr = TopoManager()
+        self._nodes: {str, Node} = {}
+        self._manager_node_namer = node_namer
+        # record all tree nodes, which is used when generating codes
+        self._tree_nodes: [Node] = []
+        # head node is always point to the first node of nodes
+        self._head = None
+        # tail node is always point to the last node of nodes
+        self._tail = None
+        # nodes of Input type
+        self._inputs: [Node] = []
+        # nodes of Output type
+        self._returns: [Node] = []
+        # ast of ast.FunctionDef
+        self._ast_functiondef = None
+        # name of manager
+        self._manager_name = "OriginNodeManager"
+
+    @property
+    def node_list(self):
+        """ Get node list. """
+        nodes = []
+        node = self._head
+        while node is not None:
+            nodes.append(node)
+            node = node.get_next()
+        return nodes
+
+    @property
+    def node_count(self):
+        """Number of nodes."""
+        node_num = 0
+        node = self._head
+        while node is not None:
+            node_num = node_num + 1
+            node = node.get_next()
+        return node_num
+
+    def insert_node(self, new_node: Node, base_node: Node, before_node: bool):
+        """
+        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.
+        """
+        # update node name
+        new_node_name = self._manager_node_namer.get_name(new_node)
+        new_node.set_name(new_node_name)
+        if isinstance(new_node, NodeManager):
+            new_node.set_manager_name(new_node_name)
+        # insert node to list table
+        if base_node is None:
+            if self._nodes:
+                raise ValueError("base_node cannot be None when node inserted is not the first node.")
+            self._head = new_node
+            self._tail = new_node
+        elif before_node:
+            base_node.insert_before(new_node)
+            if self._head == base_node:
+                self._head = new_node
+        else:
+            base_node.insert_after(new_node)
+            if self._tail == base_node:
+                self._tail = new_node
+        self._add_node_to_nodes(new_node)
+        self._topo_mgr.on_insert_node(new_node)
+        new_node.set_node_manager(self)
+        # record Input nodes, Output nodes and tree nodes
+        if new_node.get_node_type() == NodeType.Output:
+            self._returns.append(new_node)
+        elif new_node.get_node_type() == NodeType.Input:
+            self._inputs.append(new_node)
+        elif new_node.get_node_type() == NodeType.Tree:
+            self._tree_nodes.append(new_node)
+
+    def erase_node(self, node: Node):
+        """
+        Erase a node from nodes.
+
+        Args:
+            node (Node): _description_
+        """
+        self._topo_mgr.on_erase_node(node)
+        for key, value in self._nodes.items():
+            if id(value) == id(node):
+                # update self._head and self._tail
+                if self._head == node:
+                    self._head = node.get_next()
+                if self._tail == node:
+                    self._tail = node.get_prev()
+                # erase node
+                self._nodes.pop(key)
+                value.isolate()
+                break
+
+    def nodes(self):
+        """
+        Get nodes.
+
+        Returns:
+            A list of nodes.
+        """
+        # If iterating nodes directly without new list, iteration may stuck caused
+        # by node topology being modified during iteration.
+        nodes = []
+        node = self._head
+        while node is not None:
+            nodes.append(node)
+            node = node.get_next()
+        return nodes
+
+    def get_node(self, node_name: str) -> Optional[Node]:
+        """
+        Get node of current NodeManager by `node_name`.
+
+        Args:
+            node_name (str): A str represents name of node as key of query.
+
+        Returns:
+            An instance of Node if found else None.
+        """
+        return self._nodes.get(node_name)
+
+    def append_python_node(self, new_node: Node):
+        """Append python node"""
+        NodeManager.insert_node(self, new_node, self._tail, False)
+
+    def get_head(self):
+        """Get head node of nodes"""
+        return self._head
+
+    def get_tail(self):
+        """Get tail node of nodes"""
+        return self._tail
+
+    def reg_observer(self, observer):
+        """Register observer to monitor code changes."""
+        self._topo_mgr.reg_observer(observer)
+        for node in self.nodes():
+            if isinstance(node, NodeManager):
+                node.reg_observer(observer)
+            if node.get_node_type() == NodeType.Tree:
+                node.symbol_tree.reg_observer(observer)
+
+    def get_tree_nodes(self):
+        """Get tree nodes inserted into symbol tree, include nodes later erased by user."""
+        tree_nodes = []
+        tree_nodes.extend(self._tree_nodes)
+        for node in self.nodes():
+            if isinstance(node, NodeManager):
+                tree_nodes.extend(node.get_tree_nodes())
+        return tree_nodes
+
+    def set_ast_functiondef(self, ast_functiondef: ast.FunctionDef):
+        """Set _ast_functiondef."""
+        self._ast_functiondef = ast_functiondef
+
+    def get_ast_functiondef(self):
+        """Get _ast_functiondef."""
+        return self._ast_functiondef
+
+    def get_inputs(self):
+        """Get _inputs"""
+        return self._inputs
+
+    def get_returns(self):
+        """Get _returns"""
+        return self._returns
+
+    def set_manager_name(self, name: str):
+        """Set _manager_name"""
+        self._manager_name = name
+
+    def get_manager_name(self):
+        """Get _manager_name"""
+        return self._manager_name
+
+    def dump(self, title="") -> str:
+        """
+        Dump topological relation.
+
+        title (str): A string as a title will be printed before dumping topological relation.
+        """
+        try:
+            from tabulate import tabulate # pylint: disable=unused-import,reportMissingModuleSource
+        except ImportError:
+            return ""
+        dump_str = "=" * 40 + title + "=" * 40 + '\n'
+        node_specs = [[
+            n.get_node_type(),
+            n.get_name(),
+            astunparse.unparse(n.get_ast()).strip(),
+            [[key, ((value[0].get_name(), value[1]) if value else ())]
+             for key, value in n.get_arg_providers().items()],
+            [[
+                key,
+                [(val[0].get_name(), val[1]) if val else ()
+                 for val in value] if value else []
+            ] for key, value in n.get_target_users().items()]
+        ] for n in NodeManager.nodes(self)]
+        dump_str += tabulate(node_specs, headers=['node type', 'name', 'codes', 'arg providers', 'target users'])
+        dump_str += '\n' + "=" * (82 + len(title)) + '\n'
+        return dump_str
+
+    def _add_node_to_nodes(self, node: Node):
+        """
+        Add `node` to `_nodes` dict.
+
+        Args:
+            node (Node): A Node to be added into `_nodes`.
+
+        Raises:
+            RuntimeError: If name of the node is duplicated.
+        """
+        node_name = node.get_name()
+        if self._nodes.get(node_name) is not None:
+            raise ValueError(f"Duplicated node name: {node_name} in"
+                             f"{self.get_name() if isinstance(self, Node) else 'construct'}")
+        self._nodes[node_name] = node

mindspore/rewrite/node/node_topological_manager.py

@@ -0,0 +1,243 @@
+# 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.
+# ============================================================================
+"""SymbolTree nodes topological relationship manager."""
+from typing import Tuple, List
+from mindspore import log as logger
+from ..api.scoped_value import ScopedValue
+from .node import Node
+from ..common.observable import Observable
+from ..common.event import Event
+
+
+class TopoManager(Observable):
+    """SymbolTree topological-relationship manager."""
+
+    @staticmethod
+    def get_node_users(node: Node) -> List[Tuple[Node, int]]:
+        """
+        Get all nodes which depend on node.
+
+        Args:
+            node (Node): An instance of node.
+
+        Returns:
+            A list of nodes represents node users.
+        """
+        results = []
+        for users in node.get_target_users().values():
+            if not users:
+                continue
+            for user in users:
+                if user not in results:
+                    results.append(user)
+        return results
+
+    @staticmethod
+    def on_update_target(node: Node, index: int, old_target: ScopedValue, new_target: ScopedValue):
+        """
+        Update node's dicts while updating target of node.
+
+        Args:
+            node (Node): An instance of Node whose target being updated.
+            arg_idx (int): An int indicates which target of node being updated.
+            old_target (ScopedValue): An instance of ScopedValue represents old target.
+            new_target (ScopedValue): An instance of ScopedValue represents new target.
+        """
+        # Update old_target provider node's target_user dict & old arg's user nodes' arg_providers dict
+        old_provider = TopoManager._get_value_provider(node, old_target)
+        if old_provider:
+            for user in node.get_target_users(index):
+                old_provider[0].append_target_users(old_provider[1], user)
+                user[0].set_arg_providers(user[1], old_provider)
+        else:
+            for user in node.get_target_users(index):
+                user[0].set_arg_providers(user[1], ())
+        # Update new_target node's target_users dict & new user nodes' arg_providers dict
+        node.get_target_users().clear()
+        provider = TopoManager._get_value_provider(node, new_target)
+        if provider:
+            TopoManager._update_target_users_by_node(node, index, provider)
+        else:
+            TopoManager._update_target_users_by_value(node, index, new_target)
+
+    @staticmethod
+    def _update_target_users_by_value(node, index, value: ScopedValue):
+        """
+        Update node's _target_users by ScopedValue when insert a new node.
+        This function is called when target is not found in previous nodes, which means a new target name is set.
+        """
+        search_node = node.get_next()
+        while search_node is not None:
+            if search_node.get_normalized_args() is not None:
+                for arg_index, arg in enumerate(search_node.get_normalized_args().values()):
+                    if arg == value:
+                        node.append_target_users(index, (search_node, arg_index))
+                        search_node.set_arg_providers(arg_index, (node, index))
+            if search_node.get_targets() is not None:
+                for _, target in enumerate(search_node.get_targets()):
+                    if target == value:
+                        return
+            search_node = search_node.get_next()
+        return
+
+    @staticmethod
+    def _update_target_users_by_node(node, index, provider: Tuple[Node, int]):
+        """
+        Update node's _target_users by previous node when insert a new node.
+        This function is called when target is found in previous nodes, which means a repeat target name is set.
+        """
+        # Args of nodes which are between node and provider should not be changed
+        # [last provider] -> no change args -> [insert node] -> need change args -> [next provider] -> no change args
+        nodes_before_insert = []
+        search_node = provider[0].get_next()
+        while search_node is not None:
+            nodes_before_insert.append(search_node)
+            if search_node == node:
+                break
+            search_node = search_node.get_next()
+        provider_target_users = provider[0].get_target_users(provider[1])
+        for user in provider_target_users[:]: # copy list by slice to support remove item during iterating
+            if user[0] not in nodes_before_insert:
+                node.append_target_users(index, user)
+                provider_target_users.remove(user)
+                user[0].set_arg_providers(user[1], (node, index))
+
+    @staticmethod
+    def _get_value_provider(node, value: ScopedValue):
+        node = node.get_prev()
+        while node is not None:
+            if node.get_targets() is not None:
+                for index, target in enumerate(node.get_targets()):
+                    if target == value:
+                        return (node, index)
+            node = node.get_prev()
+        return ()
+
+    def topo_changed(self):
+        """
+        The function is executed when an Event.TopologicalChangeEvent event is received.
+        """
+        self.changed(Event.TopologicalChangeEvent)
+
+    def on_insert_node(self, node: Node):
+        """
+        Update provider dict and consumer dict while inserting node into SymbolTree and update inputs of node by updated
+        provider dict and consumer dict.
+
+        Args:
+            node (Node): An instance of Node which been inserted into SymbolTree.
+        """
+        if node.get_normalized_args() is not None:
+            for index, arg in enumerate(node.get_normalized_args().values()):
+                provider = TopoManager._get_value_provider(node, arg)
+                if provider:
+                    node.set_arg_providers(index, provider)
+                    provider[0].append_target_users(provider[1], (node, index))
+        if node.get_targets() is not None:
+            for index, target in enumerate(node.get_targets()):
+                provider = TopoManager._get_value_provider(node, target)
+                if provider:
+                    TopoManager._update_target_users_by_node(node, index, provider)
+                else:
+                    TopoManager._update_target_users_by_value(node, index, target)
+        self.topo_changed()
+
+    def on_erase_node(self, node: Node):
+        """
+        Update provider dict and consumer dict while erasing node from SymbolTree.
+
+        Args:
+            node (Node): An instance of Node which been erased from SymbolTree.
+        """
+        prev_providers = {}
+        # Find previous node with same target of current node.
+        for index, target_users in node.get_target_users().items():
+            if not target_users:
+                continue
+            prev_provider = TopoManager._get_value_provider(node, node.get_targets()[index])
+            if not prev_provider:
+                logger.warning(f"Node {node.get_name()}'s target {index}({node.get_targets()[index]}) is used in node "
+                               f"{target_users[0][0].get_name()}'s arg {target_users[0][1]}, "
+                               f"no other node provides this target if node {node.get_name()} is erased.")
+                prev_providers[index] = None
+            else:
+                prev_providers[index] = prev_provider
+        # Update targets topological of nodes
+        for index, prev_provider in prev_providers.items():
+            for target_user in node.get_target_users(index):
+                if prev_provider is None:
+                    target_user[0].get_arg_providers().pop(target_user[1], None)
+                else:
+                    prev_provider[0].append_target_users(prev_provider[1], target_user)
+                    target_user[0].set_arg_providers(target_user[1], prev_provider)
+        # Update arguments topological of nodes
+        for _, arg_providers in node.get_arg_providers().items():
+            if not arg_providers:
+                continue
+            provider_target_users = arg_providers[0].get_target_users(arg_providers[1])
+            for target_user in reversed(provider_target_users):
+                if target_user[0] == node:
+                    provider_target_users.remove(target_user)
+        self.topo_changed()
+
+    def on_update_arg(self, node: Node, arg_idx: int, old_arg: ScopedValue, new_arg: ScopedValue):
+        """
+        Update provider dict and consumer dict while updating argument of node and update inputs of node by updated
+        provider dict and consumer dict.
+
+        Args:
+            node (Node): An instance of Node whose arguments being updated.
+            arg_idx (int): An int indicates which argument of node being updated.
+            old_arg (ScopedValue): An instance of ScopedValue represents original argument.
+            new_arg (ScopedValue): An instance of ScopedValue represents new argument.
+        """
+        # Update old arg's provider node's target_users.
+        old_provider = TopoManager._get_value_provider(node, old_arg)
+        if old_provider:
+            old_provider_target_users = old_provider[0].get_target_users(old_provider[1])
+            for target_user in reversed(old_provider_target_users):
+                if target_user[0] == node and target_user[1] == arg_idx:
+                    old_provider_target_users.remove(target_user)
+                    break
+        # Update new arg's provider node's target_users.
+        provider = TopoManager._get_value_provider(node, new_arg)
+        if provider:
+            provider[0].append_target_users(provider[1], (node, arg_idx))
+        # Update current node's arg_providers.
+        node.set_arg_providers(arg_idx, provider)
+        self.topo_changed()
+
+    def on_update_arg_by_node(self, dst_node: Node, arg_idx: int, src_node: Node, out_idx: int):
+        """
+        Update argument of 'dst_node' by another Node.
+
+        Args:
+            dst_node (Node): Node to be modified.
+            arg_idx (int): Indicate which input being modified.
+            src_node (Node): Node as new input.
+            out_idx (int): Indicate which output of 'src_node' as new input of 'dst_node'.
+        """
+        # Update old arg's provider node's target_users.
+        if arg_idx in dst_node.get_arg_providers().keys():
+            arg_provider = dst_node.get_arg_providers()[arg_idx]
+            if arg_provider:
+                provider_target_users = arg_provider[0].get_target_users(arg_provider[1])
+                if (dst_node, arg_idx) in provider_target_users:
+                    provider_target_users.remove((dst_node, arg_idx))
+        # Update new arg's provider node's target_users.
+        src_node.append_target_users(out_idx, (dst_node, arg_idx))
+        # Update current node's arg_providers.
+        dst_node.set_arg_providers(arg_idx, (src_node, out_idx))
+        self.topo_changed()

mindspore/rewrite/parsers/arguments_parser.py

@@ -15,10 +15,11 @@
 """Parse ast.arguments to input-node of SymbolTree."""
 import ast
 
-from ..parser import Parser
-from ..parser_register import reg_parser
+from .parser import Parser
+from .parser_register import reg_parser
 from ..symbol_tree import SymbolTree
 from ..common import error_str
+from ..node.node_manager import NodeManager
 
 
 class ArgumentsParser(Parser):
@@ -28,36 +29,36 @@ class ArgumentsParser(Parser):
         """Parse target type"""
         return ast.arguments
 
-    def process(self, stree: SymbolTree, node: ast.arguments):
+    def process(self, stree: SymbolTree, ast_node: ast.arguments, node_manager: NodeManager):
         """
         Parse ast.arguments and create input-node to stree.
 
         Args:
             stree (SymbolTree): symbol tree under parsing.
-            node (ast.arguments): argument node in construct.
+            ast_node (ast.arguments): ast argument node in construct.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Raises:
-            RuntimeError: Types of node.args elements are not ast.arg.
+            RuntimeError: Types of ast_node.args elements are not ast.arg.
         """
-        if hasattr(node, "posonlyargs"):
-            stree.try_append_python_node(node, node.posonlyargs)
+        if hasattr(ast_node, "posonlyargs"):
+            stree.try_append_python_node(ast_node, ast_node.posonlyargs, node_manager)
 
-        for arg in node.args:
+        for arg in ast_node.args:
             if not isinstance(arg, ast.arg):
                 raise RuntimeError(error_str(f"only support ast.arg in arguments arg, but got '{type(arg).__name__}'",
-                                             child_node=arg, father_node=node))
-            stree.append_input_node(arg, arg.arg)
-
-        if hasattr(node, "vararg"):
-            stree.try_append_python_node(node, node.vararg)
-        if hasattr(node, "kwonlyargs"):
-            stree.try_append_python_node(node, node.kwonlyargs)
-        if hasattr(node, "kw_defaults"):
-            stree.try_append_python_node(node, node.kw_defaults)
-        if hasattr(node, "kwarg"):
-            stree.try_append_python_node(node, node.kwarg)
-        if hasattr(node, "defaults"):
-            stree.try_append_python_node(node, node.defaults)
+                                             arg, ast_node))
+            stree.append_input_node(arg, arg.arg, node_manager=node_manager)
+        if hasattr(ast_node, "vararg"):
+            stree.try_append_python_node(ast_node, ast_node.vararg, node_manager)
+        if hasattr(ast_node, "kwonlyargs"):
+            stree.try_append_python_node(ast_node, ast_node.kwonlyargs, node_manager)
+        if hasattr(ast_node, "kw_defaults"):
+            stree.try_append_python_node(ast_node, ast_node.kw_defaults, node_manager)
+        if hasattr(ast_node, "kwarg"):
+            stree.try_append_python_node(ast_node, ast_node.kwarg, node_manager)
+        if hasattr(ast_node, "defaults"):
+            stree.try_append_python_node(ast_node, ast_node.defaults, node_manager)
 
 
 g_arguments_parser = reg_parser(ArgumentsParser())