mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/rewrite/parsers/class_def_parser.py

@@ -13,16 +13,17 @@
 # limitations under the License.
 # ============================================================================
 """Parse ast.ClassDef which is subclass of Cell to SymbolTree."""
+import sys
 import ast
-
+import inspect
 import astunparse
 from mindspore import log as logger
 from mindspore._extends.parse.namespace import CellNamespace
+from mindspore.rewrite.ast_creator_register import ast_creator_registry
 from ..symbol_tree import SymbolTree
 from ..parser import Parser
 from ..parser_register import ParserRegister, reg_parser
-from ..api.scoped_value import ScopedValue
-from ..ast_helpers import AstReplacer, AstModifier
+from ..ast_helpers import AstReplacer
 from ..common import error_str
 
 
@@ -115,50 +116,88 @@ class ClassDefParser(Parser):
         """Parse target type"""
         return ast.ClassDef
 
+    def handle_father_class(self, stree, node: ast.ClassDef):
+        """Handle father class."""
+        for base in node.bases:
+            parser: Parser = ParserRegister.instance().get_parser(type(base))
+            father_class = parser.process(stree, base)
+            if "Cell" not in father_class:
+                for k, m in sys.modules.items():
+                    if k in ("_ast", "ast"):
+                        continue
+                    if hasattr(m, father_class):
+                        cls = getattr(m, father_class)
+                        source_code = inspect.getsource(cls)
+                        ast_root: ast.Module = ast.parse(source_code)
+                        stree._father_class_ast.append(ast_root) # pylint: disable=protected-access
+                        break
+
+    def process(self, stree: SymbolTree, node: ast.ClassDef):
+        """
+        Parse init and construct in ast.ClassDef.
+
+        Args:
+            stree ([SymbolTree]): Symbol Tree under parsing.
+            node ([ast.ClassDef]): An ast.ClassDef node.
+        """
+        replacer = AstReplacer(node)
+        replacer.replace_all(stree.get_ori_cls_name(), stree.get_opt_cls_name())
+
+        stree.set_class_ast(node)
+        self.handle_father_class(stree, node)
+
+        if self._need_add_init_func(node):
+            self._add_init_func(node)
+
+        for body in node.body:
+            if isinstance(body, ast.FunctionDef):
+                if body.name == "__init__":
+                    self._process_init_func_ast(stree, node, body)
+                    stree.set_init_func_ast(body)
+                elif body.name == "construct":
+                    parser: Parser = ParserRegister.instance().get_parser(ast.FunctionDef)
+                    parser.process(stree, body)
+                else:
+                    logger.info(
+                        "Ignoring ast.FunctionDef in ast.ClassDef except __init__ and construct function: %s",
+                        body.name)
+            else:
+                logger.info("Ignoring unsupported node(%s) in ast.ClassDef.", type(body).__name__)
+
     def _is_subtree_field(self, ori_net, field) -> bool:
         op = getattr(ori_net, field)
         return not type(op).__name__ in self._cell_namespace
 
-    def _process_init_func_ast(self, stree: SymbolTree, init_ast: ast.FunctionDef):
+    def _process_init_func_ast(self, stree: SymbolTree, cls_ast: ast.ClassDef, init_ast: ast.FunctionDef):
         """Process init func"""
-        super_index = ClassDefParser._find_super_expr_of_init_func(init_ast)
         ClassDefParser._modify_arguments_of_init_func(init_ast)
-        self._replace_ori_field_of_init_func(stree, init_ast.body, super_index)
-        # re-find super_index for init_func changed in _replace_ori_field_of_init_func
-        super_index = ClassDefParser._find_super_expr_of_init_func(init_ast)
-        ClassDefParser._insert_handler_to_init_func(init_ast, super_index)
+        new_bodies = self._replace_ori_field_of_init_func(stree, cls_ast, init_ast.body)
+        init_ast.body = new_bodies
 
     @staticmethod
-    def _find_super_expr_of_init_func(ast_init_fn: ast.FunctionDef) -> int:
-        """Find index of super(XXnet).__init__() in body of init ast.FunctionDef"""
-        if not ast_init_fn.body:
-            return -1
-        super_index = -1
-        while True:
-            super_index += 1
-            expr = ast_init_fn.body[super_index]
-            if not isinstance(expr, ast.Expr):
-                continue
-            expr_value = expr.value
-            if not isinstance(expr_value, ast.Call):
-                continue
-            expr_value_func = expr_value.func
-            if not isinstance(expr_value_func, ast.Attribute):
-                continue
-            expr_value_func_value = expr_value_func.value
-            if expr_value_func.attr != "__init__" or not isinstance(expr_value_func_value, ast.Call):
-                continue
-            expr_value_func_value_func = expr_value_func_value.func
-            if not isinstance(expr_value_func_value_func, ast.Name) or expr_value_func_value_func.id != "super":
-                continue
-            break
-        return super_index
+    def _is_super_expr(expr: ast.AST) -> bool:
+        """Check whether ast node is super().__init__()"""
+        if not isinstance(expr, ast.Expr):
+            return False
+        expr_value = expr.value
+        if not isinstance(expr_value, ast.Call):
+            return False
+        expr_value_func = expr_value.func
+        if not isinstance(expr_value_func, ast.Attribute):
+            return False
+        expr_value_func_value = expr_value_func.value
+        if expr_value_func.attr != "__init__" or not isinstance(expr_value_func_value, ast.Call):
+            return False
+        expr_value_func_value_func = expr_value_func_value.func
+        if not isinstance(expr_value_func_value_func, ast.Name) or expr_value_func_value_func.id != "super":
+            return False
+        return True
 
     @staticmethod
     def _modify_arguments_of_init_func(ast_init_fn: ast.FunctionDef):
         """Replace init function input parameters to self and global_vars."""
         arg_self = ast.arg(arg="self", annotation="")
-        arg_global_vars = ast.arg(arg="global_vars", annotation="")
+        arg_global_vars = ast.arg(arg="obj", annotation="")
         ast_init_fn.args = ast.arguments(args=[arg_self, arg_global_vars], posonlyargs=[], kwonlyargs=[],
                                          kw_defaults=[], defaults=[], vararg=None, kwarg=None)
         ast.fix_missing_locations(ast_init_fn)
@@ -189,7 +228,116 @@ class ClassDefParser(Parser):
         for counter, index in enumerate(body_index_to_be_deleted):
             bodies.pop(index - counter)
 
-    def _replace_ori_field_of_init_func(self, stree: SymbolTree, bodies: [], super_index: int):
+    def _handle_tuple_for_replace_ori_field(self, ast_tuple: ast.Tuple, new_bodies):
+        """ Handle ast.Assign node with target of ast.Tuple in init func to new ast nodes. """
+        for e in ast_tuple.elts:
+            if isinstance(e, ast.Attribute):
+                field_name = e.attr
+                value = ast.Call(ast.Name('getattr', ast.Load()),
+                                 [ast.Name('obj', ast.Load()),
+                                  ast.Constant(value=field_name, kind=None)], [])
+                new_assign = ast_creator_registry.get("Assign")(targets=[e], value=value)
+                new_bodies.append(new_assign)
+
+    def _handle_express_for_replace_ori_field(self, cls_ast, ast_expr: ast.Expr, stree, new_bodies):
+        """ Handle ast.Expr node in init func to new ast nodes. """
+        ast_call = ast_expr.value
+        if not isinstance(ast_call.func, ast.Attribute) or not isinstance(ast_call.func.value, ast.Name)\
+            or ast_call.func.value.id != 'self':
+            return
+        for func_def in cls_ast.body:
+            if isinstance(func_def, ast.FunctionDef) and func_def.name == ast_call.func.attr:
+                for func_def_body in func_def.body:
+                    self._handle_bodies_for_replace_ori_field(cls_ast, func_def_body, stree, new_bodies)
+                return
+
+    def _handle_assign_for_replace_ori_field(self, ast_assign: ast.Assign, stree, new_bodies):
+        """ Handle ast.Assign node in init func to new ast nodes. """
+        if len(ast_assign.targets) != 1:
+            raise RuntimeError("not support multi-targets in assign now!", father_node=ast_assign)
+        target = ast_assign.targets[0]
+        if isinstance(target, ast.Tuple):
+            self._handle_tuple_for_replace_ori_field(target, new_bodies)
+            return
+        if not isinstance(target, ast.Attribute) or not isinstance(target.value, ast.Name)\
+            or target.value.id != 'self':
+            logger.info(f"Ignoring {astunparse.unparse(target)} in __init__ function.")
+            return
+        field_name = target.attr
+        # Ensure that the instance has corresponding attribute
+        if not hasattr(stree.get_origin_network(), field_name):
+            return
+        # Check to avoid repeat code
+        for new_ast in new_bodies:
+            if isinstance(new_ast, ast.Assign) and isinstance(new_ast.targets[0], ast.Attribute)\
+                and new_ast.targets[0].attr == field_name:
+                return
+        value = ast.Call(ast.Name('getattr', ast.Load()),
+                         [ast.Name('obj', ast.Load()),
+                          ast.Constant(value=field_name, kind=None)], [])
+        new_assign = ast_creator_registry.get("Assign")(targets=[target], value=value)
+        new_bodies.append(new_assign)
+
+    def _handle_bodies_for_replace_ori_field(self, cls_ast, body, stree, new_bodies):
+        """ handle_bodies_for_replace_ori_field. """
+        if self._is_super_expr(body):
+            new_bodies.append(body)
+            return
+        if isinstance(body, ast.If):
+            for if_body in body.body + body.orelse:
+                self._handle_bodies_for_replace_ori_field(cls_ast, if_body, stree, new_bodies)
+            return
+        if isinstance(body, ast.Expr) and isinstance(body.value, ast.Call):
+            self._handle_express_for_replace_ori_field(cls_ast, body, stree, new_bodies)
+            return
+        if isinstance(body, ast.Assign):  # if not assign node, delete
+            self._handle_assign_for_replace_ori_field(body, stree, new_bodies)
+            return
+
+    def _need_add_init_func(self, cls_ast: ast.ClassDef) -> bool:
+        """If class has base nn.Cell but not have init func, then we need to add init func"""
+        base_nn_cell = False
+        for base in cls_ast.bases:
+            if isinstance(base, ast.Name) and base.id == 'Cell'\
+                or isinstance(base, ast.Attribute) and isinstance(base.value, ast.Name)\
+                and base.value.id == "nn" and base.attr == 'Cell':
+                base_nn_cell = True
+                break
+        if not base_nn_cell:
+            return False
+        for body in cls_ast.body:
+            if isinstance(body, ast.FunctionDef) and body.name == '__init__':
+                return False
+        return True
+
+    def _add_init_func(self, cls_ast: ast.ClassDef):
+        """Add init func with super().__init__()"""
+        init_func_ast = ast.FunctionDef(
+            name='__init__',
+            args=ast.arguments(
+                posonlyargs=[],
+                args=[
+                    ast.arg(arg='self')],
+                kwonlyargs=[],
+                kw_defaults=[],
+                defaults=[]),
+            body=[
+                ast.Expr(
+                    value=ast.Call(
+                        func=ast.Attribute(
+                            value=ast.Call(
+                                func=ast.Name(id='super', ctx=ast.Load()),
+                                args=[],
+                                keywords=[]),
+                            attr='__init__',
+                            ctx=ast.Load()),
+                        args=[],
+                        keywords=[]))],
+            decorator_list=[])
+        cls_ast.body.insert(0, init_func_ast)
+        ast.fix_missing_locations(cls_ast)
+
+    def _replace_ori_field_of_init_func(self, stree: SymbolTree, cls_ast: ast.ClassDef, bodies: []):
         """
         Replace original field in init func to self.XX = getattr(self._handler, "XX").
         Only keep following two kinds of ast nodes in bodies right now:
@@ -198,86 +346,16 @@ class ClassDefParser(Parser):
 
         Args:
             bodies ([]): bodied of init ast.FunctionDef.
-            super_index (int): index of super().__init__() in bodies.
 
         Raises:
             RuntimeError: Not support multi-targets in assign.
             RuntimeError: Only support target.value in [ast.Name] in assign node.
         """
-        body_index_to_be_deleted = []
-        scope_checker = AstScopeChecker("self")
-        for body_index, body in enumerate(bodies):
-            if body_index == super_index:
-                continue  # ignoring super.__init__()
-            if isinstance(body, ast.If):
-                if scope_checker.check(body.test):
-                    self._replace_ori_field_of_init_func(stree, body.body, -1)
-                    self._replace_ori_field_of_init_func(stree, body.orelse, -1)
-                    continue
-                else:
-                    logger.info("Ignoring un-eval-able if: %s", astunparse.unparse(body.test))
-            if not isinstance(body, ast.Assign):  # if not assign node, delete
-                body_index_to_be_deleted.append(body_index)
-                continue
-            if len(body.targets) != 1:
-                raise RuntimeError("not support multi-targets in assign now!", father_node=body)
-            target = body.targets[0]
-            if not isinstance(target, ast.Attribute):  # only keep class member
-                body_index_to_be_deleted.append(body_index)
-                continue
-            if not isinstance(target.value, ast.Name):
-                logger.info(f"Ignoring {astunparse.unparse(target)} in __init__ function.")
-                body_index_to_be_deleted.append(body_index)
-                continue
-            target_value: ast.Name = target.value
-            if target_value.id != "self":
-                body_index_to_be_deleted.append(body_index)
-                continue
-            field_name = target.attr
-            body.value = ast.Call(ast.Name('getattr', ast.Load()),
-                                  [ast.Attribute(ast.Name('self', ast.Load()), '_handler', ast.Load()),
-                                   ast.Constant(value=field_name, kind=None)], [])
-        for counter, index in enumerate(body_index_to_be_deleted):
-            bodies.pop(index - counter)
-        ClassDefParser._remove_empty_ast_in_init_func(bodies)
-
-    @staticmethod
-    def _insert_handler_to_init_func(ast_init_fn: ast.FunctionDef, super_index):
-        """Insert 'self._handler = global_vars.get('handler')' to init ast.FunctionDef.body"""
-        if super_index == -1:
-            super_index = 0
-        AstModifier.insert_assign_to_function(ast_init_fn, [ScopedValue.create_naming_value("_handler", "self")],
-                                              ScopedValue.create_naming_value("get", "global_vars"),
-                                              [ScopedValue.create_variable_value("handler")], None,
-                                              ast_init_fn.body[super_index], False)
-
-    def process(self, stree: SymbolTree, node: ast.ClassDef):
-        """
-        Parse init and construct in ast.ClassDef.
 
-        Args:
-            stree ([SymbolTree]): Symbol Tree under parsing.
-            node ([ast.ClassDef]): An ast.ClassDef node.
-        """
-        replacer = AstReplacer(node)
-        replacer.replace_all(stree.get_ori_cls_name(), stree.get_opt_cls_name())
-
-        stree.set_class_ast(node)
-
-        for body in node.body:
-            if isinstance(body, ast.FunctionDef):
-                if body.name == "__init__":
-                    self._process_init_func_ast(stree, body)
-                    stree.set_init_func_ast(body)
-                elif body.name == "construct":
-                    parser: Parser = ParserRegister.instance().get_parser(ast.FunctionDef)
-                    parser.process(stree, body)
-                else:
-                    logger.info(
-                        "Ignoring ast.FunctionDef in ast.ClassDef except __init__ and construct function: %s",
-                        body.name)
-            else:
-                logger.info("Ignoring unsupported node(%s) in ast.ClassDef.", type(body).__name__)
+        new_bodies = []
+        for body in bodies:
+            self._handle_bodies_for_replace_ori_field(cls_ast, body, stree, new_bodies)
+        return new_bodies
 
 
 g_classdef_parser = reg_parser(ClassDefParser())

mindspore/rewrite/parsers/for_parser.py

@@ -19,6 +19,7 @@ import astunparse
 from mindspore.rewrite.api.scoped_value import ScopedValue, ValueType
 from mindspore.rewrite.ast_helpers.ast_modifier import AstModifier
 from mindspore import log as logger
+from mindspore import nn
 from ..symbol_tree import SymbolTree
 from ..parser import Parser
 from ..parser_register import reg_parser
@@ -34,12 +35,13 @@ class ForParser(Parser):
     def modify_init_ast(stree, i, obj, iter_var_name):
         """Modify the ast node in init function."""
         target = f"{iter_var_name.strip()}_{str(i)}"
-        stree.add_global_vars(target, obj)
+        setattr(stree.get_origin_network(), target, obj)
         stree.get_origin_network().insert_child_to_cell(target, obj)
         AstModifier.insert_assign_to_function(stree.get_init_func_ast(),
                                               targets=[ScopedValue(ValueType.NamingValue, "self", target)],
-                                              expr=ScopedValue(ValueType.NamingValue, "global_vars", "get"),
-                                              args=[ScopedValue(ValueType.StringValue, "", target)])
+                                              expr=ScopedValue(ValueType.NamingValue, "", "getattr"),
+                                              args=[ScopedValue(ValueType.NamingValue, "", "obj"),
+                                                    ScopedValue(ValueType.StringValue, "", target)])
 
     @staticmethod
     def modify_construct_ast(stree, ast_node, old_name, new_name):
@@ -62,19 +64,24 @@ class ForParser(Parser):
             logger.warning(
                 f"For MindSpore Rewrtie, illegal iteration condition for For node, it must start with{EVAL_WHITE_LIST}")
             return
-        if iter_code.startswith("self"):
+        if "self" in iter_code:
             iter_code = iter_code.replace("self", "stree.get_origin_network()")
         try:
             iter_obj = eval(iter_code)
-        except Exception as e:
-            error_info = f"For MindSpore Rewrtie, when eval '{iter_code}' by using JIT Fallback feature, " \
+        except (NameError, TypeError) as e:
+            _info = f"For MindSpore Rewrtie, when eval '{iter_code}' by using JIT Fallback feature, " \
                          f"an error occurred: {str(e)}"
-            logger.error(error_info)
-            raise e
+            logger.warning(_info)
+            stree.try_append_python_node(node, node)
+            return
 
         iter_var_name = iter_code.split(".")[-1]
         index = stree.get_ast_root().body.index(node) + 1
         if isinstance(iter_obj, list):
+            for obj in iter_obj:
+                if not isinstance(obj, nn.Cell):
+                    stree.try_append_python_node(node, node)
+                    return
             for i, obj in enumerate(iter_obj):
                 ForParser.modify_init_ast(stree, i, obj, iter_var_name)
                 for body in node.body:
@@ -84,14 +91,17 @@ class ForParser(Parser):
                     index += 1
             if stree.get_ori_cls_name() == "SequentialCell":
                 stree.on_change(Event.CodeChangeEvent)
-        elif isinstance(iter_obj, range):
-            raise NotImplementedError("For MindSpore Rewrtie, range not support")
+            stree.get_ast_root().body.remove(node)
+            return
+        if isinstance(iter_obj, range):
+            logger.warning("For MindSpore Rewrite, range not support.")
         elif isinstance(iter_obj, zip):
-            raise NotImplementedError("For MindSpore Rewrtie, zip not support")
+            logger.warning("For MindSpore Rewrite, zip not support.")
         elif isinstance(iter_obj, enumerate):
-            raise NotImplementedError("For MindSpore Rewrtie, enumerate not support")
+            logger.warning("For MindSpore Rewrite, enumerate not support.")
         else:
-            raise ValueError("For MindSpore Rewrtie, not supported type: ", iter_obj)
-
+            logger.warning(f"For MindSpore Rewrite, not supported type: {type(iter_obj).__name__}")
+        stree.try_append_python_node(node, node)
+        return
 
 g_for_parser = reg_parser(ForParser())

mindspore/rewrite/parsers/function_def_parser.py

@@ -14,10 +14,11 @@
 # ============================================================================
 """Parse bodies of ast.FunctionDef which is construct function to nodes of SymbolTree."""
 import ast
-
+from mindspore import log as logger
 from ..parser_register import ParserRegister, reg_parser
 from ..parser import Parser
 from ..symbol_tree import SymbolTree
+from ..api.node_type import NodeType
 
 
 class FunctionDefParser(Parser):
@@ -27,6 +28,24 @@ class FunctionDefParser(Parser):
         """Parse target type"""
         return ast.FunctionDef
 
+    def remove_dead_code(self, stree: SymbolTree):
+        """Remove dead codes"""
+        # Find out return node position
+        return_idx = -1
+        for idx, node in enumerate(stree.nodes()):
+            if node.get_node_type() == NodeType.Output:
+                return_idx = idx
+                break
+        if return_idx == -1:
+            return
+        # Remove nodes after return node.
+        # Reverse traversal to ensure that nodes are orphaned and can be deleted.
+        for idx, node in reversed(list(enumerate(stree.nodes()))):
+            if idx <= return_idx:
+                break
+            logger.info(f"Remove dead code node:{node.get_name()}")
+            stree.erase_node(node)
+
     def process(self, stree: SymbolTree, node: ast.FunctionDef):
         """Parse bodies of ast.FunctionDef which is construct function to nodes of SymbolTree."""
         stree.set_ast_root(node)
@@ -44,13 +63,11 @@ class FunctionDefParser(Parser):
             else:
                 parser.process(stree, body)
 
-        for body in node.body:
-            if isinstance(body, (ast.For, ast.If)):
-                node.body.remove(body)
         if hasattr(node, "decorator_list"):
             stree.try_append_python_node(node, node.decorator_list)
         if hasattr(node, "returns"):
             stree.try_append_python_node(node, node.returns)
+        self.remove_dead_code(stree)
 
 
 g_functiondef_parser = reg_parser(FunctionDefParser())

mindspore/rewrite/parsers/if_parser.py

@@ -46,14 +46,18 @@ class IfParser(Parser):
         bodies = None
         try:
             test_value = eval(test_code)
-        except NameError:
+        except (NameError, TypeError):
             stree.try_append_python_node(node, node)
             return
 
         bodies = node.body if test_value else node.orelse
         index = stree.get_ast_root().body.index(node) + 1
+        info_node = ast.Name(id=f"# If node has been replaced by {bool(test_value)} branch.",
+                             lineno=0, col_offset=0, ctx=ast.Load)
+        exp_node = ast.Expr(value=info_node, lineno=0, col_offset=0, ctx=ast.Load)
+        stree.get_ast_root().body.insert(index-1, exp_node)
         for body in bodies:
             stree.get_ast_root().body.insert(index, body)
             index += 1
-
+        stree.get_ast_root().body.remove(node)
 g_if_parser = reg_parser(IfParser())