mindspore 2.4.0__cp310-cp310-macosx_10_15_x86_64.whl

mindspore/rewrite/node/node_manager.py

@@ -0,0 +1,283 @@
+# 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."""
+from typing import Optional, Union
+import ast
+from .node import Node
+from .node_topological_manager import TopoManager
+from ..api.node_type import NodeType
+from ..api.scoped_value import ScopedValue
+
+
+class NodeManager:
+    """
+    NodeManager saves nodes and manager nodes' topological relationship.
+    """
+    def __init__(self):
+        """Initializer of NodeManager"""
+        self._topo_mgr = TopoManager()
+        self._nodes: {str, Node} = {}
+        self._manager_node_namer = None
+        # 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 node manager
+        # SymbolTree    -> ast.FunctionDef
+        # CallFunction  -> ast.FunctionDef
+        # ControlFlow   -> list
+        # CellContainer -> ast.Assign
+        self._node_manager_ast: Union[ast.AST, list] = 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): Node to be erased.
+        """
+        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_manager_ast(self, node_manager_ast: Union[ast.AST, list]):
+        """Set _node_manager_ast."""
+        self._node_manager_ast = node_manager_ast
+
+    def get_manager_ast(self):
+        """Get _node_manager_ast."""
+        return self._node_manager_ast
+
+    def get_input_nodes(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 on_update_arg(self, node: Node, arg_idx: int, old_arg: ScopedValue, new_arg: ScopedValue):
+        """
+        Update node topological when node arg is modified.
+        """
+        self._topo_mgr.on_update_arg(node, arg_idx, old_arg, new_arg)
+
+    def on_update_arg_by_node(self, dst_node: Node, arg_idx: int, src_node: Node, out_idx: int):
+        """
+        Update node topological when node arg is modified by another node.
+        """
+        self._topo_mgr.on_update_arg_by_node(dst_node, arg_idx, src_node, out_idx)
+
+    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 = f"\n[{title}]\n"
+        node_specs = [[
+            n.get_node_type(),
+            n.get_name(),
+            n.get_source_code(),
+            [[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'
+        return dump_str
+
+    def get_top_manager(self) -> 'NodeManager':
+        """
+        Get the top node_manager with type of no-method CallFunction or SymbolTree this
+        node_manager belongs to.
+        """
+        from .call_function import CallFunction
+        from ..symbol_tree import SymbolTree
+        if isinstance(self, SymbolTree):
+            return self
+        if isinstance(self, CallFunction) and not self.is_method():
+            return self
+        return self.get_node_manager().get_top_manager()
+
+    def set_manager_node_namer(self, node_namer):
+        """Set manager node namer"""
+        self._manager_node_namer = node_namer
+
+    def _add_node_to_nodes(self, node: Node):
+        """
+        Add `node` to `_nodes` dict.
+
+        Args:
+            node (Node): A Node to be added into `_nodes`.
+
+        Raises:
+            ValueError: 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,223 @@
+# 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
+from mindspore import log as logger
+from .node import Node
+from ..api.scoped_value import ScopedValue
+from ..common.observable import Observable
+from ..common.event import Event
+
+
+class TopoManager(Observable):
+    """SymbolTree topological-relationship manager."""
+
+    @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(index).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/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+# ============================================================================
+"""
+Parsers for resolve ast to SymbolTree
+"""
+from .parser import Parser
+from .parser_register import ParserRegister, ParserRegistry, reg_parser
+from .module_parser import ModuleParser
+from .arguments_parser import ArgumentsParser
+from .assign_parser import AssignParser
+from .for_parser import ForParser
+from .function_def_parser import FunctionDefParser
+from .if_parser import IfParser
+from .return_parser import ReturnParser
+from .class_def_parser import ClassDefParser
+from .while_parser import WhileParser
+from .expr_parser import ExprParser # Rely on AssignParser

mindspore/rewrite/parsers/arguments_parser.py

@@ -0,0 +1,63 @@
+# 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.
+# ============================================================================
+"""Parse ast.arguments to input-node of SymbolTree."""
+import ast
+
+from . import Parser, reg_parser
+from ..symbol_tree import SymbolTree
+from ..node import NodeManager
+from ..common.error_log import error_str
+
+
+class ArgumentsParser(Parser):
+    """Parse ast.arguments to input-node of SymbolTree."""
+
+    def target(self):
+        """Parse target type"""
+        return 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.
+            ast_node (ast.arguments): ast argument node in construct.
+            node_manager (NodeManager): NodeManager those asts belong to.
+
+        Raises:
+            TypeError: Types of ast_node.args elements are not ast.arg.
+        """
+        if hasattr(ast_node, "posonlyargs"):
+            stree.try_append_python_node(ast_node, ast_node.posonlyargs, node_manager)
+
+        for arg in ast_node.args:
+            if not isinstance(arg, ast.arg):
+                raise TypeError(error_str(f"only support ast.arg in arguments arg, but got '{type(arg).__name__}'",
+                                          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"):
+            pass
+        if hasattr(ast_node, "kwarg"):
+            stree.try_append_python_node(ast_node, ast_node.kwarg, node_manager)
+        if hasattr(ast_node, "defaults"):
+            pass
+
+
+g_arguments_parser = reg_parser(ArgumentsParser())