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

mindspore/ops/operations/custom_ops.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
+# Copyright 2021-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.
@@ -28,6 +28,8 @@ import subprocess
 import numpy as np
 import mindspore as ms
 from mindspore._c_expression import Oplib, typing
+from mindspore._c_expression import pyboost_custom_ext
+from mindspore.common._stub_tensor import _convert_stub
 from mindspore import context
 from mindspore.common import Tensor
 from mindspore.common import dtype as mstype
@@ -156,6 +158,55 @@ def _compile_aot(file):
     return func_path
 
 
+class _CustomExt(ops.PrimitiveWithInfer):
+    """
+    `Custom` primitive is used for PyBoost.
+    """
+
+    def __init__(self, func, out_shape=None, out_dtype=None, bprop=None):
+        super().__init__("CustomExt")
+        self.func = func
+        self.out_shape = out_shape
+        self.out_dtype = out_dtype
+        self.bprop = bprop
+
+    def __infer__(self, *args):
+        if callable(self.out_shape):
+            infer_shape = self.out_shape(*(x["shape"] for x in args))
+        else:
+            infer_shape = self.out_shape
+
+        if callable(self.out_dtype):
+            infer_dtype = self.out_dtype(*(x["dtype"] for x in args))
+        else:
+            infer_dtype = self.out_dtype
+
+        infer_value = None
+        if infer_shape is None:
+            logger.warning("'out_shape' is None. Add a placeholder instead. "
+                           "A CPP version of infer shape function is required "
+                           "in this case.")
+            infer_shape = (1,)
+        # after all automatic infer information fulfillment, throw error if infer_shape/infer_dtype is still None
+        if not isinstance(infer_shape, (tuple, list)):
+            raise TypeError("'out_shape' must be one of [tuple, list, function], but got {}".format(type(infer_shape)))
+
+        if not isinstance(infer_dtype, (typing.Type, tuple, list)):
+            raise TypeError("'out_dtype' must be one of [mindspore.dtype, tuple, list, function], but got {}"
+                            .format(type(infer_dtype)))
+
+        out = {
+            "shape": infer_shape,
+            "dtype": infer_dtype,
+            "value": infer_value,
+        }
+        return out
+
+    def get_bprop(self):
+        """return back propagation function"""
+        return self.bprop
+
+
 class Custom(ops.PrimitiveWithInfer):
     r"""
     `Custom` primitive is used for user defined operators and is to enhance the expressive ability of built-in
@@ -164,7 +215,7 @@ class Custom(ops.PrimitiveWithInfer):
     function if needed. Then these `Custom` objects can be directly used in neural networks.
     Detailed description and introduction of user-defined operators, including correct writing of parameters,
     please refer to `Custom Operators Tutorial
-    <https://www.mindspore.cn/tutorials/experts/en/r2.2/operation/op_custom.html>`_ .
+    <https://www.mindspore.cn/docs/en/master/model_train/custom_program/op_custom.html>`_ .
 
     .. warning::
         - This is an experimental API that is subject to change.
@@ -172,19 +223,11 @@ class Custom(ops.PrimitiveWithInfer):
     .. note::
         The supported platforms are determined by the input `func_type`. The supported platforms are as follows:
 
-        - "hybrid": supports ["Ascend", "GPU", "CPU"].
-        - "akg": supports ["Ascend", "GPU", "CPU"].
-        - "tbe": supports ["Ascend"].
-        - "aot": supports ["GPU", "CPU"].
+        - "hybrid": supports ["GPU", "CPU"].
+        - "akg": supports ["GPU", "CPU"].
+        - "aot": supports ["GPU", "CPU", "Ascend"].
         - "pyfunc": supports ["CPU"].
         - "julia": supports ["CPU"].
-        - "aicpu": supports ["Ascend"].
-
-        If run on ge backend, use `CustomRegOp` to generate the registration information of "aicpu" and "tbe" operator,
-        use `custom_info_register` to bind the registration information to the `func` of the "tbe" operator,
-        then save the registration information of "aicpu" operator and the `func` implementation of "tbe" operator to
-        a file or separate files, keep these files in a separate directory, and set the absolute path of this directory
-        to environment variable "MS_DEV_CUSTOM_OPP_PATH" before running the network.
 
     Args:
         func (Union[function, str]):
@@ -193,16 +236,15 @@ class Custom(ops.PrimitiveWithInfer):
               computation logic of a user defined operator. The function can be one of the following:
 
               1. A AKG operator implementation function, which can use ir builder/tvm compute/hybrid grammar.
-              2. A TBE operator implementation function.
-              3. A pure python function
-              4. An kernel decorated function written by the Hybrid DSL.
+              2. A pure python function
+              3. An kernel decorated function written by the Hybrid DSL.
 
             - str: If func is of str type, then str should be a path of file along with a function name.
               This could be used when func_type is "aot" or "julia".
 
               1. for "aot":
 
-                 Currently "aot" supports GPU/CPU(linux only) platform.
+                 a) GPU/CPU platform.
                  "aot" means ahead of time, in which case Custom directly launches user defined "xxx.so" file as an
                  operator. Users need to compile a handwriting "xxx.cu"/"xxx.cc" file into "xxx.so" ahead of time,
                  and offer the path of the file along with a function name.
@@ -258,6 +300,19 @@ class Custom(ops.PrimitiveWithInfer):
                        (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32,
                        "aot"))
 
+                 b) Ascend platform.
+                 Before using Custom operators on the Ascend platform, users must first develop custom operators
+                 based on Ascend C and compile them. The complete development and usage process can refer to the
+                 tutorial `AOT-Type Custom Operators(Ascend) <https://www.mindspore.cn/docs/en/master/model_train/custom_program/operation/op_custom_ascendc.html>`_.
+                 By passing the name of the operator through the input parameter `func`, there are two usage methods
+                 based on the implementation of the infer shape function:
+
+                 - Python infer: If the operator's infer shape is implemented in Python, that is, the infer shape
+                   function is passed through the `out_shape` parameter, specify `func="CustomName"` .
+                 - C++ infer: If the operator's infer shape is implemented through C++, then pass the path of the
+                   infer shape implementation file in `func` and separate the operator name with `:`,
+                   for example: `func="add_custom_infer.cc:AddCustom"` .
+
               2. for "julia":
 
                  Currently "julia" supports CPU(linux only) platform.
@@ -299,18 +354,7 @@ class Custom(ops.PrimitiveWithInfer):
 
         func_type (str): The implementation type of `func`, should be one of
 
-            [ ``"hybrid"`` , ``"akg"`` , ``"tbe"`` , ``"aot"`` , ``"pyfunc"`` , ``"julia"`` , ``"aicpu"`` ].
-
-            Each `func_type` only supports specific platforms(targets). Default: ``"hybrid"`` .
-            The supported platforms of `func_type`:
-
-            - ``"hybrid"``: supports ["Ascend", "GPU", "CPU"].
-            - ``"akg"``: supports ["Ascend", "GPU", "CPU"].
-            - ``"tbe"``: supports ["Ascend"].
-            - ``"aot"``: supports ["GPU", "CPU"].
-            - ``"pyfunc"``: supports ["CPU"].
-            - ``"julia"``: supports ["CPU"].
-            - ``"aicpu"``: supports ["Ascend"].
+            [ ``"hybrid"`` , ``"akg"`` , ``"aot"`` , ``"pyfunc"`` , ``"julia"`` ].
 
         bprop (function): The back propagation function of `func`. Default: ``None`` .
         reg_info (Union[str, dict, list, tuple]): Represents the registration information(reg info) of `func` with
@@ -372,68 +416,6 @@ class Custom(ops.PrimitiveWithInfer):
         >>> # the result will be a 16 * 16 tensor with all elements 2
         >>> print(output.shape)
         (16, 16)
-        >>> # Example, func_type = "tbe"
-        >>> square_with_bias_op_info = CustomRegOp() \
-        ...     .fusion_type("OPAQUE") \
-        ...     .attr("bias", "required", "float") \
-        ...     .input(0, "x") \
-        ...     .output(0, "y") \
-        ...     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
-        ...     .dtype_format(DataType.F16_Default, DataType.F16_Default) \
-        ...     .target("Ascend") \
-        ...     .get_op_info()
-        >>>
-        >>> @custom_info_register(square_with_bias_op_info)
-        ... def square_with_bias(input_x, output_y, bias=0.0, kernel_name="square_with_bias"):
-        ...     import te.lang.cce
-        ...     from te import tvm
-        ...     from topi.cce import util
-        ...
-        ...     shape = input_x.get("shape")
-        ...     dtype = input_x.get("dtype").lower()
-        ...
-        ...     shape = util.shape_refine(shape)
-        ...     data = tvm.placeholder(shape, name="data", dtype=dtype)
-        ...
-        ...     with tvm.target.cce():
-        ...         res0 = te.lang.cce.vmul(data, data)
-        ...         res = te.lang.cce.vadds(res0, bias)
-        ...         sch = te.lang.cce.auto_schedule(res)
-        ...
-        ...     config = {"print_ir": False,
-        ...               "name": kernel_name,
-        ...               "tensor_list": [data, res]}
-        ...
-        ...     te.lang.cce.cce_build_code(sch, config)
-        >>>
-        >>> def test_tbe():
-        ...     square_with_bias = ops.Custom(square_with_bias, out_shape=lambda x, _: x, \
-        ...                                   out_dtype=lambda x, _: x, func_type="tbe")
-        ...     res = self.square_with_bias(input_x, 1.0)
-        ...     return res
-        >>>
-        >>> # Example, func_type = "aicpu"
-        >>> resize_bilinear_op_info = CustomRegOp("ResizeBilinear") \
-        ...     .fusion_type("OPAQUE") \
-        ...     .input(0, "input", "required") \
-        ...     .output(1, "output", "required") \
-        ...     .attr("align_corners", "required", "bool") \
-        ...     .attr("cust_aicpu", "optional", "str", "aicpu_kernels") \
-        ...     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
-        ...     .dtype_format(DataType.F16_Default, DataType.F32_Default) \
-        ...     .target("Ascend") \
-        ...     .get_op_info()
-        >>>
-        >>> @custom_info_register(resize_bilinear_op_info)
-        ... def resize_bilinear_aicpu():
-        ...     return
-        >>>
-        >>> def test_aicpu(x):
-        ...     resize_bilinear_op = ops.Custom(resize_bilinear_aicpu, out_shape=[1, 1, 9, 9], \
-        ...                                     out_dtype=mstype.float32, func_type="aicpu")
-        ...     res = resize_bilinear_op(x, True, "aicpu_kernels")
-        ...     return res
-        >>>
         >>> # Example, func_type = "aot"
         >>> def test_aot(x, y, out_shapes, out_types):
         ...     program = ops.Custom("./reorganize.so:CustomReorganize", out_shapes, out_types, "aot")
@@ -484,6 +466,7 @@ class Custom(ops.PrimitiveWithInfer):
         self.func_source_str = ""
         self._func_compile_attrs = {}
         self._is_ms_kernel = False
+        self.out_shape = out_shape
 
         self._check_platform()
         self._check_func()
@@ -499,7 +482,6 @@ class Custom(ops.PrimitiveWithInfer):
             add_pyfunc(func_id, self.func)
             self.add_prim_attr("fn_id", func_id)
 
-        self.out_shape = out_shape
         if self.out_shape is None and self.func_type == "aot":
             self.add_prim_attr("cpp_infer_shape", True)
         self.out_dtype = out_dtype
@@ -524,6 +506,12 @@ class Custom(ops.PrimitiveWithInfer):
 
         self.add_prim_attr("func_type", self.func_type)
         self._update_attr()
+        self.enable_pyboost = False
+        self.custom_pyboost = _CustomExt(self.func, self.out_shape, self.out_dtype, self.bprop)
+        if context.get_context("device_target") == "Ascend" and self.func_type == "aot":
+            self.enable_pyboost = True
+            for key, value in super().get_attr_dict().items():
+                self.custom_pyboost.add_prim_attr(key, value)
 
     def __infer__(self, *args):
         if callable(self.out_shape):
@@ -541,9 +529,9 @@ class Custom(ops.PrimitiveWithInfer):
         # deal with the case of ms script
         # enable auto infer function if any infer information is missing
         if self._is_ms_kernel and (infer_dtype is None or infer_shape is None):
-            logger.warning("{}, 'out_shape' or 'out_dtype' is None, infer the output shape and output dtype "
-                           "automatically. There might be some Python RuntimeWarning but it wouldn't influence the "
-                           "result.".format(self.log_prefix))
+            logger.info("{}, 'out_shape' or 'out_dtype' is None, infer the output shape and output dtype "
+                        "automatically. There might be some Python RuntimeWarning but it wouldn't influence the "
+                        "result.".format(self.log_prefix))
 
             auto_infer_result = self._auto_infer(*args)
 
@@ -585,6 +573,7 @@ class Custom(ops.PrimitiveWithInfer):
         return out
 
     def get_bprop(self):
+        """return back propagation function"""
         return self.bprop
 
     def _set_akg_kernel_type(self):
@@ -613,7 +602,37 @@ class Custom(ops.PrimitiveWithInfer):
                 raise Exception("{}, function {} is not found in source file {}!"
                                 .format(self.log_prefix, func, source_file))
 
+    def _check_aot_func(self):
+        """Check the source code and bin lib for aot type custom op"""
+        if not isinstance(self.func, str):
+            raise TypeError("{}, 'func' must be of type str, but got {}".format(
+                self.log_prefix, type(self.func)))
+        file_name_list = self.func.split(":")
+        if len(file_name_list) != 2:
+            if callable(self.out_shape):
+                return
+            raise TypeError(
+                "{}, 'func' should be like 'file_name:func_name', but got {}".format(
+                    self.log_prefix, self.func))
+        file_path = os.path.realpath(file_name_list[0])
+        if os.environ.get('MS_CUSTOM_AOT_WHITE_LIST') is None:
+            if Custom.custom_aot_warning:
+                logger.info("{}, no white list is set and it might cause problems. "
+                            "Set the legal path of the file in MS_CUSTOM_AOT_WHITE_LIST"
+                            .format(self.log_prefix))
+                Custom.custom_aot_warning = False
+        else:
+            legal_path = os.path.realpath(os.environ.get('MS_CUSTOM_AOT_WHITE_LIST'))
+            if legal_path not in file_path:
+                raise TypeError(
+                    "{}, the legal path for the file is {}, but the file is {}".format(
+                        self.log_prefix, legal_path, file_path))
+        if file_path.endswith(("cu", "cpp", "cc")):
+            file_path = _compile_aot(file_path)
+            self.func = file_path + ":" + file_name_list[1]
+
     def _check_platform(self):
+        """check the platform"""
         if platform.system() != 'Linux':
             raise Exception("Custom op only supported on Linux platform currently.")
 
@@ -623,30 +642,7 @@ class Custom(ops.PrimitiveWithInfer):
             raise ValueError("{}, 'func_type' must be one of {}, but got {}"
                              .format(self.log_prefix, self.supported_func_type, self.func_type))
         if self.func_type == "aot":
-            if not isinstance(self.func, str):
-                raise TypeError("{}, 'func' must be of type str, but got {}".format(
-                    self.log_prefix, type(self.func)))
-            file_name_list = self.func.split(":")
-            if len(file_name_list) != 2:
-                raise TypeError(
-                    "{}, 'func' should be like 'file_name:func_name', but got {}".format(
-                        self.log_prefix, self.func))
-            file_path = os.path.abspath(file_name_list[0])
-            if os.environ.get('MS_CUSTOM_AOT_WHITE_LIST') is None:
-                if Custom.custom_aot_warning:
-                    logger.warning("{}, no white list is set and it might cause problems. "
-                                   "Set the legal path of the file in MS_CUSTOM_AOT_WHITE_LIST"
-                                   .format(self.log_prefix))
-                    Custom.custom_aot_warning = False
-            else:
-                legal_path = os.path.abspath(os.environ.get('MS_CUSTOM_AOT_WHITE_LIST'))
-                if legal_path not in file_path:
-                    raise TypeError(
-                        "{}, the legal path for the file is {}, but the file is {}".format(
-                            self.log_prefix, legal_path, file_path))
-            if file_path.endswith(("cu", "cpp", "cc")):
-                file_path = _compile_aot(file_path)
-                self.func = file_path + ":" + file_name_list[1]
+            self._check_aot_func()
 
         elif self.func_type == "julia":
             self._check_julia_func()
@@ -1122,3 +1118,12 @@ class Custom(ops.PrimitiveWithInfer):
         infer_value = Tensor(fake_output) if enable_infer_value else None
 
         return infer_shape, infer_dtype, infer_value
+
+    def __call__(self, *args):
+        if self.enable_pyboost:
+            return _convert_stub(pyboost_custom_ext(self.custom_pyboost, [args]))
+        should_elim, output = self.check_elim(*args)
+        if should_elim:
+            return output
+        # pylint: disable=protected-access
+        return ops.primitive._run_op(self, self.name, args)

mindspore/ops/operations/debug_ops.py

@@ -13,19 +13,23 @@
 # limitations under the License.
 # ============================================================================
 """debug_ops"""
-from types import FunctionType, MethodType
+import os
+import stat
 
+import numpy as np
 from mindspore import log as logger
-from mindspore._c_expression import security
+from mindspore._c_expression import security, HookType
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 from mindspore.ops.primitive import prim_attr_register, Primitive, PrimitiveWithInfer
+from mindspore._checkparam import check_hook_fn
 
 
 SUMMARY_TENSOR_CACHE = []
+TENSORDUMP_ID = 0
 
 
 def _cache_summary_data(op_name, define_name, tensor):
@@ -55,10 +59,13 @@ SUMMARY_RETURN_VALUE = {'dtype': mstype.int32, 'shape': [1], 'value': None}
 
 class ScalarSummary(Primitive):
     """
-    This operator will put a scalar to a summary file with protocol buffer format. It must be used with SummaryRecord
-    or SummaryCollector, which specify the directory of the summary file. The summary file can
+    This operator will put a scalar to a summary file with protocol buffer format.
+    It must be used with :class:`mindspore.SummaryRecord` or :class:`mindspore.SummaryCollector`,
+    which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.2/index.html>`_ for details.
+    mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -75,7 +82,7 @@ class ScalarSummary(Primitive):
     Examples:
         >>> import mindspore
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore import Tensor, set_context
         >>>
         >>>
@@ -105,17 +112,20 @@ class ScalarSummary(Primitive):
 
         self.add_prim_attr("side_effect_io", True)
         self.add_prim_attr("channel_name", "ms_scalar_summary")
+        self.add_prim_attr("dyn_input_sizes", [-1, 1])
 
     def __call__(self, *args):
         _cache_summary_data(self.name, args[0], args[1])
 
 
-class ImageSummary(PrimitiveWithInfer):
+class ImageSummary(Primitive):
     """
     This operator will put an image tensor to a summary file with protocol buffer format. It must be used with
     SummaryRecord or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.2/index.html>`_ for details.
+    mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -131,7 +141,7 @@ class ImageSummary(PrimitiveWithInfer):
     Examples:
 
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>>
         >>>
         >>> class Net(nn.Cell):
@@ -155,18 +165,7 @@ class ImageSummary(PrimitiveWithInfer):
 
         self.add_prim_attr("side_effect_io", True)
         self.add_prim_attr("channel_name", "ms_image_summary")
-
-    def __infer__(self, name, value):
-        _check_summary_param(name, value, self.__class__.__name__)
-
-        # The shape dim of image should be 4.
-        v_shape = value['shape']
-        image_dim = 4
-        if len(v_shape) != image_dim:
-            raise ValueError(f"For '{self.name}', the dimension of 'value' must be {image_dim},"
-                             f" but got {len(v_shape)}.")
-
-        return SUMMARY_RETURN_VALUE
+        self.add_prim_attr("dyn_input_sizes", [-1, 1])
 
     def __call__(self, *args):
         _cache_summary_data(self.name, args[0], args[1])
@@ -177,7 +176,9 @@ class TensorSummary(Primitive):
     This operator will put a tensor to a summary file with protocol buffer format. It must be used with SummaryRecord
     or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.2/index.html>`_ for details.
+    mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable.
@@ -194,7 +195,7 @@ class TensorSummary(Primitive):
     Examples:
         >>> import mindspore
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore import Tensor, set_context
         >>>
         >>>
@@ -229,12 +230,123 @@ class TensorSummary(Primitive):
         _cache_summary_data(self.name, args[0], args[1])
 
 
-class HistogramSummary(PrimitiveWithInfer):
+class TensorDump(Primitive):
+    """
+    Save the Tensor as an npy file in numpy format.
+
+    .. warning::
+        - If a large amount of data is stored within a short period, it may lead to memory overflow on the device side.
+          Consider slicing the data to reduce the data scale.
+        - Since data saving is processed asynchronously, when the amount of data is too large or the main process exits
+          too quickly, data loss may occur. You need to actively control the destruction time of the main process,
+          such as using sleep.
+
+    Args:
+        input_output (str, optional): Used to control Tensordump behavior.
+            Available value is one of ['in', 'out', 'all']. Default value is ``out``.
+
+            In case of OpA --> RedistributionOps --> OpB,
+            The dump data of OpA's output is not equal to OpB's input (Due to the redistribution operators).
+            So the parameter input_output is to handle this situation.
+
+            Assuming OpA's output is used as both Tensordump's input parameter and OpB's input parameter.
+            Different requirements of saving dump data can be achieved by configuring parameter input_output:
+
+            - If the input_output is 'out', the dump data contains only OpA's output slice.
+            - If the input_output is 'all', the dump data contains both OpA's output slice and OpB's input slice.
+            - If the input_output is 'in', the dump data contains only OpB's input slice.
+
+            For input_output is 'all' or 'in', the input slice npy file format is:
+            id_fileName_cNodeID_dumpMode_rankID.npy.
+
+            For input_output is 'out' or 'all' the output slice npy file format is:
+            id_fileName.npy.
+
+            - id: An auto increment ID.
+            - fileName: Value of the parameter file
+              (if parameter file_name is a user-specified path, the value of fileName is the last level of the path).
+            - cNodeID: The node ID of the Tensordump node in the step_parallel_end.ir file.
+            - dumpMode: Value of the parameter input_output.
+            - rankID: Logical device id.
+
+    Inputs:
+        - **file** (str) - The path of the file to be saved.
+        - **input_x** (Tensor) - Input Tensor of any dimension.
+
+    Raises:
+        TypeError: If `file` is not a str.
+        TypeError: If `input_x` is not a Tensor.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import time
+        >>> from mindspore import nn, Tensor, ops
+        >>> ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")
+        >>> class Net(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net, self).__init__()
+        ...         self.dump = ops.TensorDump()
+        ...
+        ...     def construct(self, x):
+        ...         x += 1.
+        ...         self.dump('add', x)
+        ...         x /= 2.
+        ...         self.dump('div', x)
+        ...         x *= 5.
+        ...         self.dump('mul', x)
+        ...         return x
+        ...
+        >>> x = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]).astype(np.float32)
+        >>> input_x = Tensor(x)
+        >>> net = Net()
+        >>> out = net(input_x)
+        >>> time.sleep(0.5)
+        >>> add = np.load('0_add.npy')
+        >>> print(add)
+        [[2. 3. 4. 5.]
+         [6. 7. 8. 9.]]
+    """
+    @prim_attr_register
+    def __init__(self, input_output='out'):
+        """Initialize TensorDump."""
+        if security.enable_security():
+            raise ValueError('The TensorDump is not supported, please without `-s on` and recompile source.')
+        self.add_prim_attr("side_effect_io", True)
+        self.add_prim_attr("channel_name", "ms_tensor_dump")
+
+    def __call__(self, file, input_x):
+        validator.check_value_type('file', file, [str], self.__class__.__name__)
+        if not file:
+            raise ValueError("For 'TensorDump', the input argument[file] cannot be an empty string.")
+        validator.check_value_type('input_x', input_x, [Tensor], self.__class__.__name__)
+        global TENSORDUMP_ID
+        npy_suffix = ".npy"
+        directory, filename = os.path.split(file)
+        if directory and not os.path.exists(directory):
+            os.makedirs(directory, mode=0o700, exist_ok=True)
+        new_filename = f"{TENSORDUMP_ID}_{filename}"
+        new_file = os.path.join(directory, new_filename)
+        if not new_file.endswith(npy_suffix):
+            new_file += npy_suffix
+        if os.path.exists(new_file):
+            os.chmod(new_file, stat.S_IWUSR)
+        np.save(new_file, input_x.asnumpy())
+        os.chmod(new_file, stat.S_IRUSR)
+        TENSORDUMP_ID += 1
+
+
+class HistogramSummary(Primitive):
     """
     This operator will calculate the histogram of a tensor and put it to a summary file with protocol buffer format.
     It must be used with SummaryRecord or SummaryCollector, which specify the directory of the summary file.
     The summary file can be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.2/index.html>`_ for details.
+    mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable.
@@ -250,7 +362,7 @@ class HistogramSummary(PrimitiveWithInfer):
     Examples:
         >>> import mindspore
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore import Tensor, set_context
         >>>
         >>>
@@ -280,17 +392,7 @@ class HistogramSummary(PrimitiveWithInfer):
 
         self.add_prim_attr("side_effect_io", True)
         self.add_prim_attr("channel_name", "ms_histogram_summary")
-
-    def __infer__(self, name, value):
-        _check_summary_param(name, value, self.__class__.__name__)
-
-        v_shape = value['shape']
-        # In the summary, the histogram value should be a tensor whose shape is not [].
-        if not v_shape:
-            raise ValueError(f"For '{self.name}', the type of 'value' must be tensor, "
-                             f"its shape should not be [], but got {v_shape}.")
-
-        return SUMMARY_RETURN_VALUE
+        self.add_prim_attr("dyn_input_sizes", [-1, 1])
 
     def __call__(self, *args):
         _cache_summary_data(self.name, args[0], args[1])
@@ -430,16 +532,15 @@ class HookBackward(PrimitiveWithInfer):
     def __init__(self, hook_fn, cell_id=""):
         """Initialize HookBackward."""
         super(HookBackward, self).__init__(self.__class__.__name__)
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"For '{self.name}', the type of 'hook_fn' must be python function, "
-                            f"but got {type(hook_fn)}.")
+        if not check_hook_fn("HookBackward", hook_fn):
+            return
         if cell_id != "":
             logger.warning(f"The args 'cell_id' of HookBackward will be removed in a future version. If the value of "
                            f"'cell_id' is set, the hook function will not work.")
         self.add_prim_attr("cell_id", cell_id)
         self.init_attrs["cell_id"] = cell_id
         self.cell_id = cell_id
-        self.add_backward_hook_fn(hook_fn)
+        self.set_hook_fn(hook_fn, HookType.HookBackward)
 
     def infer_shape(self, *inputs_shape):
         if len(inputs_shape) == 1: