mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/train/amp.py

@@ -14,6 +14,11 @@
 # ============================================================================
 """Auto mixed precision."""
 from __future__ import absolute_import
+import inspect
+import types
+from typing import Any
+import functools
+import collections
 
 import mindspore as ms
 from mindspore import nn
@@ -27,8 +32,9 @@ from mindspore.train.loss_scale_manager import DynamicLossScaleManager, LossScal
 from mindspore import boost, context
 from mindspore.ops import operations as P
 from mindspore.ops import Primitive
+from mindspore.ops import auto_generate as gen
 from mindspore import log as logger
-
+from mindspore._c_expression.amp import pop_amp_strategy, push_amp_strategy, create_amp_strategy, AmpLevel
 
 AMP_WHITE_LIST = [
     nn.Conv1d,
@@ -50,19 +56,81 @@ AMP_WHITE_LIST = [
     P.BatchMatMul,
     P.PReLU,
     P.ReLU,
-    P.Ger
+    P.Ger,
 ]
 
-
 AMP_BLACK_LIST = [
     nn.BatchNorm1d,
     nn.BatchNorm2d,
     nn.BatchNorm3d,
-    nn.LayerNorm
+    nn.LayerNorm,
+]
+
+AMP_AUTO_WHITE_LIST = [
+    P.Conv2D,
+    P.Conv3D,
+    P.Conv2DTranspose,
+    P.Conv3DTranspose,
+    gen.Convolution,
+    P.MatMul,
+    gen.MatMulExt,
+    P.BatchMatMul,
+    gen.BatchMatMulExt,
+    gen.PReLU,
+    P.Einsum,
+    gen.Dense,
+    gen.Addmm,
 ]
 
+AMP_AUTO_BLACK_LIST = [
+    gen.Pow,
+    gen.ACos,
+    gen.Asin,
+    gen.Cosh,
+    P.Erfinv,
+    P.Exp,
+    P.Expm1,
+    P.Log,
+    P.Log1p,
+    P.Reciprocal,
+    P.Rsqrt,
+    P.Sinh,
+    P.Tan,
+    P.Softplus,
+    gen.SoftplusExt,
+    P.LayerNorm,
+    gen.LayerNormExt,
+    P.BatchNorm,
+    gen.GroupNorm,
+    P.KLDivLoss,
+    P.SmoothL1Loss,
+    P.MultilabelMarginLoss,
+    P.SoftMarginLoss,
+    P.TripletMarginLoss,
+    P.MultiMarginLoss,
+    P.BCEWithLogitsLoss,
+    P.Pdist,
+    P.Cdist,
+    P.Renorm,
+]
+
+# Indicates which inputs of primitives need to be converted
+AMP_PRIM_ARG_TABLE = collections.defaultdict(list, {})
+
+# Primitives in inner amp black list will not be converted in O2/O3
+_INNER_AMP_BLACK_LIST = []
+
 MS_AMP_BY_REWRITE = False
-_amp_cast_op = P.Cast
+
+
+def amp_cast(value, dtype):
+    """This function is used to insert cast operators for tensors during auto mixed precision."""
+    if isinstance(value, ms.Tensor) and value.dtype in mstype.float_type:
+        return P.Cast()(value, dtype)
+    return value
+
+_amp_cast_op = amp_cast
+
 
 class _OutputTo16(nn.Cell):
     """Wrap cell for amp. Cast network output back to float16."""
@@ -88,278 +156,185 @@ class _OutputTo32(nn.Cell):
         return F.mixed_precision_cast(mstype.float32, out)
 
 
-
-def _allow_mix_precision(node, allowed_list, dtype) -> bool:
+def _operator_need_cast(node, force_cast: bool, white_list=None, black_list=None) -> bool:
     """
-    Check whether current node need do mix precision. Follow conditions need to be satisfied:
-        1) Type of node is one of (Primitive, nn.Cell)
-        2) Node is not Cast Op
-        3) to_float(mindspore.float16) is not set in Cell
+    Check whether current node is a operator that need to be casted. Follow conditions need to be satisfied:
+        1) Type of node is CallPrimitive and type of instance is Primitive
+        2) Type of instance is not P.Cast
+        3) force_cast is True, which means one of upper layer cells is under casting
+        4) white_list exist and type of node is in white_list
+        5) black_list exist and type of node is in not black_list
     """
-    node_inst = node.get_instance()
-    if node_inst in allowed_list:
-        return True
-    if node.get_targets() is None:
+    if node.get_node_type() != ms.rewrite.NodeType.CallPrimitive:
         return False
-    if not issubclass(node.get_instance_type(), (Primitive, nn.Cell)):
+    if not inspect.isclass(node.get_instance_type()):
         return False
-    if isinstance(node_inst, _amp_cast_op):
+    if not issubclass(node.get_instance_type(), Primitive):
         return False
-    if issubclass(node.get_instance_type(), nn.Cell):
-        # if cell is already in allowed_list, it means to_float() is set by amp.
-        # if cell is not in allowed_list, but has to_float(),
-        # it means to_float() is set by user.
-        to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-        if hasattr(node_inst, to_float_flag) and getattr(node_inst, to_float_flag):
-            return False
-    allowed_list.append(node.get_instance())
-    return True
+    if issubclass(node.get_instance_type(), P.Cast):
+        return False
+    if node.get_instance_type() in _INNER_AMP_BLACK_LIST:
+        return False
+    if force_cast:
+        return True
+    if white_list is not None and node.get_instance_type() in white_list:
+        return True
+    if black_list is not None and node.get_instance_type() not in black_list:
+        return True
+    return False
 
 
-def _insert_cast_operator_process(node, dtype):
-    """insert cast for operators in white_list."""
-    dtype_str = "mindspore.bfloat16" if dtype == mstype.bfloat16 else "mindspore.float16"
-    new_cast_node = None
-    stree = node.get_symbol_tree()
-    # insert cast fp16/bf16 before the primitive operators
-    if issubclass(node.get_instance_type(), Primitive):
-        for idx, arg in enumerate(node.get_args()):
-            position = stree.before(node)
-            new_node = _amp_cast_op()
-            cast_args = ms.rewrite.ScopedValue.create_name_values([arg.value, dtype_str], [arg.scope, ""])
-            arg_provider = node.get_handler().get_arg_providers()[idx]
-            if arg_provider and len(arg_provider[0].get_target_users(arg_provider[1])) > 1:
-                cast_targets = [stree.unique_name(str(arg))]
-            else:
-                cast_targets = ms.rewrite.ScopedValue.create_name_values([arg.value], [arg.scope])
-            new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
-                                                             targets=cast_targets,
-                                                             args=cast_args,
-                                                             name='incast_{}{}'.format(node.get_name(), idx))
-            stree.insert(position, new_cast_node)
-            node.set_arg_by_node(idx, new_cast_node)
-    # insert cast fp16/bf16 before the Cell operators
-    elif issubclass(node.get_instance_type(), nn.Cell):
-        node.get_instance().to_float(dtype)
-    # ignore if subclass is not one of (Primitive, nn.Cell)
-    else:
-        return
-
-    # insert cast float32 after the operators
-    position = stree.after(node)
-    new_node = _amp_cast_op()
-    cast_args = ms.rewrite.ScopedValue.create_name_values([node.get_targets()[0].value,
-                                                           "mindspore.float32"])
-    new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
-                                                     targets=[node.get_targets()[0]],
-                                                     args=cast_args,
-                                                     name='outcast_{}'.format(node.get_name()))
-    # insert node & unique names
-    stree.insert(position, new_cast_node)
-    # update argument names
-    for user in node.get_users():
-        if user.get_name() == new_cast_node.get_name():
-            continue
-        for idx, arg in enumerate(user.get_args()):
-            if arg == node.get_targets()[0]:
-                user.set_arg_by_node(idx, new_cast_node)
-
-
-def _insert_cast_operator_white_list(stree, white_list, dtype):
-    """insert cast for operators in white_list."""
-    allowed_list = []
-    # Ignore if net called ".to_float(dtype)"
-    net = stree.get_handler().get_origin_network()
-    to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-    if isinstance(net, nn.Cell) and hasattr(net, to_float_flag) and getattr(net, to_float_flag):
-        return
-    node_list = []
-    node_list.extend(list(stree.nodes()))
-    while node_list:
-        node = node_list.pop()
-        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
-            if MS_AMP_BY_REWRITE:
-                _insert_cast_for_cell_container(node, dtype, allowed_list, white_list=white_list)
-            for n in node.get_handler().node_list:
-                if n.get_node_type() == ms.rewrite.NodeType.Tree:
-                    _insert_cast_operator_white_list(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)),
-                                                     white_list, dtype)
-        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
-            substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
-            _insert_cast_operator_white_list(substree, white_list, dtype)
-        elif node.get_node_type() in [ms.rewrite.NodeType.CallFunction, ms.rewrite.NodeType.ControlFlow]:
-            if isinstance(node.get_handler(), ms.rewrite.node.NodeManager):
-                nodes = [ms.rewrite.Node(n) for n in node.get_handler().nodes()]
-                node_list.extend(nodes)
-        elif node.get_instance_type() in white_list and _allow_mix_precision(node, allowed_list, dtype):
-            _insert_cast_operator_process(node, dtype)
+def _precision_set_by_user(cell_inst: nn.Cell) -> bool:
+    """Check whether cell precision is set by user."""
+    for flag in ["fp32", "fp16", "bf16"]:
+        if hasattr(cell_inst, flag) and getattr(cell_inst, flag):
+            return True
+    return False
 
 
-def _insert_cast_for_cell_container(cell_container, dtype, allowed_list, *, white_list=None, black_list=None):
+def _net_need_cast(node, force_cast: bool, white_list=None, black_list=None) -> bool:
     """
-    Insert cast for cell containers.
-    Only one of white_list and black_list can be set.
+    Check whether current node is type of tree whose network needs to be casted. Follow conditions need to
+    be satisfied:
+        1) Type of node is Tree and type of instance is Cell
+        2) Cell.to_float(xxx) is not set by user
+        3) force_cast is True, which means one of upper layer networks is under casting
+        4) white_list exist and type of node is in white_list
+        5) black_list exist and type of node is in not black_list
     """
+    if node.get_node_type() != ms.rewrite.NodeType.Tree:
+        return False
+    if not inspect.isclass(node.get_instance_type()):
+        return False
+    if not issubclass(node.get_instance_type(), nn.Cell):
+        return False
+    if node.get_instance_type() in _INNER_AMP_BLACK_LIST:
+        return False
+    if _precision_set_by_user(node.get_instance()):
+        return False
+    if force_cast:
+        return True
+    if white_list is not None and node.get_instance_type() in white_list:
+        return True
+    if black_list is not None and node.get_instance_type() not in black_list:
+        return True
+    return False
+
+
+def _insert_cast_for_operator(node, dtype):
+    """insert cast pair for node."""
+    dtype_str = "bfloat16" if dtype == mstype.bfloat16 else "float16"
+    stree = node.get_symbol_tree()
+    # insert cast fp16/bf16 for inputs of node
+    for idx, arg in enumerate(node.get_args()):
+        if arg.type != ms.rewrite.ValueType.NamingValue:
+            continue
+        incast_args = ms.rewrite.ScopedValue.create_name_values([arg.value, dtype_str], [arg.scope, "mindspore"])
+        arg_providers = node.get_arg_providers()
+        if not arg_providers or idx not in arg_providers or \
+            len(arg_providers[idx][0].get_target_users(arg_providers[idx][1])) > 1:
+            # create new target names when argument is used by other node
+            incast_targets = [stree.unique_name(f"{arg.value}_var")]
+        else:
+            incast_targets = ms.rewrite.ScopedValue.create_name_values([arg.value], [arg.scope])
+        incast_node = ms.rewrite.Node.create_call_function(_amp_cast_op, targets=incast_targets, args=incast_args)
+        stree.insert(stree.before(node), incast_node)
+        node.set_arg_by_node(idx, incast_node)
+    # insert cast fp32 for outputs of node
+    for _, target in enumerate(node.get_targets()):
+        if target.type != ms.rewrite.ValueType.NamingValue:
+            continue
+        outcast_args = ms.rewrite.ScopedValue.create_name_values([target.value, "float32"],
+                                                                 [target.scope, "mindspore"])
+        outcast_targets = ms.rewrite.ScopedValue.create_name_values([target.value], [target.scope])
+        outcast_node = ms.rewrite.Node.create_call_function(_amp_cast_op, targets=outcast_targets, args=outcast_args)
+        stree.insert(stree.after(node), outcast_node)
 
-    class CastNet(nn.Cell):
-        """Cast net"""
-        def __init__(self, dtype):
-            super().__init__()
-            self.cast = _amp_cast_op()
-            self.dtype = dtype
-
-        def construct(self, x):
-            return self.cast(x, self.dtype)
-
-    cast_flag = False
-    current_node = None
-    stree = cell_container.get_symbol_tree()
-    for node in cell_container.get_handler().nodes():
-        current_node = ms.rewrite.Node(node)
-        if (white_list is not None and current_node.get_instance_type() in white_list) or \
-           (black_list is not None and current_node.get_instance_type() not in black_list) and \
-           (_allow_mix_precision(current_node, allowed_list, dtype)):
-            cast_flag = True
-            current_node.get_instance().to_float(dtype)
-        elif cast_flag:
-            # cast next node back to float32
-            current_node.get_instance().to_float(mstype.float32)
-            cast_flag = False
-    if cast_flag and current_node:
-        # if last node in cell_container is casted to fp16/bf16, insert a cast node to cast value back to fp32
-        cast_node = ms.rewrite.Node.create_call_cell(cell=CastNet(mstype.float32),
-                                                     args=[current_node.get_targets()[0]],
-                                                     targets=[current_node.get_targets()[0]],
-                                                     name=f"outcast_{cell_container.get_name()}")
-        stree.insert(stree.after(current_node), cast_node)
+
+def _insert_cast_for_operators(stree, dtype, force_cast, *, white_list=None, black_list=None):
+    """insert cast for operators not in black_list."""
+    # get all nodes of stree exclude nodes in subtree.
+    all_nodes = stree.all_nodes(False)
+    for node in all_nodes:
+        if not node.get_targets():
+            continue
+        if _operator_need_cast(node, force_cast, white_list, black_list):
+            _insert_cast_for_operator(node, dtype)
+        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
+            force_cast_ = force_cast or _net_need_cast(node, force_cast, white_list, black_list)
+            if not _precision_set_by_user(node.get_instance()):
+                subtree = node.get_sub_tree()
+                _insert_cast_for_operators(subtree, dtype, force_cast_, white_list=white_list, black_list=black_list)
 
 
 def _need_removed_cast_pair(node, dtype):
     """check whether the cast pairs should be removed."""
-    dtype_str = "mindspore.bfloat16" if dtype == mstype.bfloat16 else "mindspore.float16"
-    cast_dtypes = ms.rewrite.ScopedValue.create_name_values([dtype_str, "mindspore.float32"])
+    dtype_str = "bfloat16" if dtype == mstype.bfloat16 else "float16"
+    cast_dtypes = ms.rewrite.ScopedValue.create_name_values([dtype_str, "float32"], ["mindspore", "mindspore"])
     cast_dtype_f16 = cast_dtypes[0]
     cast_dtype_f32 = cast_dtypes[1]
-    # current node should be Cast Op to float32
+    # current node should be cast fp32
     if node.get_instance_type() != _amp_cast_op:
         return False
     node_cast_type = node.get_args()[1]
     if node_cast_type != cast_dtype_f32:
         return False
-    # all user nodes should be Cast Op to dtype or Cell with to_float(dtype)
+    # all user nodes should be cast fp16/bf16
     if not node.get_users():
         return False
     all_nodes = [ms.rewrite.Node(n) for n in node.get_handler().get_node_manager().nodes()]
     for user in node.get_users():
-        # If ControlFlow node(if statement) exists between current node and user node,
+        # If ControlFlow node(e.g. if, for, while) exists between current node and user node,
         # cast pair should not be removed.
         middle_nodes = all_nodes[all_nodes.index(node): all_nodes.index(user)]
         if any([n.get_node_type() == ms.rewrite.NodeType.ControlFlow for n in middle_nodes]):
             return False
-        if isinstance(user.get_instance(), nn.Cell):
-            to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-            if not (hasattr(user.get_instance(), to_float_flag) and getattr(user.get_instance(), to_float_flag)):
-                return False
-        elif user.get_instance_type() == _amp_cast_op:
-            user_cast_type = user.get_args()[1]
-            if user_cast_type != cast_dtype_f16:
-                return False
-        else:
+        if user.get_instance_type() != _amp_cast_op:
             return False
+        user_cast_type = user.get_args()[1]
+        if user_cast_type != cast_dtype_f16:
+            return False
+        # cast pair detected, check next user
+        continue
     return True
 
 
-def _removed_cast_pair_process(cast_f32_node):
-    """remove the duplicated cast operators."""
-    stree = cast_f32_node.get_symbol_tree()
-    cast_f32_users = cast_f32_node.get_users()
-    # remove cast f16 nodes
-    for user_node in cast_f32_users:
-        if user_node.get_instance_type() == _amp_cast_op:
-            cast_f16_node = user_node
-            # modify arguments using cast_f16's target[0] to cast_f32's args[0], which is f16 type
-            for cast_f16_user in cast_f16_node.get_users():
-                for idx, arg in enumerate(cast_f16_user.get_args()):
-                    if arg == cast_f16_node.get_targets()[0]:
-                        cast_f16_user.set_arg(idx, cast_f32_node.get_args()[0])
-            stree.erase(cast_f16_node)
-        # update args of cell f16 nodes
-        elif isinstance(user_node.get_instance(), nn.Cell):
-            cell_f16_node = user_node
-            for idx, arg in enumerate(cell_f16_node.get_args()):
-                if arg == cast_f32_node.get_targets()[0]:
-                    cell_f16_node.set_arg(idx, cast_f32_node.get_args()[0])
-    # remove the cast f32 node
-    stree.erase(cast_f32_node)
-
-
 def _remove_duplicated_cast(stree, dtype):
     """remove the duplicated cast operators."""
-    node_list = []
-    node_list.extend(list(stree.nodes()))
-    while node_list:
-        node = node_list.pop()
-        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
-            for n in node.get_handler().node_list:
-                if n.get_node_type() == ms.rewrite.NodeType.Tree:
-                    _remove_duplicated_cast(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)), dtype)
-        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
-            substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
-            _remove_duplicated_cast(substree, dtype)
-        elif node.get_node_type() in [ms.rewrite.NodeType.CallFunction, ms.rewrite.NodeType.ControlFlow]:
-            if isinstance(node.get_handler(), ms.rewrite.node.NodeManager):
-                nodes = [ms.rewrite.Node(n) for n in node.get_handler().nodes()]
-                node_list.extend(nodes)
-        elif _need_removed_cast_pair(node, dtype):
-            _removed_cast_pair_process(node)
-
-
-def _auto_white_list(network, white_list, dtype):
-    """process the white list of network."""
-    stree = ms.rewrite.SymbolTree.create(network)
-    _insert_cast_operator_white_list(stree, white_list, dtype)
-    _remove_duplicated_cast(stree, dtype)
-    return stree.get_network()
-
-
-def _insert_cast_operator_black_list(stree, black_list, dtype):
-    """insert cast for operators not in black_list."""
-    allowed_list = []
-    # Ignore if net called ".to_float(dtype)"
-    net = stree.get_handler().get_origin_network()
-    to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-    if isinstance(net, nn.Cell) and hasattr(net, to_float_flag) and getattr(net, to_float_flag):
-        return
-    for node in stree.nodes(all_nodes=True):
-        if node.get_targets() is None:
-            continue
-        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
-            _insert_cast_for_cell_container(node, dtype, allowed_list, black_list=black_list)
-        elif isinstance(node.get_handler().get_node_manager(), ms.rewrite.node.CellContainer):
-            # nodes in CellContainer are processed by _insert_cast_for_cell_container
-            continue
-        elif node.get_instance_type() not in black_list and _allow_mix_precision(node, allowed_list, dtype):
-            _insert_cast_operator_process(node, dtype)
-
-
-def _remove_duplicated_cast_rewrite(stree, dtype):
-    """remove the duplicated cast operators."""
-    for node in stree.nodes(all_nodes=True):
+    all_nodes = list(stree.nodes(all_nodes=True))
+    for node in all_nodes:
         if _need_removed_cast_pair(node, dtype):
-            user_nodes = node.get_users()
-            # remove cast f16 nodes
-            for user_node in user_nodes:
-                if user_node.get_instance_type() == _amp_cast_op:
-                    stree.erase(user_node)
-            # remove the cast f32 node
+            incast_nodes = node.get_users()
+            # remove cast fp16/bf16 nodes
+            for incast_node in incast_nodes:
+                # get_target_users() return {target0: [(user0, arg_idx), ...], ...}
+                target_users = list(incast_node.get_target_users().values())
+                if not target_users or not target_users[0]:
+                    continue
+                for user_node, arg_idx in target_users[0]:
+                    user_node.set_arg(arg_idx, incast_node.get_args()[0])
+                stree.erase(incast_node)
+            # remove the cast fp32 node
             stree.erase(node)
 
 
-def _auto_black_list_rewrite(network, black_list, dtype):
+def _auto_mixed_precision_rewrite(network, dtype, *, white_list=None, black_list=None):
+    """Implement auto mixed precision by rewrite"""
+    if (white_list is None and black_list is None) or (white_list is not None and black_list is not None):
+        raise ValueError("For _auto_mixed_precision_rewrite, one of white_list and black_list must be provided.")
+    # enable rewrite configs for amp
+    ms.rewrite.common.namespace._ms_cells_to_subtree = True
+    ms.rewrite.parsers.assign_parser.AssignParser._share_one_implementation = True
+    # insert casts by rewrite
     stree = ms.rewrite.SymbolTree.create(network)
-    _insert_cast_operator_black_list(stree, black_list, dtype)
-    _remove_duplicated_cast_rewrite(stree, dtype)
-    return stree.get_network()
+    _insert_cast_for_operators(stree, dtype, False, white_list=white_list, black_list=black_list)
+    _remove_duplicated_cast(stree, dtype)
+    new_net = stree.get_network()
+    # disable rewrite configs
+    ms.rewrite.parsers.assign_parser.AssignParser._share_one_implementation = False
+    ms.rewrite.common.namespace._ms_cells_to_subtree = False
+    ms.rewrite.common.config.clear_caches()
+    return new_net
 
 
 def _auto_black_list(network, black_list, dtype):
@@ -381,6 +356,42 @@ def _auto_black_list(network, black_list, dtype):
     return network
 
 
+class amp_decorator:
+    """
+    Auto mixed precision decorator.
+    Type of lists: List[Tuple[str, List[int]]]
+    """
+    def __init__(self, amp_level, amp_dtype, white_list, black_list):
+        self.amp_level = amp_level
+        self.amp_dtype = amp_dtype
+        self.white_list = white_list
+        self.black_list = black_list
+
+    def __enter__(self):
+        push_amp_strategy(self.amp_level, self.amp_dtype, self.white_list, self.black_list)
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
+        pop_amp_strategy()
+
+
+def _set_amp_decorator(obj, amp_level, amp_dtype, white_list, black_list):
+    """
+    Set auto mixed precision context decorator for object.
+    Type of lists: List[Tuple[str, List[int]]]
+    """
+    if inspect.isfunction(obj) or inspect.ismethod(obj):
+        @functools.wraps(obj)
+        def wrapper(*args, **kwargs):
+            with amp_decorator(amp_level, amp_dtype, white_list, black_list):
+                return obj(*args, **kwargs)
+        return wrapper
+    if isinstance(obj, nn.Cell):
+        obj.construct = types.MethodType(
+            _set_amp_decorator(obj.construct.__func__, amp_level, amp_dtype, white_list, black_list), obj)
+        return obj
+    raise TypeError(f"For amp_level '{amp_level}', the network type should be Cell or function, bot got {type(obj)}.")
+
+
 def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
     """
     Returns a network processed with auto mixed precision.
@@ -391,26 +402,44 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
     converted to lower precision float, and calculation results are converted back to full precision float,
     i.e. ``mstype.float32`` .
 
-    The framework has a set of built-in blacklists and whitelists, and the `amp_level` determines which cells and
-    operators are specifically converted.
+    The `amp_level` and its corresponding lists determine which cells and operators are converted.
 
-    The current built-in whitelist contents are:
+    When `amp_level` is set to ``O0``, no cells and operators are converted.
 
-    [:class:`mindspore.nn.Conv1d`, :class:`mindspore.nn.Conv2d`, :class:`mindspore.nn.Conv3d`,
-    :class:`mindspore.nn.Conv1dTranspose`, :class:`mindspore.nn.Conv2dTranspose`,
-    :class:`mindspore.nn.Conv3dTranspose`, :class:`mindspore.nn.Dense`, :class:`mindspore.nn.LSTMCell`,
-    :class:`mindspore.nn.RNNCell`, :class:`mindspore.nn.GRUCell`, :class:`mindspore.ops.Conv2D`,
-    :class:`mindspore.ops.Conv3D`, :class:`mindspore.ops.Conv2DTranspose`,
-    :class:`mindspore.ops.Conv3DTranspose`, :class:`mindspore.ops.MatMul`, :class:`mindspore.ops.BatchMatMul`,
-    :class:`mindspore.ops.PReLU`, :class:`mindspore.ops.ReLU`, :class:`mindspore.ops.Ger`]
+    When `amp_level` is set to ``O1``, cells and operators in whitelist will be converted to lower precision
+    operations. For details on whitelist, refer to :func:`mindspore.amp.get_white_list`.
 
-    The current built-in blacklist contents are:
+    When `amp_level` is set to ``O2``, cells in blacklist will maintain full precision, and cells outside the
+    list will be converted to low precision. For details on blacklist, refer to :func:`mindspore.amp.get_black_list`.
 
-    [:class:`mindspore.nn.BatchNorm1d`, :class:`mindspore.nn.BatchNorm2d`, :class:`mindspore.nn.BatchNorm3d`,
-    :class:`mindspore.nn.LayerNorm`]
+    When `amp_level` is set to ``O3``, all cells will be converted to low precision.
+
+    When `amp_level` is set to ``auto``, operators in `auto_whitelist` will be converted to lower precision
+    operations, operators in `auto_blacklist` will be converted to full precision  operations, operators in
+    `promote_list` will be converted to the higher accuracy float type of the operator inputs, and operators
+    not listed will run in the type defined by their inputs.
+
+    Operators in `auto_whitelist` are:
+
+    ``Conv2D``, ``Conv3D``, ``Conv2DTranspose``, ``Conv3DTranspose``, ``Convolution``, ``MatMul``, ``MatMulExt``,
+    ``BatchMatMul``, ``BatchMatMulExt``, ``PReLU``, ``Einsum``, ``Dense``, ``Addmm``
+
+    Operators in `auto_blacklist` are:
+
+    ``Pow``, ``ACos``, ``Asin``, ``Cosh``, ``Erfinv``, ``Exp``, ``Expm1``, ``Log``, ``Log1p``, ``Reciprocal``,
+    ``Rsqrt``, ``Sinh``, ``Tan``, ``Softplus``, ``SoftplusExt``, ``LayerNorm``, ``LayerNormExt``, ``BatchNorm``,
+    ``GroupNorm``, ``KLDivLoss``, ``SmoothL1Loss``, ``MultilabelMarginLoss``, ``SoftMarginLoss``,
+    ``TripletMarginLoss``, ``MultiMarginLoss``, ``BCEWithLogitsLoss``, ``Pdist``, ``Cdist``, ``Renorm``,
+    ``ReduceProd``, ``Softmax``, ``LogSoftmax``, ``CumProd``, ``CumSum``, ``CumsumExt``, ``ProdExt``, ``SumExt``,
+    ``Norm``
+
+    Operators in `promote_list` are:
+
+    ``Addcdiv``, ``Addcmul``, ``Cross``, ``_PyboostCrossPrim``, ``Dot``, ``GridSampler2D``, ``GridSampler3D``,
+    ``BiasAdd``
 
     For details on automatic mixed precision, refer to
-    `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html>`_ .
+    `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/master/beginner/mixed_precision.html>`_ .
 
     Note:
         - Repeatedly calling mixed-precision interfaces, such as `custom_mixed_precision` and `auto_mixed_precision`,
@@ -418,10 +447,18 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
         - If interfaces like `Model` and `build_train_network` is used to train the network which is converted by
           mixed-precision interfaces such as `custom_mixed_precision` and `auto_mixed_precision`, `amp_level`
           need to be configured to ``O0`` to avoid the duplicated accuracy conversion.
+        - When `amp_level` is set to ``auto``, the output of the network may be lower precision. In this case, you
+          may need to manually convert the type to avoid type inconsistency errors of the loss function.
+        - When `amp_level` is set to ``auto``, and cells in the network are configured with `to_float`, the accuracy
+          specified by `to_float` takes effect first.
+
+    .. warning::
+        ``auto`` level of `amp_level` is an experimental API that is subject to change or deletion.
 
     Args:
-        network (Cell): Definition of the network.
-        amp_level (str): Supports ["O0", "O1", "O2", "O3"]. Default: ``"O0"`` .
+        network (Union[Cell, function]): Definition of the network. Function type is supported only when `amp_level`
+            is set to ``auto`` .
+        amp_level (str): Supports ["O0", "O1", "O2", "O3", "auto"]. Default: ``"O0"`` .
 
             - "O0": Do not change.
             - "O1": Convert cells and operators in whitelist to lower precision operations, and keep full
@@ -429,25 +466,34 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
             - "O2": Keep full precision operations for cells and operators in blacklist, and convert the rest
               to lower precision operations.
             - "O3": Cast network to lower precision.
+            - "auto": Operators in `auto_whitelist` will be converted to lower precision operations, operators in
+              `auto_blacklist` will be converted to full precision, operators in `promote_list` will be converted
+              to the higher accuracy float type of the operator inputs, and operators not listed will run in the
+              type defined by their inputs.
 
         dtype (Type): The type used in lower precision calculations, can be ``mstype.float16`` or ``mstype.bfloat16`` ,
             default: ``mstype.float16`` .
 
     Raises:
-        TypeError: If `network` is not a Cell.
+        TypeError: If `network` is not a Cell or a function.
         ValueError: If `dtype` is not one of ``mstype.float16`` , ``mstype.bfloat16`` .
         ValueError: If `amp_level` is not within the supported range.
 
     Examples:
         >>> from mindspore import amp
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> network = LeNet5()
         >>> amp_level = "O1"
         >>> net = amp.auto_mixed_precision(network, amp_level)
     """
     if not isinstance(network, nn.Cell):
-        raise TypeError("The network type should be Cell.")
+        if amp_level == "auto":
+            if not inspect.isfunction(network) and not inspect.ismethod(network):
+                raise TypeError("For amp_level 'auto', the network type should be Cell or function.")
+            # function is supported for amp_level 'auto'
+        else:
+            raise TypeError(f"For amp_level '{amp_level}', the network type should be Cell.")
 
     if dtype not in (mstype.float16, mstype.bfloat16):
         raise ValueError(f"The dtype should be one of (mstype.float16, mstype.bfloat16), but got {dtype}.")
@@ -456,27 +502,35 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
         return network
 
     # Return network if the same amp level has already been configurated
-    if getattr(network, "_amp_level") in ("O1", "O2", "O3"):
+    if hasattr(network, "_amp_level") and getattr(network, "_amp_level") in ("O1", "O2", "O3", "auto"):
         logger.warning(f"The network's auto mixed-precision level is adjusted from {getattr(network, '_amp_level')} "
                        f"to {amp_level}, and repeated calls to mixed-precision interfaces can cause performance "
                        f"degradation.")
 
     if amp_level == "O1":
-        network = _auto_white_list(network, AMP_WHITE_LIST, dtype)
+        network = _auto_mixed_precision_rewrite(network, dtype, white_list=AMP_WHITE_LIST)
     elif amp_level == "O2":
         if MS_AMP_BY_REWRITE:
-            network = _auto_black_list_rewrite(network, AMP_BLACK_LIST, dtype)
+            network = _auto_mixed_precision_rewrite(network, dtype, black_list=AMP_BLACK_LIST)
         else:
             network = _auto_black_list(network, AMP_BLACK_LIST, dtype)
             network = _OutputTo32(network)
     elif amp_level == "O3":
         if MS_AMP_BY_REWRITE:
-            network = _auto_black_list_rewrite(network, [], dtype)
+            network = _auto_mixed_precision_rewrite(network, dtype, black_list=[])
         else:
             network.to_float(dtype)
             network = _OutputTo32(network)
+    elif amp_level == "auto":
+        white_list = [(prim.__name__, AMP_PRIM_ARG_TABLE[prim]) for prim in AMP_AUTO_WHITE_LIST]
+        black_list = [(prim.__name__, AMP_PRIM_ARG_TABLE[prim]) for prim in AMP_AUTO_BLACK_LIST]
+        # set amp_strategy attribute for the object
+        amp_strategy = create_amp_strategy(AmpLevel.AmpAuto, dtype, white_list, black_list)
+        setattr(network, "amp_strategy", amp_strategy)
+        # set amp_strategy context decorator for the object
+        network = _set_amp_decorator(network, AmpLevel.AmpAuto, dtype, white_list, black_list)
     else:
-        raise ValueError("The amp level {} is not supported".format(amp_level))
+        raise ValueError(f"The amp level {amp_level} is not supported")
 
     setattr(network, "_amp_level", amp_level)
 
@@ -516,6 +570,10 @@ _config_level = {
     "O3": {
         "keep_batchnorm_fp32": False,
         "cast_model_type": mstype.float16,
+        "loss_scale_manager": None},
+    "auto": {
+        "keep_batchnorm_fp32": False,
+        "cast_model_type": mstype.float32,
         "loss_scale_manager": None}}
 
 
@@ -540,20 +598,11 @@ def _check_kwargs(key_words):
 def _check_level(level, boost_level):
     """Check level."""
     if not isinstance(level, str):
-        raise TypeError("The argument `level` must be a string in ['O0', 'O1', 'O2', 'O3', 'auto'], \
-                         but got type {}.".format(type(level)))
+        raise TypeError(f"The argument `level` must be a string in ['O0', 'O1', 'O2', 'O3', 'auto'],"
+                        f"but got type {type(level)}.")
     validator.check('level', level, "", ['O0', 'O1', 'O2', 'O3', 'auto'], validator.IN)
     validator.check('boost_level', boost_level, "", ['O0', 'O1', 'O2'], validator.IN)
 
-    if level == "auto":
-        device_target = context.get_context('device_target')
-        if device_target == "GPU":
-            level = "O2"
-        elif device_target == "Ascend":
-            level = "O3"
-        else:
-            raise ValueError("Level `auto` only support when `device_target` is GPU or Ascend.")
-
     enable_boost = False
     if boost_level in ["O1", "O2"]:
         enable_boost = True
@@ -578,7 +627,8 @@ def _add_loss_network(network, loss_fn, cast_model_type):
             return self._loss_fn(F.mixed_precision_cast(mstype.float32, out), label)
 
     validator.check_value_type('loss_fn', loss_fn, nn.Cell)
-    if cast_model_type == mstype.float16:
+    if cast_model_type in (mstype.float16, mstype.bfloat16) or \
+       (hasattr(network, "_amp_level") and getattr(network, "_amp_level") in ("O2", "O3", "auto")):
         network = WithLossCell(network, loss_fn)
     else:
         network = nn.WithLossCell(network, loss_fn)
@@ -634,20 +684,10 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
             Default: ``None`` .
         level (str): Supports ['O0', 'O1', 'O2', 'O3', 'auto']. Default: ``'O0'`` .
 
-            - 'O0': Do not change.
-            - 'O1': Cast the operators in white_list to float16, the remaining operators are kept in float32.
-              The operators in the whitelist: [Conv1d, Conv2d, Conv3d, Conv1dTranspose, Conv2dTranspose,
-              Conv3dTranspose, Dense, LSTMCell, RNNCell, GRUCell, MatMul, BatchMatMul, PReLU, ReLU, Ger].
-            - 'O2': Cast network to float16, keep batchnorm and `loss_fn` (if set) run in float32,
-              using dynamic loss scale.
-            - 'O3': Cast network to float16, with additional property `keep_batchnorm_fp32=False` .
-            - 'auto': Set to level to recommended level in different devices. Set level to 'O2' on GPU, Set
-              level to 'O3' Ascend. The recommended level is chosen by the export experience, not applicable to all
-              scenarios. User should specify the level for special network.
-
-            'O2' is recommended on GPU, 'O3' is recommended on Ascend. Property of `keep_batchnorm_fp32`,
-            `cast_model_type` and `loss_scale_manager` determined by `level` setting may be overwritten by settings in
-            `kwargs`.
+            For details on amp level, refer to :func:`mindspore.amp.auto_mixed_precision`.
+
+            Property of `keep_batchnorm_fp32`, `cast_model_type` and `loss_scale_manager` determined by `level`
+            setting may be overwritten by settings in `kwargs`.
 
         boost_level (str): Option for argument `level` in `mindspore.boost` , level for boost mode
             training. Supports ['O0', 'O1', 'O2']. Default: ``'O0'`` .
@@ -670,13 +710,13 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
             take no effect on this property.
 
     Raises:
-        ValueError: If device is CPU, property `loss_scale_manager` is not `None` or `FixedLossScaleManager`
-            (with property `drop_overflow_update=False` ).
+        ValueError: If device is CPU, property `loss_scale_manager` is not `None` or
+            :class:`mindspore.amp.FixedLossScaleManager` (with property `drop_overflow_update=False` ).
 
     Examples:
         >>> from mindspore import amp, nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> network = LeNet5()
         >>> net_loss = nn.SoftmaxCrossEntropyWithLogits(reduction="mean")
         >>> net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
@@ -728,7 +768,7 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
 
 def get_white_list():
     """
-    Provide a copy of internal white list used by auto mixed precision.
+    Provide a copy of internal white list used by auto mixed precision with `amp_level` set to ``O1``.
 
     The current built-in whitelist contents are:
 
@@ -766,7 +806,7 @@ def get_white_list():
 
 def get_black_list():
     """
-    Provide a copy of internal black list used by auto mixed precision.
+    Provide a copy of internal black list used by auto mixed precision with `amp_level` set to ``O2``.
 
     The current built-in blacklist contents are:
 
@@ -789,7 +829,6 @@ def get_black_list():
 
 def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=mstype.float16):
     """
-    Custom mixed precision by setting whitelist or blacklist.
     When the `white_list` is provided, primitives and cells in `white_list` will perform the precision conversion.
     When the `black_list` is provided, cells that are not in `black_list` will perform the pereision conversion.
     Only one of `white_list` and `black_list` should be provided.
@@ -823,7 +862,7 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=m
     Examples:
         >>> from mindspore import amp, nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> custom_white_list = amp.get_white_list()
         >>> custom_white_list.append(nn.Flatten)
@@ -844,11 +883,11 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=m
 
     if white_list is not None:
         _list_check(white_list, "white_list")
-        network = _auto_white_list(network, white_list, dtype)
+        network = _auto_mixed_precision_rewrite(network, dtype, white_list=white_list)
     else:
         _list_check(black_list, "black_list")
         if MS_AMP_BY_REWRITE:
-            network = _auto_black_list_rewrite(network, black_list, dtype)
+            network = _auto_mixed_precision_rewrite(network, dtype, black_list=black_list)
         else:
             network = _auto_black_list(network, black_list, dtype)
             network = _OutputTo32(network)
@@ -883,7 +922,8 @@ def _list_check(custom_list: list, list_name: str):
             if elem not in custom_list:
                 logger.warning(f"{elem} is removed from internal black list.")
 
-def _config_amp(*, enable_rewrite: bool = None, cast_op: type = None): # pylint: disable=unused-variable
+
+def _config_amp(*, enable_rewrite: bool = None, cast_op: types.FunctionType = None): # pylint: disable=unused-variable
     """Configure auto mixed precision."""
     global MS_AMP_BY_REWRITE
     global _amp_cast_op