mindspore 2.0.0a0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/rewrite/sparsify/sparse_transformer.py

@@ -0,0 +1,448 @@
+# 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.
+# ============================================================================
+"""Sparsify transformer"""
+import ast
+import inspect
+import textwrap
+from collections import deque
+import astunparse
+
+from mindspore import ops, nn
+from mindspore import log as logger
+from mindspore.rewrite.parsers.assign_parser import AssignParser
+from mindspore.rewrite.sparsify.utils import ArgType, SparseFunc, sparse_rules, get_sparse_func, builtin_ops, \
+    get_binop_name, get_sparse_method_outputs, arg_type_to_prefix_map, get_inputs_outputs
+
+
+OPS_MODULE = "mindspore.ops."
+MAX_RECURSION_DEPTH = 10
+
+
+def sparsify_helper(f, arg_types, user_defined_rules=None, sparse_name="", full_sparse_rules=None, depth=0):
+    """Calls sparse_transformer from raw function."""
+    if isinstance(f, nn.Cell):
+        tree = ast.parse(textwrap.dedent(inspect.getsource(f.construct)))
+        # remove self
+        tree.body[0].args.args.pop(0)
+        global_vars = f.construct.__globals__
+        # pylint: disable=protected-access
+        init_vars = f._cells
+    else:
+        tree = ast.parse(textwrap.dedent(inspect.getsource(f)))
+        global_vars = f.__globals__
+        init_vars = {}
+    functiondef = tree.body[0]
+    args = [arg.arg for arg in functiondef.args.args]
+    type_map = dict(zip(args, arg_types))
+
+    sparse_transformer = SparseTransformer(
+        type_map, global_vars, init_vars, user_defined_rules, full_sparse_rules, depth)
+    sparse_tree = []
+    if not sparse_name:
+        sparse_name = functiondef.name
+    changed = False
+    for body in functiondef.body:
+        sparse_body = sparse_transformer.transform(body)
+        changed |= sparse_transformer.has_changed()
+        sparse_tree.append(sparse_body)
+    return_types = sparse_transformer.return_types
+
+    if changed:
+        sparse_tree = list(x[0] for x in sparse_transformer.sparse_functiondef.values()) + sparse_tree
+        ast_module = ast.Module([ast.FunctionDef(
+            sparse_name, functiondef.args, sparse_tree, functiondef.decorator_list, functiondef.returns)])
+        return ast_module, True, return_types
+    return tree, False, return_types
+
+
+class SparseTransformer(ast.NodeTransformer):
+    """Transformer class for sparsify."""
+    def __init__(self, type_map, global_vars, init_vars, user_defined_rules=None, full_sparse_rules=None, depth=0):
+        """Init method."""
+        super().__init__()
+        self.type_map = type_map
+        self.global_vars = global_vars
+        self.init_vars = init_vars
+        self.depth = depth
+        self.return_types = (ArgType.NONSPARSE,)
+        # maps function name and arg types to sparsified ast and return types, which are then inserted into module
+        self.sparse_functiondef = {}
+        # maps function name and arg types to return types for ast that do not change after sparsify
+        self.origin_functiondef = {}
+
+        # keeps track of arg_type for each operand on the call stack recursively
+        self._frames = deque()
+        self._changed = False
+        # variables for which arg_types diverge with control flow are not supported, and are considered dead
+        # after exiting the block
+        self._dead_vars = {}
+        # full_sparse_rules are inherited from caller cell and takes precedence over generic rules
+        if full_sparse_rules:
+            self.full_sparse_rules = full_sparse_rules
+        else:
+            self.full_sparse_rules = {}
+            user_defined_rules = user_defined_rules or {}
+            self.get_sparse_rules(user_defined_rules)
+
+    @staticmethod
+    def make_call(node, name="", args=None):
+        """Returns a call node with given name and args, if provided."""
+        if name:
+            func = ast.Name(name, ast.Load())
+        else:
+            func = node.func
+        if args is None:
+            args = node.args
+        return ast.Call(func, args, node.keywords)
+
+    def get_sparse_rules(self, user_defined_rules):
+        """Generates sparse rules for the transformer from generic sparse rules and user-defined sparse rules."""
+        for func, rules in {**sparse_rules, **user_defined_rules}.items():
+            for r in rules:
+                sparse_func = get_sparse_func(r)
+                # sparse rules are accessed by the function object and input arg_types pair
+                sparse_func_map = self.full_sparse_rules.get(func, {})
+                sparse_func_map[tuple(sparse_func.inputs)] = sparse_func
+                self.full_sparse_rules[func] = sparse_func_map
+
+    def transform(self, node):
+        """Transforms a single node which represents a stmt in the ast."""
+        self.clear_stack()
+        self._changed = False
+        stmt = self.visit(node)
+        return stmt
+
+    def has_changed(self):
+        """Whether the SparseTransformer has changed"""
+        return self._changed
+
+    def add_frame(self):
+        """Add a frame into deque."""
+        self._frames.append([])
+
+    def pop_frame(self):
+        """Pop a frame in deque."""
+        return tuple(self._frames.pop())
+
+    def push_onto_frame(self, t):
+        """Push an arg_type into frame deque."""
+        if not self._frames:
+            raise ValueError("Current frame not initialized!")
+        self._frames[-1].append(t)
+
+    def push_all_onto_frame(self, t):
+        """Push all arg_types into frame deque."""
+        if not self._frames:
+            raise ValueError("Current frame not initialized!")
+        for i in t:
+            self._frames[-1].append(i)
+
+    def clear_stack(self):
+        """Clear frame deque"""
+        self._frames.clear()
+
+    def make_sparse_func(self, func, node_type, inputs):
+        """Returns SparseFunc by looking up sparse_rules."""
+        rules = {}
+        if node_type == ast.Call:
+            if isinstance(func, nn.Cell):
+                func_name = func.__class__.__name__.lower()
+            else:
+                func_name = getattr(func, "__name__", func)
+        elif node_type == ast.BinOp:
+            func_name = func
+        rules = self.full_sparse_rules.get(func, {})
+
+        if ArgType.ANY in rules:
+            sparse_func = rules[ArgType.ANY]
+        elif inputs in rules:
+            sparse_func = rules[inputs]
+        else:
+            # attempts to find sparse op based on sparse prefix if sparse rules not found
+            sparse_func_name = arg_type_to_prefix_map.get(inputs[0], "$") + "_" + func_name
+            sparse_op = getattr(ops, sparse_func_name, None)
+            if sparse_op is None:
+                if any(input_type != ArgType.NONSPARSE for input_type in inputs):
+                    return None
+                outputs = (ArgType.NONSPARSE,)
+            else:
+                func_name = sparse_func_name
+                _, outputs = get_inputs_outputs(sparse_op)
+            sparse_func = SparseFunc(func_name, inputs, outputs)
+
+        if sparse_func.fn != func:
+            self._changed = True
+        return sparse_func
+
+    def get_sparse_node(self, node, args, func, arg_types):
+        """
+        Retrieves target from sparse rules if matches, otherwise sparsify the node by recursively expanding `func`
+        until maximum recursion depth is reached. Functions in mindspore.ops are not expanded.
+        If no matching sparse rule is found, an error is raised.
+        """
+        sparse_func = self.make_sparse_func(func, type(node), arg_types)
+        if sparse_func is not None:
+            if self._changed:
+                func_node = ast.Name(sparse_func.fn, ast.Load())
+                if sparse_func.fn in self.global_vars:
+                    func_node = ast.Name(sparse_func.fn, ast.Load())
+                else:
+                    func_node = ast.Name("ops", ast.Load())
+                    func_node = ast.Attribute(func_node, sparse_func.fn, ast.Load())
+                node = ast.Call(func_node, args, node.keywords)
+            self.push_all_onto_frame(sparse_func.outputs)
+            return node
+
+        if func.__module__[:len(OPS_MODULE)] == OPS_MODULE:
+            raise ValueError(f"Sparse rules not registered for {func}!")
+
+        if isinstance(func, nn.Cell):
+            class_name = func.__class__.__name__
+            func_name = class_name.lower()
+            init_args = inspect.getfullargspec(func).args
+            if len(init_args) != 1:
+                raise ValueError(f"Nested cell {class_name} with arguments for init supported!")
+        else:
+            func_name = func.__name__
+        sparse_func_name = f"sparse_{'_'.join(arg_type_to_prefix_map.get(t, 'default') for t in arg_types)}_{func_name}"
+        if (func_name, arg_types) in self.sparse_functiondef:
+            self._changed = True
+            # pylint: disable=get-dict-value-exception
+            self.push_all_onto_frame(self.sparse_functiondef[(func_name, arg_types)][1])
+            return SparseTransformer.make_call(node, sparse_func_name, args)
+        if (func_name, arg_types) in self.origin_functiondef:
+            # pylint: disable=get-dict-value-exception
+            self.push_all_onto_frame(self.origin_functiondef[(func_name, arg_types)])
+            return node
+        if self.depth == MAX_RECURSION_DEPTH:
+            raise RuntimeError(f"Maximum recursion depth {MAX_RECURSION_DEPTH} for sparsify reached at {func}!")
+        functiondef, changed, return_types = sparsify_helper(
+            func, arg_types, sparse_name=sparse_func_name, full_sparse_rules=self.full_sparse_rules,
+            depth=self.depth + 1)
+        self.push_all_onto_frame(return_types)
+        if changed:
+            self._changed = True
+            self.sparse_functiondef[(func_name, arg_types)] = (functiondef, return_types)
+            return SparseTransformer.make_call(node, sparse_func_name, args)
+        self.origin_functiondef[(func_name, arg_types)] = return_types
+        return SparseTransformer.make_call(node, args=args)
+
+    def map_type_to_target(self, node_target, value_types):
+        """Records arg_type for each target."""
+        if isinstance(node_target, (ast.Tuple, ast.List)):
+            targets = node_target.elts
+            if len(targets) != len(value_types):
+                raise ValueError(f"Target {astunparse.unparse(node_target)} size and value size not match for "
+                                 f"ast.Assign {len(targets)} != {len(value_types)}")
+            target_vars = []
+            for target in targets:
+                if not isinstance(target, ast.Name):
+                    raise ValueError(f"Each target {ast.dump(target)} for ast.Assign should be ast.Name!")
+                target_vars.append(target.id)
+            for var, t in zip(target_vars, value_types):
+                self.type_map[var] = t
+        elif isinstance(node_target, ast.Name):
+            var = node_target.id
+            if len(value_types) == 1:
+                self.type_map[var] = value_types[0]
+            else:
+                self.type_map[var] = value_types
+        else:
+            raise ValueError(f"Targets for ast.Assign not supported for {type(node_target)}!")
+
+    def visit_method(self, node):
+        """Visits each node based on node class."""
+        method = "visit_" + node.__class__.__name__
+        visitor = getattr(self, method, None)
+        if visitor is None:
+            raise ValueError(f"{type(node)} is not supported in SparseTransformer!")
+        return visitor(node)
+
+    def visit(self, node):
+        """Visitor interface for all nodes."""
+        if not node._fields:
+            return node
+        if isinstance(node, (ast.AugAssign, ast.Expr)):
+            return self.visit_generic_stmt(node)
+        if isinstance(node, (ast.BoolOp, ast.Compare, ast.Subscript)):
+            # node always evaluates to non-sparse values
+            return self.visit_generic_expr(node)
+        if isinstance(node, (ast.Tuple, ast.List, ast.UnaryOp)):
+            # node contains multiple expressions but is not composable
+            return self.visit_composite_generic_expr(node)
+        if isinstance(node, (ast.Attribute, ast.Num, ast.Str)):
+            return self.visit_scalar_expr(node)
+        if isinstance(node, (ast.Index, ast.Slice)):
+            # node forms only a part of an expression and does not exist as standalone expression
+            return self.visit_partial_expr(node)
+        return self.visit_method(node)
+
+    def visit_generic_stmt(self, node):
+        """Visitor for generic statement."""
+        self.add_frame()
+        node = self.generic_visit(node)
+        self.pop_frame()
+        return node
+
+    def visit_scalar_expr(self, node):
+        """Visitor for scalar expression."""
+        self.push_onto_frame(ArgType.NONSPARSE)
+        return node
+
+    def visit_generic_expr(self, node):
+        """Visitor for generic expression."""
+        self.add_frame()
+        node = self.generic_visit(node)
+        self.pop_frame()
+        self.push_onto_frame(ArgType.NONSPARSE)
+        return node
+
+    def visit_composite_generic_expr(self, node):
+        """Visitor for composite generic expression."""
+        return self.generic_visit(node)
+
+    def visit_partial_expr(self, node):
+        """Visitor for a part of an expression."""
+        return node
+
+    def visit_Assign(self, node):     # pylint: disable=invalid-name
+        """Visitor for ast.Assign."""
+        self.add_frame()
+        value = self.visit(node.value)
+        value_types = self.pop_frame()
+        for node_target in node.targets:
+            self.map_type_to_target(node_target, value_types)
+        return ast.Assign(node.targets, value)
+
+    def visit_BinOp(self, node):     # pylint: disable=invalid-name
+        """Visitor for ast.Binop."""
+        self.add_frame()
+        node = self.generic_visit(node)
+        arg_types = self.pop_frame()
+        if len(arg_types) != 2:
+            raise ValueError(f"Binary op {astunparse.unparse(node)} values for arg_type len({arg_types}) != 2")
+        func = get_binop_name(node.op)
+        if func:
+            sparse_func = self.make_sparse_func(func, type(node), arg_types)
+            if sparse_func is None:
+                raise ValueError(f"Sparse rules not defined for {arg_types[0]} {func} {arg_types[1]}!")
+            outputs = sparse_func.outputs
+        else:
+            outputs = (ArgType.NONSPARSE,)
+        self.push_all_onto_frame(outputs)
+        return node
+
+    def visit_Call(self, node):     # pylint: disable=invalid-name
+        """Visitor for ast.Call."""
+        self.add_frame()
+        args = []
+        for arg in node.args:
+            args.append(self.visit(arg))
+        arg_types = self.pop_frame()
+
+        if all(t == ArgType.NONSPARSE for t in arg_types):
+            # if none of the arguments is sparse, do nothing
+            self.push_onto_frame(ArgType.NONSPARSE)
+            return node
+
+        # pylint: disable=protected-access
+        func_name = AssignParser._get_func_name(node)
+        if func_name is None or func_name == "":
+            raise RuntimeError(f"Function not exist for {ast.dump(node)}!")
+        # pylint: disable=protected-access
+        func_scope = AssignParser._get_func_scope(node)
+
+        if not func_scope:
+            if func_name in builtin_ops:
+                self.push_onto_frame(ArgType.NONSPARSE)
+                return node
+            if func_name in self.global_vars:
+                # external function with sparse arguments are inlined and cached
+                func = self.global_vars[func_name]
+                return self.get_sparse_node(node, args, func, arg_types)
+            raise ValueError(f"Call to undefined {func_name}!")
+
+        if func_scope in self.global_vars:
+            namespace = self.global_vars[func_scope]
+            func = getattr(namespace, func_name, None)
+            if func is None:
+                raise ValueError(f"{func_name} not defined in {namespace}!")
+            return self.get_sparse_node(node, args, func, arg_types)
+
+        if func_scope == "self":
+            func = self.init_vars.get(func_name, None)
+            if func is None:
+                raise ValueError(f"{func_name} not defined in in Cell.__init__!")
+            return self.get_sparse_node(node, args, func, arg_types)
+
+        func_scope_type = self.type_map.get(func_scope, None)
+        if func_scope_type is not None:
+            # tensor methods
+            if func_scope_type == ArgType.NONSPARSE:
+                outputs = (ArgType.NONSPARSE,)
+            else:
+                outputs = get_sparse_method_outputs(func_name, func_scope_type)
+            self.push_all_onto_frame(outputs)
+            return node
+        raise ValueError(f"Undefined var {func_scope}!")
+
+    def visit_Name(self, node):     # pylint: disable=invalid-name
+        """Visitor for ast.Name."""
+        if node.id in self.type_map:
+            tensor_type = self.type_map[node.id]
+        elif node.id in self.global_vars:
+            logger.warning(f"Global variable {node.id} treaded as nonsparse value by default.")
+            tensor_type = ArgType.NONSPARSE
+        elif node.id in self._dead_vars:
+            raise ValueError(f"Divergent arg_types {self._dead_vars.get(node.id)} for {node.id} are currently not "
+                             f"supported in control flow and the variable is considered dead upon leaving "
+                             f"the block")
+        else:
+            raise ValueError(f"Undefined variable {node.id}!")
+
+        if isinstance(tensor_type, tuple):
+            self.push_all_onto_frame(tensor_type)
+        else:
+            self.push_onto_frame(tensor_type)
+        return node
+
+    def visit_Return(self, node):     # pylint: disable=invalid-name
+        """Visitor for ast.Return."""
+        self.add_frame()
+        node = self.generic_visit(node)
+        self.return_types = self.pop_frame()
+        return node
+
+    def visit_While(self, node):     # pylint: disable=invalid-name
+        """
+        Visitor for ast.While.
+        Variables for which arg_types diverge with control flow are not supported, and as a fallback routine,
+        unsupported variables are treated as out-of-scope after leaving the control flow body.
+        """
+        self.add_frame()
+        test = self.visit(node.test)
+        self.pop_frame()
+        orig_type_map = self.type_map.copy()
+        body = list(self.visit(expr) for expr in node.body)
+        for var, t in self.type_map.items():
+            if var not in orig_type_map:
+                # new variables in while body are considered active after the leaving the block
+                orig_type_map[var] = t
+            elif orig_type_map[var] != t:
+                # variables for which arg_types diverge are considered dead after leaving the block
+                self._dead_vars[var] = (t, orig_type_map.pop(var))
+        self.type_map = orig_type_map
+        orelse = list(self.visit(expr) for expr in node.orelse)
+        return ast.While(test, body, orelse)

mindspore/rewrite/sparsify/sparsify.py

@@ -0,0 +1,109 @@
+# 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.
+# ============================================================================
+"""sparsify implementation"""
+import os
+
+from mindspore import ops
+from mindspore.rewrite import SymbolTree, ScopedValue
+from mindspore.rewrite.ast_helpers import AstModifier
+from mindspore.rewrite.sparsify.sparse_transformer import SparseTransformer
+from mindspore.rewrite.sparsify.utils import SparseFunc, ArgType
+
+
+op_vars = vars(ops)
+
+
+def get_user_defined_rules(sparse_rules, global_vars, tree):
+    """Register user-defined sparse rules."""
+    user_defined_rules = {}
+
+    def register_callable(fn):
+        func_name = fn.__name__
+        if global_vars.get(func_name, None) is fn:
+            init_targets = [ScopedValue.create_naming_value(func_name, "self")]
+            AstModifier.append_global_vars_expr_to_init(tree.get_init_func_ast(), init_targets, func_name)
+        elif not op_vars.get(func_name, None) is fn:
+            raise ValueError(f"{fn} not found in globals or mindspore.ops!")
+
+    for source, targets in sparse_rules.items():
+        if not isinstance(targets, (tuple, list)) or isinstance(targets, SparseFunc):
+            targets = [targets]
+        else:
+            targets = list(targets)
+        for sparse_func in targets:
+            if isinstance(sparse_func, SparseFunc) and callable(sparse_func.fn):
+                register_callable(sparse_func.fn)
+            elif callable(sparse_func):
+                register_callable(sparse_func)
+            rule = user_defined_rules.get(source, [])
+            rule.append(sparse_func)
+            user_defined_rules[source] = rule
+
+    return user_defined_rules
+
+
+def sparsify_tree(tree, arg_types, sparse_rules, f):
+    """Sparsify SymbolTree object."""
+    global_vars = f.construct.__globals__
+    user_defined_rules = get_user_defined_rules(sparse_rules, global_vars, tree)
+
+    # skip self
+    args = [arg.arg for arg in tree.get_ast_root().args.args[1:]]
+    if isinstance(arg_types, tuple):
+        if len(args) != len(arg_types):
+            raise ValueError(f"arg_types should have the same length as function parameters, but "
+                             f"{len(arg_types)} != {len(args)}!")
+        type_map = dict(zip(args, arg_types))
+    elif isinstance(arg_types, dict):
+        if all(isinstance(i, int) for i in arg_types.keys()):
+            type_map = {args[i]: arg_types[i] if i in arg_types else ArgType.NONSPARSE for i in range(len(args))}
+        elif all(isinstance(i, str) for i in arg_types.keys()):
+            type_map = {arg: arg_types[arg] if arg in arg_types else ArgType.NONSPARSE for arg in args}
+        else:
+            raise ValueError(f"Keys for arg_types {list(arg_types.keys())} should be all ints or all strings!")
+    else:
+        raise ValueError(f"Unsupported type for arg_types {type(arg_types)}!")
+
+    # pylint: disable=protected-access
+    init_vars = f._cells
+    sparse_transformer = SparseTransformer(type_map, global_vars, init_vars, user_defined_rules)
+    for i, node_ast in enumerate(tree.get_ast_root().body):
+        sp_ast = sparse_transformer.transform(node_ast)
+        if sparse_transformer.has_changed():
+            tree.get_ast_root().body[i] = sp_ast
+    for module, _ in sparse_transformer.sparse_functiondef.values():
+        tree.get_module_ast().body.append(module)
+
+
+def sparsify(f, arg_types, sparse_rules=None):
+    """
+    Sparsify a Cell object by inferring the appropriate sparse function calls to replace the original function calls by
+    propagating sparse properties provided in `arg_types`.
+
+    Args:
+        f (Cell): Cell object to be sparsified.
+        arg_types (Tuple[ArgType] | Dict[int, ArgType]): The type of argument (sparse csr, sparse coo,
+            non-sparse etc.) expected by `f`. If `arg_type` is a tuple, its length should be the same as the number of
+            arguments for `f`; if `arg_type` is a dictionary, each key represents an index into the arguments, and
+            arguments not referenced by the dictionary are considered to be non-sparse.
+        sparse_rules (Dict[str, SparseFunc], Optional): Additional sparse rules.
+    """
+    os.environ["STREE_PYTHON_FALLBACK"] = "1"
+    tree = SymbolTree.create(f)
+    handler = tree.get_handler()
+    sparse_rules = sparse_rules or {}
+    sparsify_tree(handler, arg_types, sparse_rules, f)
+    os.unsetenv("STREE_PYTHON_FALLBACK")
+    return tree.get_network()