mindspore 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/ops/operations/custom_ops.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2024 Huawei Technologies Co., Ltd
+# Copyright 2021-2025 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.
@@ -21,6 +21,7 @@ import re
 import ast
 import hashlib
 import stat
+import copy
 import inspect
 import importlib
 import platform
@@ -29,7 +30,6 @@ 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
@@ -40,6 +40,7 @@ from mindspore.communication.management import get_rank, GlobalComm
 from ._ms_kernel import determine_variable_usage
 from ._custom_grad import autodiff_bprop
 from ._pyfunc_registry import add_pyfunc
+from ._custom_ops_utils import ExtensionBuilder
 
 if platform.system() != "Windows":
     import fcntl
@@ -109,12 +110,19 @@ def _compile_aot(file):
     func_path = cache_path + file_name + ".so"
     include_file = "{} -I{}".format(include_file, file[:file.rindex('/')])
 
+    if context.get_context("device_target") == "Ascend":
+        ascend_cann_path = os.getenv("ASCEND_OPP_PATH").split('opp')[0]
+        ascend_include = os.path.join(ascend_cann_path, "include")
+        include_file = "{} -I{}".format(include_file, ascend_include)
+
+    include_file = include_file.split(" ")
     if func_path not in Custom.compiled_bin:
         Custom.compiled_bin.append(func_path)
 
         if file.endswith("cpp") or file.endswith("cc"):
             cmd = ["g++", "-std=c++17", "--shared", "-fPIC", "-D_GLIBCXX_USE_CXX11_ABI=0"]
-            cmd += [include_file, "-o", func_path, file]
+            cmd += include_file
+            cmd += ["-o", func_path, file]
         elif file.endswith("cu"):
             cmd = ["nvcc"]
             cmd += ["--shared", "-Xcompiler", "-fPIC", "-O3", "-gencode", "arch=compute_70, code=sm_70"]
@@ -141,12 +149,13 @@ def _compile_aot(file):
                 logger.warning("The current version of nvcc, V{}.{}.{},  might have unfixed issues with std string, "
                                "which will lead to errors in aot custom op with attrs."
                                "The version higher than V10.1.168 is recommended".format(v_major, v_mid, v_minor))
-            cmd += [include_file, "-o", func_path, file]
+            cmd += include_file
+            cmd += ["-o", func_path, file]
         else:
             raise ValueError("The source file must be a cc/cpp/cu file, but get: {}".format(file))
 
         proc = subprocess.Popen(
-            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=False)
 
         (out, _) = proc.communicate(timeout=30)
 
@@ -183,10 +192,16 @@ class _CustomExt(ops.PrimitiveWithInfer):
 
         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.")
+            logger.debug("'out_shape' is None. Add a placeholder instead. "
+                         "A CPP version of infer shape function is required "
+                         "in this case.")
             infer_shape = (1,)
+        if infer_dtype is None:
+            logger.debug("'out_dtype' is None. Add a placeholder instead. "
+                         "A CPP version of infer type function is required "
+                         "in this case.")
+            infer_dtype = ms.float16
+
         # 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)))
@@ -215,7 +230,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/docs/en/master/model_train/custom_program/op_custom.html>`_ .
+    <https://www.mindspore.cn/tutorials/en/master/custom_program/op_custom.html>`_ .
 
     .. warning::
         - This is an experimental API that is subject to change.
@@ -297,14 +312,17 @@ class Custom(ops.PrimitiveWithInfer):
                  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>`_.
+                 tutorial `AOT-Type Custom Operators(Ascend)
+                 <https://www.mindspore.cn/tutorials/en/master/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 `:`,
+                 based on the implementation of the infer function:
+
+                 - Python infer: If the operator's infer function is implemented in Python, that is, the infer shape
+                   function is passed through the `out_shape` parameter, and the infer type is passed throuht the
+                   `out_dtype`, then the `func` should be specified as the operator name, for example,
+                   `func="CustomName"`.
+                 - C++ infer: If the operator's infer function is implemented through C++, then pass the path of the
+                   infer function implementation file in `func` and separate the operator name with `:`,
                    for example: `func="add_custom_infer.cc:AddCustom"` .
 
               2. for "julia":
@@ -425,6 +443,8 @@ class Custom(ops.PrimitiveWithInfer):
         self._func_compile_attrs = {}
         self._is_ms_kernel = False
         self.out_shape = out_shape
+        self.out_dtype = out_dtype
+        self.is_ascend_c = context.get_context("device_target") == "Ascend" and self.func_type == "aot"
 
         self._check_platform()
         self._check_func()
@@ -440,15 +460,17 @@ class Custom(ops.PrimitiveWithInfer):
             add_pyfunc(func_id, self.func)
             self.add_prim_attr("fn_id", func_id)
 
-        if self.out_shape is None and self.func_type == "aot":
-            self.add_prim_attr("cpp_infer_shape", True)
-        self.out_dtype = out_dtype
+        self.set_infer_flag()
+
+        self.multi_output = (reg_info is not None and (len(reg_info.get("outputs", [])) > 1))
+        self.add_prim_attr("multi_output", self.multi_output)
+
         self.bprop = bprop
         self.fake_output = False
         self.single_scalar_output = False
-        if not self.out_dtype:
+        if not self.out_dtype and not self.func_type == "pyfunc":
             self.fake_output = True
-        elif not self.out_shape:
+        elif not self.out_shape and self.func_type == "pyfunc":
             self.single_scalar_output = True
         self.add_prim_attr("fake_output", self.fake_output)
         self.add_prim_attr("single_scalar_output", self.single_scalar_output)
@@ -464,13 +486,28 @@ 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
+
+        if self.is_ascend_c:
+            self.set_inputs_type(reg_info)
+            self.custom_pyboost = _CustomExt(self.func, self.out_shape, self.out_dtype, self.bprop)
             for key, value in super().get_attr_dict().items():
                 self.custom_pyboost.add_prim_attr(key, value)
 
+    def set_infer_flag(self):
+        """set cpp infer attr"""
+        if self.out_shape is None and self.func_type == "aot":
+            self.add_prim_attr("cpp_infer_shape", True)
+        if self.out_dtype is None and self.func_type == "aot":
+            self.add_prim_attr("cpp_infer_type", True)
+
+    def set_inputs_type(self, reg_info):
+        """set custom_inputs_type attr"""
+        if not self.is_ascend_c or not reg_info.get('attr'):
+            return
+        inputs_type = ["tensor"] * len(reg_info.get("inputs", [])) + \
+                      [attr.get("type") for attr in reg_info.get("attr", [])]
+        self.add_prim_attr("custom_inputs_type", inputs_type)
+
     def __infer__(self, *args):
         if callable(self.out_shape):
             infer_shape = self.out_shape(*(x["shape"] for x in args))
@@ -510,10 +547,15 @@ class Custom(ops.PrimitiveWithInfer):
                 infer_dtype = mstype.int32
         if self.func_type == "aot":
             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.".format(self.log_prefix))
+                logger.debug("{}, 'out_shape' is None. Add a placeholder instead. "
+                             "A CPP version of infer shape function is required "
+                             "in this case.".format(self.log_prefix))
                 infer_shape = (1,)
+            if infer_dtype is None:
+                logger.debug("{}, 'out_dtype' is None. Add a placeholder instead. "
+                             "A CPP version of infer type function is required "
+                             "in this case.".format(self.log_prefix))
+                infer_dtype = ms.float16
         # 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 {}"
@@ -757,6 +799,26 @@ class Custom(ops.PrimitiveWithInfer):
                     if isinstance(item, dict) and item.get("value") is not None:
                         self.add_prim_attr(item[KEY_NAME], item["value"])
 
+    def _convert_attr_to_input(self, ori_reg_info):
+        """convert attr to input"""
+        if not self.is_ascend_c or not ori_reg_info.get("attr"):
+            return ori_reg_info
+
+        reg_info = copy.deepcopy(ori_reg_info)
+        start_index = len(reg_info.get("inputs", []))
+        for i, attr_item in enumerate(reg_info.get("attr", [])):
+            new_input = {
+                'index': start_index + i,
+                'name': attr_item['name'],
+                'paramType': attr_item['paramType']}
+            reg_info['inputs'].append(new_input)
+            for dtype_format_item in reg_info.get("dtype_format", []):
+                new_dtype_format_item = list(dtype_format_item)
+                new_dtype_format_item.insert(start_index + i, DataType.None_None)
+                reg_info['dtype_format'][reg_info['dtype_format'].index(dtype_format_item)] = new_dtype_format_item
+        reg_info['attr'] = []
+        return reg_info
+
     def _register_info(self, info):
         """Register reg_info."""
         reg_info = info
@@ -787,14 +849,15 @@ class Custom(ops.PrimitiveWithInfer):
                 continue
             # Register
             reg_info = self._reformat_reg_info(reg_info, target)
-            reg_info_str = json.dumps(reg_info)
+            new_reg_info = self._convert_attr_to_input(reg_info)
+            reg_info_str = json.dumps(new_reg_info)
             op_lib = Oplib()
             if not op_lib.reg_op(reg_info_str, self.imply_path):
                 raise ValueError("{}, the registration information is registered failed. Use 'CustomRegOp' to "
                                  "generate the registration information, then pass it to 'reg_info' or use "
                                  "'custom_info_register' to bind it to 'func' if 'func' is a function."
                                  .format(self.log_prefix))
-            self._save_attr(reg_info)
+            self._save_attr(new_reg_info)
             self._save_register_status(target)
 
     def _get_expanded_list(self, data):
@@ -1078,10 +1141,202 @@ class Custom(ops.PrimitiveWithInfer):
         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]))
+        if self.is_ascend_c:
+            res = pyboost_custom_ext(self.custom_pyboost, [args])
+            return res if self.multi_output else res[0]
         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)
+
+
+class CustomOpBuilder:
+    r"""
+    CustomOpBuilder is used to initialize and configure custom operators for MindSpore.
+    Users can define and load custom operator modules through this class and apply them to the network.
+
+    In most cases, users only need to provide the source files and additional compilation options in the constructor
+    and call the `load` method to complete the compilation and loading of the operator.
+    If users have specific customization requirements, they can inherit this class and override certain methods.
+    It is important to note that if methods are overridden, some parameters passed to the constructor may be ignored.
+
+    .. warning::
+        This is an experimental API that is subject to change.
+
+    Args:
+        name (str): The unique name of the custom operator module, used to identify the operator.
+        sources (Union[str, list[str]]): The source file(s) of the custom operator. It can be a single file path or
+                                    a list of file paths.
+        backend (str, optional): The target backend for the operator, such as "CPU" or "Ascend". Default: ``None``.
+        include_paths (list[str], optional): Additionally included paths needed during compilation. Default: ``None``.
+        cflags (str, optional): Extra C++ compiler flags to be used during compilation. Default: ``None``.
+        ldflags (str, optional): Extra linker flags to be used during linking. Default: ``None``.
+        kwargs (dict, optional): Additional keyword arguments for future extensions or specific custom requirements.
+
+    .. note::
+        - If the `backend` argument is provided, additional default flags will be automatically added to
+          the compilation and linking steps to support the operator's target backend. The default options
+          can be referenced in the implementation of the `get_cflags` and `get_ldflags` methods in the `CustomOpBuilder
+          <https://gitee.com/mindspore/mindspore/blob/master/mindspore/python/mindspore/ops/operations/custom_ops.py>`_.
+        - The `sources` argument must point to valid source files for the custom operator.
+
+    Supported Platforms:
+        ``Ascend`` ``CPU``
+
+    Examples:
+        >>> from mindspore import ops
+        >>> builder = ops.CustomOpBuilder(
+        ...     name="custom_op_cpu",
+        ...     sources="custom_ops_impl/pybind_op_cpu.cpp",
+        ...     backend="CPU"
+        ... )
+        >>> my_ops = builder.load()
+    """
+    _mindspore_path = None
+    _loaded_ops = {}
+    _ms_code_base = None
+
+    def __init__(self, name, sources, backend=None, include_paths=None, cflags=None, ldflags=None, **kwargs):
+        self.name = name
+        self.source = sources
+        self.backend = backend
+        self.include_paths = include_paths
+        self.cflags = cflags
+        self.ldflags = ldflags
+        self.build_dir = kwargs.get("build_dir")
+        if CustomOpBuilder._mindspore_path is None:
+            CustomOpBuilder._mindspore_path = os.path.dirname(os.path.abspath(ms.__file__))
+            CustomOpBuilder._ms_code_base = os.path.join(CustomOpBuilder._mindspore_path, "include")
+        if self.backend == "Ascend":
+            self.ascend_cann_path = os.getenv("ASCEND_OPP_PATH").split('opp')[0]
+
+    def get_sources(self):
+        """
+        Get the source files for the custom operator.
+
+        Returns:
+            str or list[str], The source file(s) for the operator.
+        """
+        return self.source
+
+    def get_include_paths(self):
+        """
+        Get the include paths required for compiling the custom operator.
+
+        Returns:
+            list[str], A list of include paths.
+        """
+        include_list = self.include_paths if self.include_paths is not None else []
+        include_list.append(CustomOpBuilder._mindspore_path)
+        include_list.append(os.path.join(CustomOpBuilder._mindspore_path, "include"))
+        include_list.append(os.path.join(CustomOpBuilder._mindspore_path, "include/third_party"))
+        include_list.append(os.path.join(CustomOpBuilder._mindspore_path, "include/third_party/robin_hood_hashing"))
+        include_list.append(os.path.join(CustomOpBuilder._mindspore_path, "include/third_party/securec/include"))
+
+        if self.backend == "Ascend":
+            include_list.append(os.path.join(self.ascend_cann_path, "include"))
+        include_list += self._get_ms_inner_includes()
+        return include_list
+
+    def _get_ms_inner_includes(self):
+        """include paths for inner module interface."""
+        ms_inner_code_base = os.path.join(CustomOpBuilder._mindspore_path, "include", "mindspore")
+        include_list = []
+        include_list.append(ms_inner_code_base + "/core/include")
+        include_list.append(ms_inner_code_base + "/core/mindrt/include")
+        include_list.append(ms_inner_code_base + "/core/mindrt")
+        include_list.append(ms_inner_code_base + "/ops")
+        include_list.append(ms_inner_code_base + "/ops/kernel/include")
+        include_list.append(ms_inner_code_base + "/ccsrc")
+        include_list.append(ms_inner_code_base + "/ccsrc/include")
+        include_list.append(ms_inner_code_base + "/ccsrc/minddata/mindrecord/include")
+        return include_list
+
+    def get_cflags(self):
+        """
+        Get the C++ compiler flags for building the custom operator.
+
+        Returns:
+            list[str], A list of C++ compiler flags.
+        """
+        flags = ['-fstack-protector-all', '-fPIC', '-pie']
+        flags += ['-DENABLE_FAST_HASH_TABLE=1']
+        if self.backend == "Ascend":
+            flags.append('-DCUSTOM_ASCEND_OP')
+        if self.cflags is not None:
+            flags.append(self.cflags)
+        return flags
+
+    def get_ldflags(self):
+        """
+        Get the linker flags for building the custom operator.
+
+        Returns:
+            list[str], A list of linker flags.
+        """
+        flags = ['-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath', '-s']
+        flags += [
+            '-L' + os.path.abspath(os.path.join(CustomOpBuilder._mindspore_path, 'lib')),
+            '-lmindspore_core',
+            '-lmindspore_ms_backend',
+            '-lmindspore_pynative'
+        ]
+        if self.backend == "Ascend":
+            flags.append('-L' + os.path.abspath(os.path.join(CustomOpBuilder._mindspore_path, 'lib/plugin')))
+            flags.append('-L' + os.path.abspath(os.path.join(self.ascend_cann_path, "lib64")))
+            flags.append('-lascendcl')
+            flags.append('-l:libmindspore_ascend.so.2')
+        if self.ldflags is not None:
+            flags.append(self.ldflags)
+        return flags
+
+    def build(self):
+        """
+        Build the custom operator module.
+
+        This method generates a dynamic library file for the custom operator based on the provided source files,
+        include paths, compilation flags, and link flags.
+
+        Returns:
+            str, The path to the compiled module.
+        """
+        return ExtensionBuilder(self._get_build_directory()).build(
+            module_name=self.name,
+            sources=self.get_sources(),
+            extra_include_paths=self.get_include_paths(),
+            extra_cflags=self.get_cflags(),
+            extra_ldflags=self.get_ldflags())
+
+    def load(self):
+        """
+        Build and load the custom operator module.
+
+        Returns:
+            Module, The loaded custom operator module.
+        """
+        if self.name in CustomOpBuilder._loaded_ops:
+            return CustomOpBuilder._loaded_ops[self.name]
+        module_path = self.build()
+        mod = self._import_module(module_path)
+        CustomOpBuilder._loaded_ops[self.name] = mod
+        return mod
+
+    def _import_module(self, module_path):
+        """Import module from library."""
+        spec = importlib.util.spec_from_file_location(self.name, module_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+
+    def _get_build_directory(self):
+        """Get build directory."""
+        if self.build_dir is None:
+            build_root = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', "./kernel_meta"))
+            self.build_dir = os.path.join(build_root, self.name)
+        else:
+            self.build_dir = os.path.realpath(self.build_dir)
+        logger.info(f'Build {self.name} in directory {self.build_dir}')
+        if not os.path.exists(self.build_dir):
+            os.makedirs(self.build_dir, exist_ok=True)
+        return self.build_dir

mindspore/ops/operations/debug_ops.py

@@ -19,12 +19,13 @@ from pathlib import Path
 import numpy as np
 from mindspore import log as logger
 from mindspore._c_expression import security, HookType
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._c_expression import _tensordump_process_file
 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.common.jit_context import jit_context
 from mindspore.ops.primitive import prim_attr_register, Primitive, PrimitiveWithInfer
 from mindspore._checkparam import check_hook_fn
 from mindspore.ops import operations as P
@@ -62,9 +63,7 @@ class ScalarSummary(Primitive):
     """
     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/master/index.html>`_ for details.
+    which specify the directory of the summary file.
     In Ascend platform with graph mode, the environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
     can be set to solve operator execution failure when calling this operator intensively.
 
@@ -122,11 +121,9 @@ class ScalarSummary(Primitive):
 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/master/index.html>`_ for details.
+    SummaryRecord or SummaryCollector, which specify the directory of the summary file.
     In Ascend platform with graph mode, the environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
-    can be set to solve operator execution failure when calling this operator intensively.
+    can be set to solve execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -175,9 +172,7 @@ class ImageSummary(Primitive):
 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/master/index.html>`_ for details.
+    or SummaryCollector, which specify the directory of the summary file.
     In Ascend platform with graph mode, the environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
     can be set to solve operator execution failure when calling this operator intensively.
 
@@ -236,7 +231,7 @@ class TensorDump(Primitive):
     Save the Tensor as an npy file in numpy format.
 
     .. warning::
-        - The parameter input_output will no longer support the value 'all'.
+        The parameter input_output will no longer support the value 'all'.
 
     .. note::
         In Ascend platform with graph mode, the environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
@@ -344,8 +339,6 @@ 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/master/index.html>`_ for details.
     In Ascend platform with graph mode, the environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
     can be set to solve operator execution failure when calling this operator intensively.
 
@@ -403,6 +396,10 @@ class InsertGradientOf(Primitive):
     """
     Attaches callback to the graph node that will be invoked on the node's gradient.
 
+    .. warning::
+        In the callback, exercise caution when using side-effect operators,
+        such as the TensorDump operator, as current support is incomplete.
+
     Args:
         f (Function): MindSpore's Function. Callback function.
 
@@ -466,6 +463,107 @@ class InsertGradientOf(Primitive):
         self.f = f
 
 
+class Morph(PrimitiveWithInfer):
+    """
+    The `Morph` Primitive is used to encapsulate a user-defined function `fn`, allowing it to be used as a custom
+    Primitive.
+    The primary application scenario of the `Morph` Primitive is in the auto-parallel case after `GRAPH_MODE` mode,
+    where collective communication operators are used within the user-defined `fn` to implement custom parallel
+    computation logic, especially in scenarios where `fn` involves dynamic shapes.
+    When the `Morph` Primitive is applied to inputs, it is actually the encapsulated user-defined function `fn` that is
+    applied to the inputs.
+    The main difference between the `Morph` Primitive and :func:`mindspore.ops.Custom` is that the former is expanded
+    and replaced by the user-defined `fn` before automatic differentiation, so there is no need to implement a backward
+    function.
+
+    .. note::
+        - This primitive is only supported in GRAPH_MODE.
+        - `fn` must satisfy the syntax constraints of the graph mode.
+        - Users do not need to implement a custom backward function.
+        - `vararg`, `kwarg`, `kwonlyargs` and free variables are not supported in user-defined function.
+
+    Args:
+        fn (Function): Mindspore's function, user-defined function.
+        infer_shape (Function): Mindspore's function, user-defined infer_shape function.
+        infer_dtype (Function): Mindspore's function, user-defined infer_dtype function.
+
+    Inputs:
+        The inputs of user-defined `fn`.
+
+    Outputs:
+        The outputs of user-defined `fn`.
+
+    Raises:
+        RuntimeError: if not used in GRAPH_MODE.
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import context, nn, ops, Tensor, Parameter
+        >>>
+        >>> np_weight0 = np.array([1.0, 2.0, 3.0])
+        >>> np_weight1 = np.array([4.0, 5.0, 6.0])
+        >>> np_input_x = np.array([7.0, 8.0, 9.0])
+        >>>
+        >>> def infer_dtype(args):
+        ...     return args
+        >>>
+        >>> def infer_shape(args):
+        ...     return args
+        >>>
+        >>> def mul_by(*args):
+        ...     def inner(x):
+        ...         return args[0] * x
+        ...     return inner
+        >>>
+        >>> NUMBER_100 = 100
+        >>> class MorphNet(nn.Cell):
+        ...     def __init__(self):
+        ...         super(MorphNet, self).__init__()
+        ...         self.weight0 = Parameter(Tensor(np_weight0, ms.float32), name="weight0")
+        ...         self.weight1 = Parameter(Tensor(np_weight1, ms.float32), name="weight1")
+        ...         self.mul_by_100 = ops.Morph(mul_by(NUMBER_100), infer_shape, infer_dtype)
+        ...     def construct(self, x):
+        ...         a = x * self.weight0
+        ...         b = self.mul_by_100(a)
+        ...         out = b * self.weight1
+        ...         return out
+        >>>
+        >>> context.set_context(mode=context.GRAPH_MODE)
+        >>> input_x = Tensor(np_input_x, ms.float32)
+        >>> net = MorphNet()
+        >>> grad_op = ops.GradOperation(get_all=True, get_by_list=True)
+        >>> grad_net = grad_op(net, net.trainable_params())
+        >>> bwd_out = grad_net(input_x)
+        >>> x_grad = bwd_out[0][0].asnumpy()
+        >>> weight0_grad = bwd_out[1][0].asnumpy()
+        >>> weight1_grad = bwd_out[1][1].asnumpy()
+        >>> print("x_grad", x_grad)
+        >>> print("weight0_grad", weight0_grad)
+        >>> print("weight1_grad", weight1_grad)
+        x_grad [ 400. 1000. 1800.]
+        weight0_grad [2800. 4000. 5400.]
+        weight1_grad [ 700. 1600. 2700.]
+    """
+    @prim_attr_register
+    def __init__(self, fn, infer_shape, infer_dtype):
+        self.add_prim_attr('side_effect_backprop', True)
+        self.add_prim_attr('side_effect_mem', True)
+        self.add_prim_attr('side_effect_io', True)
+        self.add_prim_attr('__metamorphosis__', fn)
+        self._infer_shape = infer_shape
+        self._infer_dtype = infer_dtype
+
+    def infer_shape(self, *args):
+        return self._infer_shape(*args)
+
+    def infer_dtype(self, *args):
+        return self._infer_dtype(*args)
+
+    def __call__(self, *args):
+        raise RuntimeError("Morph is only supported in GRAPH_MODE.")
+
+
 class HookBackward(PrimitiveWithInfer):
     """
     This operation is used as a tag to hook gradient in intermediate variables. Note that this function
@@ -527,8 +625,7 @@ class HookBackward(PrimitiveWithInfer):
     def __init__(self, hook_fn, cell_id=""):
         """Initialize HookBackward."""
         super(HookBackward, self).__init__(self.__class__.__name__)
-        if not check_hook_fn("HookBackward", hook_fn):
-            return
+        check_hook_fn(hook_fn)
         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.")
@@ -600,6 +697,9 @@ class Print(Primitive):
         self.add_prim_attr("side_effect_io", True)
 
     def __call__(self, *args):
+        # Add for jit context.
+        if jit_context() and jit_context().compiled:
+            return
         for arg in args:
             if isinstance(arg, Parameter):
                 print(Tensor_.__repr__(arg))
@@ -607,6 +707,9 @@ class Print(Primitive):
                 print(arg.__repr__())
             else:
                 print(arg)
+        # Add for jit context.
+        if jit_context():
+            jit_context().run_op(self, None, *args)
 
 
 class Assert(PrimitiveWithInfer):