mindspore 2.2.11__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

mindspore/{ops/_op_impl/tbe/atomic_addr_clean.py → _extends/pijit/__init__.py}

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2024 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.
@@ -12,22 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
+"""
+Helper module for pijit analyze
+"""
 
-"""AtomicAddrClean op"""
-from mindspore.ops.op_info_register import op_info_register, TBERegOp
 
-atomic_addr_clean_op_info = TBERegOp("AtomicAddrClean") \
-    .fusion_type("ELEMWISE") \
-    .async_flag(False) \
-    .binfile_name("atomic_addr_clean.so") \
-    .compute_cost(10) \
-    .kernel_name("atomic_addr_clean") \
-    .partial_flag(True) \
-    .attr("automic_add_mem_size", "required", "listInt64", "all", "[]") \
-    .get_op_info()
+from .pijit_func_white_list import _func_map as pijit_func_white_list_map
 
 
-@op_info_register(atomic_addr_clean_op_info)
-def _atomic_addr_clean_tbe():
-    """AtomicAddrClean TBE register"""
-    return
+__all__ = ['pijit_func_white_list_map']

mindspore/_extends/pijit/pijit_func_white_list.py

@@ -0,0 +1,343 @@
+# This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
+#
+# Copyright 2020-2024 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.
+# ============================================================================
+"""The module of parser python object, called by c++."""
+
+import collections
+import types
+import math
+import os
+import numpy
+from mindspore.nn import GraphCell, Cell
+from mindspore.ops.primitive import Primitive, constexpr, _primexpr
+from mindspore.ops.composite.base import GradOperation, _Grad
+from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.common.api import jit
+from mindspore.common.tensor import Tensor
+from mindspore.common._register_for_tensor import Registry, tensor_operator_registry
+from mindspore._c_expression import MetaFuncGraph_, function_id, Primitive_, PrimitiveFunction_
+from mindspore._c_expression import Tensor as Tensor_
+from mindspore._extends.parse.resources import convert_object_map
+from mindspore import _checkparam as validator
+from mindspore import Parameter, ParameterTuple
+from mindspore.common.initializer import Zero
+
+
+def _get_after_grad_code():
+    """Get the code object of 'after_grad'"""
+    name = "after_grad"
+    codes = []
+    for cnst in GradOperation.__call__.__code__.co_consts:
+        if isinstance(cnst, types.CodeType) and cnst.co_name == name:
+            codes.append(cnst)
+    for cnst in _Grad.__call__.__code__.co_consts:
+        if isinstance(cnst, types.CodeType) and cnst.co_name == name:
+            codes.append(cnst)
+    if not codes:
+        raise RuntimeError("check GradOperation, can't find the code of 'after_grad'")
+    return codes
+
+
+def _get_psjit_code():
+    """Get the code object of 'staging_specialize'"""
+    @jit
+    def inner():
+        pass
+    return inner.__code__
+
+
+def _get_constexpr_code():
+    """Get the code object of '@constexpr'"""
+    @constexpr
+    def inner():
+        pass
+    code = inner.__call__.__code__
+    # check it before c++ use it
+    if not isinstance(inner, Primitive) or code is Primitive.__call__.__code__:
+        raise RuntimeError("@constexpr not isinstance(inner, Primitive) or code is Primitive.__call__.__code__")
+    return code
+
+
+def _get_primexpr_code():
+    """Get the code object of '@_primexpr'"""
+    @_primexpr
+    def inner():
+        pass
+    code = inner.__call__.__code__
+    # check it before c++ use it
+    if not isinstance(inner, Primitive) or code is Primitive.__call__.__code__:
+        raise RuntimeError("@_primexpr not isinstance(inner, Primitive) or code is Primitive.__call__.__code__")
+    return code
+
+
+def _pijit_constexpr():
+    """Placeholder for uniqure id"""
+
+
+def _get_pijit_constexpr_code():
+    codes = []
+    for cnst in validator.check_transpose_axis.__code__.co_consts:
+        if isinstance(cnst, types.CodeType) and cnst.co_name == "_check_dim":
+            codes.append(cnst)
+    return codes
+
+
+def _get_ms_api():
+    """Get ms api"""
+    target_types = Cell, types.FunctionType, Primitive_, PrimitiveFunction_
+    results = []
+    from mindspore.ops import operations as P
+    from mindspore.ops import functional as F
+    from mindspore.ops import composite as C
+    mods = P, F, C
+    for mod in mods:
+        for i in mod.__all__:
+            f = getattr(mod, i)
+            if isinstance(f, target_types):
+                results.append(f)
+    for f in tensor_operator_registry.values():
+        if isinstance(f, target_types):
+            results.append(f)
+    return results
+
+
+psjit_code = _get_psjit_code()
+constexpr_code = _get_constexpr_code()
+primexpr_code = _get_primexpr_code()
+
+primitive_key = id(Primitive.__call__)
+constexpr_key = id(constexpr_code)
+primexpr_key = id(primexpr_code)
+meta_func_graph_key = id(MetaFuncGraph_)
+pijit_forbidden_key = id(NotImplemented)
+pijit_constexpr_key = id(_pijit_constexpr)
+
+
+# check WrapperDescriptor: function_id(tuple.__getitem__) == function_id(tuple().__getitem__)
+# check MethodDescriptor: function_id(list.__getitem__) == function_id(list().__getitem__)
+# check instancemethod: function_id(Tensor_.from_numpy) == function_id(Tensor_(1).from_numpy)
+# check cfunction filter: function_id(Tensor_.from_numpy) != function_id(Tensor_._is_test_stub)
+# check function id: function_id(Tensor.astype) == function_id(Tensor(1).astype) == id(Tensor.astype)
+# check user defined object id: function_id(Primitive) == function_id(Primitive) == id(Primitive)
+
+
+FUNC_KEY_EMPTY = 0  # ""
+FUNC_KEY_PIJIT_CONSTEXPR = 1  # "pijit.constexpr"
+FUNC_KEY_PIJIT_FORBIDDEN = 2  # "pijit.forbidden"
+FUNC_KEY_BUILTIN_FUNC = 3  # "builtin.func"
+FUNC_KEY_LIST_APPEND = 4  # "list.append"
+FUNC_KEY_DICT_POP = 5  # "dict.pop"
+FUNC_KEY_PRIMITIVE = 6  # "mindspore._c_expression.Primitive_"
+FUNC_KEY_META_FUNCG_RAPH = 7  # "mindspore._c_expression.MetaFuncGraph_"
+FUNC_KEY_PSJIT_CODE = 8  # "mindspore.common.api.jit.<locals>.staging_specialize"
+FUNC_KEY_CONSTEXPR = 9  # "mindspore.ops.primitive.constexpr"
+FUNC_KEY_PRIMEXPR = 10  # "mindspore.ops.primitive._primexpr"
+FUNC_KEY_GET_CACHE_PRIM = 11  # "mindspore.ops._primitive_cache._get_cache_prim"
+FUNC_KEY_REGISTRY_GET = 12  # "mindspore.common._register_for_tensor.Registry.get"
+FUNC_KEY_TENSOR_ASTYPE = 13  # "mindspore.common.tensor.Tensor.astype"
+FUNC_KEY_GRAD_OPERATIONS_CODE = 14 # "mindspore.ops.composite.base._Grad.__call__.<locals>.after_grad"
+FUNC_KEY_PSJIT_CONVERTMAP = 15 # "mindspore._extends.parse.resources.convert_object_map"
+FUNC_KEY_GRAPH_CELL = 16  # "mindspore.nn.cell.GraphCell"
+FUNC_KEY_MS_API = 17  # mindspore common api
+FUNC_KEY_MAPPING_GET = 18 # collections.abc.Mapping.get
+
+# Initialized only once. This map will initialize by c++ when start pijit.
+# key is customer if fuzzy match. (Primitive, constexpr, primexpr, MetaFuncGraph)
+# key is id of code for nest object. (jit.<locals>.staging_specialize, GradOperation.__call__.<locals>.after_grad)
+# key is id of object for callalbe object.
+# key is cfunction pointer for builtin_function or method. (isinstance, tuple.__getitem__, Tensor_.asnumpy)
+_func_map = {
+    # special function
+    pijit_constexpr_key: FUNC_KEY_PIJIT_CONSTEXPR,
+    pijit_forbidden_key: FUNC_KEY_PIJIT_FORBIDDEN,
+    primitive_key: FUNC_KEY_PRIMITIVE,
+    constexpr_key: FUNC_KEY_CONSTEXPR,
+    primexpr_key: FUNC_KEY_PRIMEXPR,
+    meta_func_graph_key: FUNC_KEY_META_FUNCG_RAPH,
+    id(GraphCell.__call__): FUNC_KEY_GRAPH_CELL,
+    id(psjit_code): FUNC_KEY_PSJIT_CODE,
+    id(_get_cache_prim): FUNC_KEY_GET_CACHE_PRIM,
+    id(Registry.get): FUNC_KEY_REGISTRY_GET,
+
+    # Tensor method
+    id(Tensor.astype): FUNC_KEY_TENSOR_ASTYPE,
+
+    # types.BuiltinFunctionType
+    function_id(isinstance): FUNC_KEY_BUILTIN_FUNC,
+    function_id(issubclass): FUNC_KEY_BUILTIN_FUNC,
+    function_id(len): FUNC_KEY_BUILTIN_FUNC,
+    function_id(abs): FUNC_KEY_BUILTIN_FUNC,
+    function_id(max): FUNC_KEY_BUILTIN_FUNC,
+    function_id(all): FUNC_KEY_BUILTIN_FUNC,
+    function_id(any): FUNC_KEY_BUILTIN_FUNC,
+    function_id(hash): FUNC_KEY_BUILTIN_FUNC,
+    function_id(id): FUNC_KEY_BUILTIN_FUNC,
+    function_id(ord): FUNC_KEY_BUILTIN_FUNC,
+    function_id(callable): FUNC_KEY_BUILTIN_FUNC,
+    function_id(getattr): FUNC_KEY_BUILTIN_FUNC,
+    function_id(hasattr): FUNC_KEY_BUILTIN_FUNC,
+
+    # types.MethodDescriptorType, types.WrapperDescriptorType
+    function_id(tuple.__getitem__): FUNC_KEY_BUILTIN_FUNC,
+    function_id(tuple.count): FUNC_KEY_BUILTIN_FUNC,
+    function_id(tuple.index): FUNC_KEY_BUILTIN_FUNC,
+    function_id(list.__getitem__): FUNC_KEY_BUILTIN_FUNC,
+    function_id(list.copy): FUNC_KEY_BUILTIN_FUNC,
+    function_id(list.index): FUNC_KEY_BUILTIN_FUNC,
+    function_id(list.count): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.__contains__): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.__getitem__): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.get): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.keys): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.values): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.items): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.fromkeys): FUNC_KEY_BUILTIN_FUNC,
+    function_id(dict.copy): FUNC_KEY_BUILTIN_FUNC,
+    function_id(set.__contains__): FUNC_KEY_BUILTIN_FUNC,
+    function_id(set.copy): FUNC_KEY_BUILTIN_FUNC,
+    function_id(set.issubset): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.find): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.count): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.index): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.rfind): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.rindex): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.startswith): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.endswith): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isascii): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.islower): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isupper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.istitle): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isspace): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isdecimal): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isdigit): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isnumeric): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isalpha): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isalnum): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isidentifier): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.isprintable): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.format): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.format_map): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.__format__): FUNC_KEY_BUILTIN_FUNC,
+    function_id(list.append): FUNC_KEY_LIST_APPEND,
+    function_id(dict.pop): FUNC_KEY_DICT_POP,
+
+    # instancemethod
+    function_id(Tensor_._flatten_tensors): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_._is_flattened): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_._get_flattened_tensors): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_._get_fusion_size): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_._is_test_stub): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_.__str__): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_.__repr__): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
+    function_id(Tensor_.convert_bytes_to_tensor): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.dim): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.from_numpy): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.getitem_index_info): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.get_bytes): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.is_init): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.is_contiguous): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.stride): FUNC_KEY_BUILTIN_FUNC,
+    # Tensor_.asnumpy need real tensor value
+
+    # other builtin function
+    function_id(collections.abc.Mapping.get): FUNC_KEY_MAPPING_GET,
+    function_id(math.log): FUNC_KEY_BUILTIN_FUNC,
+
+    function_id(numpy.isinf): FUNC_KEY_BUILTIN_FUNC,
+    function_id(numpy.isnan): FUNC_KEY_BUILTIN_FUNC,
+    function_id(numpy.abs): FUNC_KEY_BUILTIN_FUNC,
+    function_id(numpy.log): FUNC_KEY_BUILTIN_FUNC,
+
+    # const function
+    function_id(os.getenv): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(validator.check_number_range): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(validator.check_is_int): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(validator.check_is_number): FUNC_KEY_PIJIT_CONSTEXPR,
+}
+
+for after_grad in _get_after_grad_code():
+    _func_map[id(after_grad)] = FUNC_KEY_GRAD_OPERATIONS_CODE
+
+for k, v in convert_object_map.items():
+    key = id(k)
+    if key not in _func_map and isinstance(v, Primitive):
+        if key is print:
+            continue
+        _func_map[key] = FUNC_KEY_PSJIT_CONVERTMAP
+
+for const_code in _get_pijit_constexpr_code():
+    _func_map[id(const_code)] = FUNC_KEY_PIJIT_CONSTEXPR
+
+GUARD_KEY_RELAX_FUNC = 1
+_guard_func_map = dict()
+
+
+def infer_after_grad(after_grad_func, inputs, outputs):
+    """ Infer after_grad """
+    contents = {}
+    for index, name in enumerate(after_grad_func.__code__.co_freevars):
+        contents[name] = after_grad_func.__closure__[index].cell_contents
+    grad = contents.get('grad_') if 'grad_' in contents else contents.get('self')
+    grad_position = contents.get('grad_position')
+    weights = contents.get('weights')
+
+    def _validation_check():
+        if not grad:
+            raise ValueError("Failed to get grad")
+        if grad.has_aux and len(inputs) <= 1:
+            raise ValueError("Hax_aux must have more than one input")
+        if grad_position is None and weights is None:
+            raise ValueError("Ethither grad_position or weights must not be empty")
+        if grad.return_ids and grad_position is None:
+            raise ValueError("return_ids must have grad_position not empty")
+
+    def _to_result(grads):
+        if not grads:
+            return None
+        if len(grads) == 1:
+            return grads[0]
+        return tuple(grads)
+
+    _validation_check()
+
+    input_grads = []
+    if isinstance(grad_position, tuple):
+        for pos in grad_position:
+            input_grads.append((pos, inputs[pos]) if grad.return_ids else inputs[pos])
+    elif isinstance(grad_position, int):
+        input_grads.append((grad_position, inputs[grad_position]) if grad.return_ids else inputs[grad_position])
+
+    param_grads = []
+    if isinstance(weights, Parameter):
+        param_grads.append(Tensor(dtype=weights.dtype, shape=weights.shape, init=Zero()))
+    elif isinstance(weights, (ParameterTuple, list)):
+        param_grads = [Tensor(dtype=weight.dtype, shape=weight.shape, init=Zero()) for weight in weights]
+    grads_output = None
+    get_all = input_grads and param_grads
+    if get_all:
+        grads_output = (_to_result(input_grads), _to_result(param_grads))
+    elif input_grads:
+        grads_output = _to_result(input_grads)
+    elif param_grads:
+        grads_output = _to_result(param_grads)
+
+    if grad.get_value:
+        return outputs, grads_output
+
+    if grad.has_aux:
+        aux_output = tuple(list(inputs)[1:])
+        return grads_output, aux_output
+    return grads_output

mindspore/_extends/remote/kernel_build_server.py

@@ -154,6 +154,8 @@ class AkgBuilderBase():
             messager.send_ack()
             wait_time_str = messager.get_message()
             messager.send_ack()
+            self.platform = messager.get_message()
+            messager.send_ack()
             self.create(int(process_num_str), int(wait_time_str))
         elif arg == 'AKG/ATTR':
             messager.send_ack()

mindspore/_profiler.py

@@ -0,0 +1,30 @@
+# Copyright 2023 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.
+# ============================================================================
+"""python profiler"""
+
+
+from mindspore._c_expression import PythonProfilerRecorder
+
+
+def ms_profile(fn):
+    def wrap_profile_function(*args, **kwargs):
+        recorder = PythonProfilerRecorder(fn.__qualname__)
+
+        recorder.record_start()
+        ret = fn(*args, **kwargs)
+        recorder.record_end()
+        return ret
+
+    return wrap_profile_function

mindspore/amp.py

@@ -20,6 +20,7 @@ from abc import ABC, abstractmethod
 from mindspore.common import mutable
 from mindspore.ops._primitive_cache import _get_cache_prim
 from mindspore.ops.operations.math_ops import NPUGetFloatStatusV2, NPUClearFloatStatusV2
+from mindspore.ops.operations.nn_ops import AllFinite
 from mindspore import _checkparam as validator
 from mindspore._c_expression import MSContext
 from .common import dtype as mstype
@@ -39,13 +40,18 @@ _partial = ops.Partial()
 
 
 @constexpr
-def _ascend_910A_target():
+def _ascend_target():
+    return context.get_context("device_target") == "Ascend"
+
+
+@constexpr
+def _ascend_910a_target():
     return MSContext.get_instance().get_ascend_soc_version() == "ascend910"
 
 
 @constexpr
-def _ascend_910B_target():
-    return MSContext.get_instance().get_ascend_soc_version() == "ascend910b"
+def _ascend_910bc_target():
+    return MSContext.get_instance().get_ascend_soc_version() in ["ascend910b", "ascend910c"]
 
 
 @constexpr
@@ -53,6 +59,22 @@ def _gpu_target():
     return context.get_context("device_target") == "GPU"
 
 
+@constexpr
+def _enable_all_finite():
+    """check whether enable all finite"""
+    runtime_conf = os.environ.get('MS_DEV_RUNTIME_CONF')
+    global_jit_config = context.get_jit_config()
+    if runtime_conf is not None and ("all_finite:True" in runtime_conf or "all_finite:true" in runtime_conf):
+        return True
+
+    if runtime_conf is not None and ("all_finite:False" in runtime_conf or "all_finite:false" in runtime_conf):
+        return False
+
+    if global_jit_config:
+        return global_jit_config["jit_level"] == "O0" or global_jit_config["jit_level"] == "O1"
+    return False
+
+
 def _grad_unscale(scale, grad):
     return grad * ops.Reciprocal()(scale).astype(grad.dtype)
 
@@ -79,22 +101,27 @@ def _overflow(inputs):
 
 
 @jit
-def _all_finite(inputs, check_overflow_mode):
+def _all_finite(inputs, check_overflow_mode, enable_allfinite):
     """all finite check"""
-    if (_ascend_910A_target()) or \
-       (_ascend_910B_target() and check_overflow_mode != "INFNAN_MODE"):
-        status = Tensor([0] * 8, mstype.int32)
-        status = ops.depend(status, inputs)
-        get_status = _get_cache_prim(NPUGetFloatStatusV2)()(status)
-        status = ops.depend(status, get_status)
-        clear_status = _get_cache_prim(NPUClearFloatStatusV2)()(status)
-        get_status = ops.depend(get_status, clear_status)
-        status_finite = get_status.equal(Tensor(0, mstype.int32)).all()
-        return status_finite
-
-    outputs = _hypermap(_partial(_overflow), inputs)
-    flag_sum = ops.addn(outputs).reshape(())
-    status_finite = ops.less(flag_sum, 1)
+    if _ascend_target():
+        if (_ascend_910a_target()) or \
+           (_ascend_910bc_target() and check_overflow_mode == "SATURATION_MODE"):
+            status = Tensor([0] * 8, mstype.int32)
+            status = ops.depend(status, inputs)
+            get_status = _get_cache_prim(NPUGetFloatStatusV2)()(status)
+            status = ops.depend(status, get_status)
+            clear_status = _get_cache_prim(NPUClearFloatStatusV2)()(status)
+            get_status = ops.depend(get_status, clear_status)
+            status_finite = get_status.equal(Tensor(0, mstype.int32)).all()
+            return status_finite
+
+    status_finite = False
+    if enable_allfinite:
+        status_finite = ~AllFinite()(inputs)  # pylint: disable=invalid-unary-operand-type
+    else:
+        outputs = _hypermap(_partial(_overflow), inputs)
+        flag_sum = ops.addn(outputs).reshape(())
+        status_finite = ops.less(flag_sum, 1)
     return status_finite
 
 
@@ -126,11 +153,11 @@ def all_finite(inputs):
 
     Tutorial Examples:
         - `Automatic Mix Precision - Loss Scaling
-          <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+          <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
     """
     inputs = mutable(inputs)
     _check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
-    return _all_finite(inputs, _check_overflow_mode)
+    return _all_finite(inputs, _check_overflow_mode, _enable_all_finite())
 
 
 @jit_class
@@ -142,7 +169,7 @@ class LossScaler(ABC):
     to scale and unscale the loss value and gradients to avoid overflow, `adjust` is used to update the
     loss scale value.
 
-    For more information, refer to the `tutorials  <https://mindspore.cn/tutorials/en/r2.2/advanced/
+    For more information, refer to the `tutorials  <https://mindspore.cn/tutorials/en/master/advanced/
     mixed_precision.html#loss-scaling>`_.
 
     .. warning::
@@ -334,7 +361,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
         """
         inputs = mutable(inputs)
         return _grad_scale_map(self.scale_value, inputs)
@@ -351,7 +378,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
         """
         inputs = mutable(inputs)
         return _grad_unscale_map(self.scale_value, inputs)
@@ -365,7 +392,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
         """
         one = ops.ones((), self.scale_value.dtype)
         scale_mul_factor = self.scale_value * self.scale_factor

mindspore/boost/adasum.py

@@ -23,7 +23,7 @@ from mindspore.communication.management import create_group
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
-from mindspore.ops.operations._inner_ops import Send, Receive
+from mindspore.ops import Send, Receive
 
 
 __all__ = ["AdaSum"]

mindspore/boost/base.py

@@ -183,7 +183,7 @@ class ParameterProcess:
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.boost import ParameterProcess
         >>>
         >>> class Net(nn.Cell):

mindspore/boost/boost_cell_wrapper.py

@@ -136,7 +136,7 @@ class BoostTrainOneStepCell(TrainOneStepCell):
         >>> from mindspore import boost
         >>> from mindspore import 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()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> optim = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
@@ -435,7 +435,7 @@ class BoostTrainOneStepWithLossScaleCell(BoostTrainOneStepCell):
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.nn import WithLossCell
         >>> from mindspore import dtype as mstype
         >>> from mindspore import boost

mindspore/boost/grad_freeze.py

@@ -206,7 +206,7 @@ class GradientFreeze:
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.nn import WithLossCell
         >>> from mindspore import dtype as mstype
         >>> from mindspore import boost
@@ -347,7 +347,7 @@ def freeze_cell(reducer_flag, network, optimizer, sens, grad, use_grad_accumulat
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.boost.grad_freeze import freeze_cell, FreezeOpt
         >>>
         >>> class Net(nn.Cell):

mindspore/boost/group_loss_scale_manager.py

@@ -94,7 +94,7 @@ class GroupLossScaleManager(Cell):
         ...                        loss_scale_manager=loss_scale_manager,
         ...                        boost_level="O1", boost_config_dict=boost_config_dict)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> model.train(2, dataset)
     """

mindspore/boost/less_batch_normalization.py

@@ -96,7 +96,7 @@ class LessBN(Cell):
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.nn import WithLossCell
         >>> from mindspore import dtype as mstype
         >>> from mindspore import boost
@@ -125,11 +125,14 @@ class LessBN(Cell):
     def __init__(self, network, fn_flag=False):
         super(LessBN, self).__init__()
         self.network = network
-        self.network.set_boost("less_bn")
-        self.network.update_cell_prefix()
-        if fn_flag:
-            self._convert_to_less_bn_net(self.network)
-        self.network.add_flags(defer_inline=True)
+        if "less_bn" not in self.network.get_flags():
+            self.network.set_boost("less_bn")
+            self.network.update_cell_prefix()
+            if fn_flag:
+                self._convert_to_less_bn_net(self.network)
+            # Please use this interface with caution as it may cause uncontrollable problem.
+            # The purpose here is to prevent the "less_bn" flag from disappearing.
+            self.network.add_flags(defer_inline=True)
 
     @staticmethod
     def _convert_dense(subcell):