mindspore 2.4.0__cp310-cp310-macosx_11_0_arm64.whl

mindspore/rewrite/api/pattern_engine.py

@@ -0,0 +1,490 @@
+# 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.
+# ============================================================================
+"""PatternEngine for modifying SymbolTree by pattern."""
+from collections import OrderedDict
+from typing import Tuple, Union, List, Type
+import abc
+
+from mindspore.nn import Cell
+from mindspore.ops.primitive import Primitive
+from mindspore import log as logger
+from mindspore import _checkparam as Validator
+from .node_type import NodeType
+from .node import Node
+from .symbol_tree import SymbolTree
+
+
+class PatternNode:
+    """
+    `PatternNode` is defined as a node while defining pattern.
+
+    .. warning::
+        This is a set of experimental APIs that is subject to change or deletion.
+
+    Args:
+        pattern_node_name (str): Name of current node.
+        match_type (Type): A type represents what type would be matched of current node. Default: ``Type[None]`` .
+        inputs (list[PatternNode]): Input nodes of current node. Default: ``None`` .
+    """
+
+    def __init__(self, pattern_node_name: str, match_type: Type = Type[None], inputs: ['PatternNode'] = None):
+        Validator.check_value_type("pattern_node_name", pattern_node_name, [str], "PatternNode")
+        if inputs is not None:
+            Validator.check_element_type_of_iterable("inputs", inputs, [PatternNode], "PatternNode")
+        self._name = pattern_node_name
+        self._type = match_type
+        if inputs is None:
+            self._inputs = []
+        else:
+            self._inputs = inputs
+
+    @staticmethod
+    def from_node(node: Node) -> 'PatternNode':
+        """
+        Create a `PatternNode` from `node`.
+
+        Args:
+            node (Node): Input rewrite node.
+
+        Returns:
+            A `PatternNode` created from `node`.
+
+        Raises:
+            TypeError: If `node` is not a `Node` instance.
+        """
+        Validator.check_value_type("node", node, [Node], "PatternNode")
+        pattern_node: PatternNode = PatternNode(node.get_targets()[0])
+        if node.get_node_type() is NodeType.CallCell:
+            pattern_node._type = node.get_instance_type()
+        return pattern_node
+
+    @staticmethod
+    def create_pattern_from_node(node: Node) -> 'PatternNode':
+        """
+        Create a Pattern from `node` with its inputs.
+
+        Args:
+            node (Node): Input rewrite node.
+
+        Returns:
+            A `PatternNode` as root of pattern created from rewrite node.
+
+        Raises:
+            TypeError: If `node` is not a `Node` instance.
+        """
+        Validator.check_value_type("node", node, [Node], "PatternNode")
+        pattern_node: PatternNode = PatternNode.from_node(node)
+        inputs = []
+        for node_input in node.get_inputs():
+            inputs.append(PatternNode.create_pattern_from_node(node_input))
+        pattern_node._inputs = inputs
+        return pattern_node
+
+    @staticmethod
+    def create_pattern_from_list(type_list: []) -> 'PatternNode':
+        """
+        Create a Pattern from a cell type list.
+
+        Args:
+            type_list (list[type]): Input cell type list.
+
+        Returns:
+            A `PatternNode` as root of pattern created from cell type list.
+
+        Raises:
+            TypeError: If `type_list` is not a `list`.
+        """
+        Validator.check_value_type("type_list", type_list, [list], "PatternNode")
+        last_node = None
+        for i, cell_type in enumerate(type_list):
+            cur_node: PatternNode = PatternNode(str(i) + "-" + str(cell_type), cell_type, [])
+            if last_node is not None:
+                cur_node._inputs = [last_node]
+            else:
+                cur_node._inputs = []
+            last_node = cur_node
+        return last_node
+
+    def add_input(self, node):
+        """
+        Add an input for current `PatternNode`.
+
+        Args:
+            node (PatternNode): Cell type as an input.
+
+        Raises:
+            TypeError: If `node` is not a `PatternNode` instance.
+        """
+        Validator.check_value_type("node", node, [PatternNode], "PatternNode")
+        self._inputs.append(node)
+
+    def set_inputs(self, inputs):
+        """
+        Set inputs for current `PatternNode`.
+
+        Args:
+            inputs (list[PatternNode]) : Inputs to be set as inputs of current `PatternNode`.
+
+        Raises:
+            TypeError: If `inputs` is not a `list` or input in `inputs` is not `PatternNode` instance.
+        """
+        Validator.check_element_type_of_iterable("inputs", inputs, [PatternNode], "PatternNode")
+        self._inputs = inputs
+
+    def match(self, node: Node) -> bool:
+        """
+        Check if current `PatternNode` can match with `node`.
+
+        Args:
+            node (Node) : A rewrite node to be match.
+
+        Raises:
+            TypeError: If `node` is not a `Node` instance.
+        """
+        Validator.check_value_type("node", node, [Node], "PatternNode")
+        return self._type == node.get_instance_type()
+
+    def get_inputs(self):
+        """
+        Getter of inputs.
+
+        Returns:
+            A list of `PatternNode`, the inputs of current node.
+        """
+
+        return self._inputs
+
+    def name(self) -> str:
+        """
+        Getter of PatternNode name.
+        """
+        return self._name
+
+    def type(self):
+        """
+        Getter of PatternNode type.
+        """
+        return self._type
+
+
+class VarNode(PatternNode):
+    """
+    VarNode is a subclass of `PatternNode` whose `match` method is always return True.
+
+    .. warning::
+        This is a set of experimental APIs that is subject to change or deletion.
+    """
+
+    def __init__(self):
+        super(VarNode, self).__init__("placeholder", Cell, [])
+
+    def match(self, node: Node) -> bool:
+        return node is not None and node.get_handler() is not None
+
+
+class Replacement(abc.ABC):
+    """
+    Interface of replacement function.
+
+    .. warning::
+        This is a set of experimental APIs that is subject to change or deletion.
+
+    Examples:
+        >>> from mindspore.rewrite import Replacement, Node
+        >>> import mindspore.nn as nn
+        >>> class BnReplacement(Replacement):
+        ...     def build(self, pattern, is_chain_pattern: bool, matched):
+        ...         bn_node: Node = matched.get(pattern.name())
+        ...         conv = nn.Conv2d(16, 16, 3)
+        ...         conv_node = Node.create_call_cell(conv, ['x1'], bn_node.get_args(), bn_node.get_kwargs())
+        ...         return [conv_node]
+    """
+    @abc.abstractmethod
+    def build(self, pattern: PatternNode, is_chain_pattern: bool, matched: OrderedDict) -> [Node]:
+        """
+        Interface define for creating replacement nodes from matched result.
+
+        Note:
+            Return value will be delivered into replace api of `SymbolTree` as argument, return value should follow
+            restraint of parameter `new_nodes` of `replace` api if `SymbolTree`. See detail in docstring of `replace`
+            api of `SymbolTree`.
+
+        Args:
+            pattern (PatternNode): A `PatternNode` represents root node of current pattern.
+            is_chain_pattern (bool): A bool indicated if pattern is a chain pattern or a tree pattern.
+            matched (OrderedDict): An OrderedDict map from pattern_node name to node represents matched result.
+
+        Returns:
+            A list of instance of `Node` as replacement nodes.
+        """
+
+        raise NotImplementedError
+
+    def __call__(self, pattern: PatternNode, is_chain_pattern: bool, matched: OrderedDict) -> [Node]:
+        return self.build(pattern, is_chain_pattern, matched)
+
+
+class PatternEngine:
+    """
+    `PatternEngine` is defined how to transform a `SymbolTree` by `PattenNode`.
+
+    .. warning::
+        This is a set of experimental APIs that is subject to change or deletion.
+
+    Args:
+        pattern (Union[PatternNode, List]): An instance of `PatternNode` or a cell-type-list to construct `PatternNode`
+            as root of a pattern.
+        replacement (callable): A callable define how to generate new_node. Default: ``None`` .
+    """
+
+    def __init__(self, pattern: Union[PatternNode, List], replacement: Replacement = None):
+        Validator.check_value_type("replacement", replacement, [Replacement], "PatternEngine")
+        if isinstance(pattern, List):
+            for p in pattern:
+                if not issubclass(p, (Cell, Primitive)):
+                    raise TypeError(f"The type of 'pattern' should be one of [Cell, Primitive]"
+                                    f"but got '{p}' with type '{type(p).__name__}'.")
+        else:
+            Validator.check_value_type("pattern", pattern, [PatternNode], "PatternEngine")
+        if isinstance(pattern, PatternNode):
+            self._is_chain = False
+            self._replacement: Replacement = replacement
+            self._pattern: PatternNode = pattern
+        elif isinstance(pattern, list):
+            self._is_chain = True
+            self._replacement: Replacement = replacement
+            self._pattern: PatternNode = PatternNode.create_pattern_from_list(pattern)
+        else:
+            raise RuntimeError("Unsupported pattern define")
+
+    def pattern(self) -> PatternNode:
+        """
+        Getter of pattern.
+
+        Returns:
+            A instance of `PatternNode`, used to indicate the type that the current pattern needs to match.
+        """
+
+        return self._pattern
+
+    @staticmethod
+    def _multi_to_multi_replace(stree: SymbolTree, old_root: Node, matched_dict: OrderedDict,
+                                new_nodes: [Node]) -> Node:
+        """
+        Replace multi-nodes in `stree` by another list of nodes.
+
+        Note:
+            Call replace api of `SymbolTree`, so parameter `new_nodes` has same restraint with parameter `new_nodes` of
+            `replace` api if `SymbolTree`. See detail in docstring of `replace` api of `SymbolTree`.
+
+        Args:
+            stree (SymbolTree): A `SymbolTree` which replacement will apply on.
+            old_root (Node): A `Node` represents root of original nodes.
+            matched_dict (OrderedDict): An instance of OrderedDict as match result, where key is the pattern name, value
+                is the matched node.
+            new_nodes (list[Node]): A list of instance of Node as replacement.
+        """
+        # call replace of SymbolTree
+        new_root = stree.replace(old_root, new_nodes)
+        # replace only support one-to-one replace or one-to-multi replace, we need to erase nodes except
+        # cur_node manually
+        to_erase_list = matched_dict.values()
+        for node in reversed(to_erase_list):
+            if node != old_root:
+                stree.erase(node)
+        return new_root
+
+    @staticmethod
+    def _multi_replace_cellcontainer(stree, cellcontainer, node, matched_dict, new_nodes):
+        """Replace node in CellContainer."""
+        to_erase_list = list(matched_dict.values())
+        stree.replace(Node(node), new_nodes)
+        for n in reversed(to_erase_list):
+            if n.get_handler() is node:
+                continue
+            stree.erase(n)
+
+    def apply(self, stree: SymbolTree) -> bool:
+        """
+        Apply current pattern to a `SymbolTree`.
+
+        Note:
+            Sub-tree node will be supported in the near feature.
+
+        Args:
+            stree (SymbolTree): A `SymbolTree` to be transformed.
+
+        Returns:
+            A bool represents if `stree` been changed.
+
+        Raises:
+            TypeError: If `stree` is not a `SymbolTree` instance.
+        """
+        Validator.check_value_type("stree", stree, [SymbolTree], "PatternEngine")
+        changed = False
+        # IR match
+        queue: [Node] = stree.get_inputs()
+        # Why need visited: we don't need or should not to visit same node multi-times because pre-visited node may
+        # already been erased from SymbolTree.
+        # When will we visit same node multi-times:
+        #      a
+        #     / \
+        #    /   \
+        #   b     c
+        #   |     |
+        #   |     d
+        #   \    /
+        #    \  /
+        #     e
+        #  1. Visit e, e does not match pattern, add b, d to queue.
+        #  2. Visit b, b does not match pattern, add a to queue.
+        #  3. Visit d, d does not match pattern, add c to queue.
+        #  4. Visit a, a matches pattern and erased from SymbolTree, add xx to queue.
+        #  5. Visit c, c does not match pattern, add a to queue.
+        #  At step 5, a is visited at second time but a is erased from SymbolTree at step 4.
+        visited: [Node] = []
+        while queue:
+            cur_node: Node = queue.pop(0)
+            if cur_node is None:  # Because inputs of node is allowed to be None in replacement.
+                continue
+            if cur_node in visited:
+                continue
+            if cur_node.get_node_type() == NodeType.Tree:
+                subtree = cur_node.get_sub_tree()
+                self.apply(subtree)
+                visited.append(cur_node)
+                queue.extend(cur_node.get_users())
+                continue
+            if cur_node.get_node_type() == NodeType.CellContainer:
+                self._process_cellcontainer(stree, cur_node.get_handler())
+                continue
+            visited.append(cur_node)
+            matched, matched_dict = self._match(self._pattern, cur_node)
+            # not matched
+            if not matched or not PatternEngine._check_match(self._pattern, matched_dict):
+                queue.extend(cur_node.get_users())
+                continue
+            # matched
+            new_nodes: [Node] = []
+            if self._replacement is not None:
+                new_nodes: [Node] = self._replacement(self._pattern, self._is_chain, matched_dict)
+            if not new_nodes:  # if replacement is empty, do nothing
+                queue.extend(cur_node.get_users())
+            else:  # replace cur_node with new_nodes
+                changed = True
+                root = PatternEngine._multi_to_multi_replace(stree, cur_node, matched_dict, new_nodes)
+                queue.append(root)
+        return changed
+
+    @staticmethod
+    def _merge_ordered_dict(dict1: OrderedDict, dict2: OrderedDict) -> OrderedDict:
+        """
+        A static util method to merge two OrderedDict.
+
+        Args:
+            dict1 (OrderedDict): First dict to be merged.
+            dict2 (OrderedDict): Second dict to be merged.
+
+        Returns:
+            Merged OrderedDict.
+        """
+
+        merged = dict1.copy()
+        merged.update(dict2)
+        return merged
+
+    def _match(self, pattern: PatternNode, node: Node) -> Tuple[bool, OrderedDict]:
+        """
+        Match `pattern` with a `node` with all inputs of the `pattern`.
+
+        Args:
+            pattern (PatternNode): Pattern to be match.
+            node (Node): Node to be match.
+
+        Returns:
+            A bool value to indicate if matched.
+            An instance of OrderedDict as match result, where key is the pattern name, value is the matched node.
+        """
+
+        # Don't iterate into subgraph node, pattern should not be matched across sub-tree
+        if node.get_node_type() not in (NodeType.CallCell, NodeType.CallPrimitive, NodeType.Input):
+            logger.debug("Pattern match failed: node(%s) is not a CallCell, CallPrimitive or Input", str(node))
+            return False, OrderedDict()
+        if not pattern.match(node):
+            logger.debug("Pattern match failed: node(%s)'s type is %s while pattern type is %s", str(node),
+                         node.get_instance_type(), pattern.type())
+            return False, OrderedDict()
+        if isinstance(pattern, VarNode):
+            return True, OrderedDict()
+        pattern_inputs = pattern.get_inputs()
+        cur_inputs = node.get_inputs()
+        input_num = len(pattern_inputs)
+        if input_num == 0:
+            return True, OrderedDict({pattern.name(): node})
+        if input_num != len(cur_inputs):
+            logger.debug("Pattern match failed: node(%s)'s has %d inputs while pattern has %d inputs", str(node),
+                         len(node.get_inputs()), input_num)
+            return False, OrderedDict()
+        result = OrderedDict()
+        for i in range(0, input_num):
+            is_matched, tmp_result = self._match(pattern_inputs[i], cur_inputs[i])
+            if not is_matched:
+                return False, OrderedDict()
+            result = PatternEngine._merge_ordered_dict(result, tmp_result)
+        result[pattern.name()] = node
+        return True, result
+
+    @staticmethod
+    def _check_match(pattern: PatternNode, match_dict: OrderedDict) -> bool:
+        """
+        Check if matched result is a leak result.
+
+        A leak result means that the result is matched the `pattern`, but some nodes in result which is not
+        corresponding to root of pattern have outputs used by nodes outside of result.
+
+        Args:
+            pattern (PatternNode): A `PatternNode` represents pattern to be match.
+            match_dict (OrderedDict): A OrderedDict represents matched result.
+
+        Returns:
+            A bool value to indicate if matched result leaked.
+        """
+        matched_nodes = match_dict.values()
+        for key in match_dict:
+            if key == pattern.name():
+                continue
+            node: Node = match_dict[key]
+            for output in node.get_users():
+                if output not in matched_nodes:
+                    logger.debug("Check match failed, pattern leaked")
+                    return False
+        return True
+
+    def _process_cellcontainer(self, stree, cellcontainer):
+        """Process CellContainer node."""
+        for node in cellcontainer.nodes():
+            if node.get_node_type() == NodeType.Tree:
+                subtree = node.symbol_tree
+                self.apply(SymbolTree(subtree))
+                continue
+            matched, matched_dict = self._match(self._pattern, Node(node))
+            if not matched:
+                continue
+            new_nodes = []
+            if self._replacement is not None:
+                new_nodes = self._replacement(self._pattern, self._is_chain, matched_dict)
+            if not new_nodes:  # if replacement is empty, do nothing
+                continue
+            PatternEngine._multi_replace_cellcontainer(stree, cellcontainer, node, matched_dict, new_nodes)

mindspore/rewrite/api/scoped_value.py

@@ -0,0 +1,181 @@
+# 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.
+# ============================================================================
+"""Rewrite module api: ValueType and ScopedValue."""
+from enum import Enum
+from typing import Optional, Union, List, Tuple
+from mindspore import _checkparam as Validator
+
+
+class ValueType(Enum):
+    """
+    ValueType represents type of `ScopedValue`.
+
+    - A `NamingValue` represents a reference to another variable.
+    - A `CustomObjValue` represents an instance of custom class or an object whose type is out of range of base-type
+      and container-type of ValueType.
+    """
+
+    # constant type
+    ConstantValue = 0
+    # container type
+    TupleValue = 20
+    ListValue = 21
+    DictValue = 22
+    # variable type
+    NamingValue = 40
+    CustomObjValue = 41
+    # unsupported type
+    UnsupportedValue = 50
+
+
+class ScopedValue:
+    """
+    `ScopedValue` represents a value with its full-scope.
+
+    `ScopedValue` is used to express: a left-value such as target of an assign statement, or a callable object such as
+    func of a call statement, or a right-value such as args and kwargs of an assign statement.
+
+    Args:
+        arg_type (ValueType): A `ValueType` represents type of current value.
+        scope (str, optional): A string represents scope of current value. Take "self.var1" as an example,
+            `scope` of this var1 is "self". Default: ``""`` .
+        value: A handler represents value of current value. The type of value is corresponding to `arg_type`.
+            Default: ``None`` .
+    """
+
+    def __init__(self, arg_type: ValueType, scope: str = "", value=None):
+        Validator.check_value_type("arg_type", arg_type, [ValueType], "ScopedValue")
+        Validator.check_value_type("scope", scope, [str], "ScopedValue")
+        self.type = arg_type
+        self.scope = scope
+        self.value = value
+
+    @classmethod
+    def create_variable_value(cls, value) -> Optional['ScopedValue']:
+        """
+        Create `ScopedValue` from a variable.
+
+        `ScopedValue`'s type is determined by type of value. `ScopedValue`'s scope is empty.
+
+        Args:
+            value: The value to be converted to `ScopedValue`.
+
+        Returns:
+            An instance of `ScopedValue`.
+
+        Examples:
+            >>> from mindspore.rewrite import ScopedValue
+            >>> variable = ScopedValue.create_variable_value(2)
+            >>> print(variable)
+            2
+        """
+        if isinstance(value, (type(None), int, float, str, bool)):
+            return cls(ValueType.ConstantValue, "", value)
+        if isinstance(value, tuple):
+            return cls(ValueType.TupleValue, "",
+                       tuple(cls.create_variable_value(single_value) for single_value in value))
+        if isinstance(value, list):
+            return cls(ValueType.ListValue, "",
+                       list(cls.create_variable_value(single_value) for single_value in value))
+        if isinstance(value, dict):
+            return cls(ValueType.DictValue, "",
+                       dict((cls.create_variable_value(key),
+                             cls.create_variable_value(single_value)) for key, single_value in value.items()))
+        return cls(ValueType.CustomObjValue, "", value)
+
+    @classmethod
+    def create_naming_value(cls, name: str, scope: str = "") -> 'ScopedValue':
+        """
+        Create a naming `ScopedValue`. A `NamingValue` represents a reference to another variable.
+
+        Args:
+            name (str): A string represents the identifier of another variable.
+            scope (str, optional): A string represents the scope of another variable. Default: ``""`` .
+
+        Returns:
+            An instance of `ScopedValue`.
+
+        Raises:
+            TypeError: If `name` is not `str`.
+            TypeError: If `scope` is not `str`.
+
+        Examples:
+            >>> from mindspore.rewrite import ScopedValue
+            >>> variable = ScopedValue.create_naming_value("conv", "self")
+            >>> print(variable)
+            self.conv
+        """
+        Validator.check_value_type("name", name, [str], "ScopedValue")
+        Validator.check_value_type("scope", scope, [str], "ScopedValue")
+        return cls(ValueType.NamingValue, scope, name)
+
+    @staticmethod
+    def create_name_values(names: Union[List[str], Tuple[str]],
+                           scopes: Union[List[str], Tuple[str]] = None) -> List['ScopedValue']:
+        """
+        Create a list of naming `ScopedValue`.
+
+        Args:
+            names (List[str] or Tuple[str]): List or tuple of `str` represents names of referenced variables.
+            scopes (List[str] or Tuple[str], optional): List or tuple of `str` represents scopes of
+                referenced variables. Default: ``None`` .
+
+        Returns:
+            An list of instance of `ScopedValue`.
+
+        Raises:
+            TypeError: If `names` is not `list` or `tuple` and name in `names` is not `str`.
+            TypeError: If `scopes` is not `list` or `tuple` and scope in `scopes` is not `str`.
+            ValueError: If the length of names is not equal to the length of scopes when scopes are not None.
+
+        Examples:
+            >>> from mindspore.rewrite import ScopedValue
+            >>> variables = ScopedValue.create_name_values(names=["z", "z_1"], scopes=["self", "self"])
+            >>> print(variables)
+            [self.z, self.z_1]
+        """
+        Validator.check_element_type_of_iterable("names", names, [str], "ScopedValue")
+        if scopes is not None:
+            Validator.check_element_type_of_iterable("scopes", scopes, [str], "ScopedValue")
+            if len(names) != len(scopes):
+                raise ValueError("Length of names should be equal to length of scopes")
+        result = []
+        for index, name in enumerate(names):
+            if scopes is not None:
+                scope = scopes[index]
+            else:
+                scope = ""
+            result.append(ScopedValue.create_naming_value(name, scope))
+        return result
+
+    def __str__(self):
+        if self.type == ValueType.ConstantValue:
+            return str(self.value)
+        if self.type == ValueType.NamingValue:
+            return f"{self.scope}.{self.value}" if self.scope else str(self.value)
+        if self.type == ValueType.CustomObjValue:
+            return f"CustomObj: {str(self.value)}"
+        return f"Illegal ValueType: {str(self.type)}"
+
+    def __eq__(self, other):
+        if id(self) == id(other):
+            return True
+        return self.type == other.type and self.scope == other.scope and self.value == other.value
+
+    def __repr__(self):
+        return str(self)
+
+    def __hash__(self):
+        return hash((self.type, self.scope, self.value))