mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0rc1__cp311-cp311-win_amd64.whl

mindspore/ops/operations/custom_ops.py

@@ -29,7 +29,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 +39,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
@@ -183,10 +183,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 +221,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.
@@ -223,8 +229,6 @@ class Custom(ops.PrimitiveWithInfer):
     .. note::
         The supported platforms are determined by the input `func_type`. The supported platforms are as follows:
 
-        - "hybrid": supports ["GPU", "CPU"].
-        - "akg": supports ["GPU", "CPU"].
         - "aot": supports ["GPU", "CPU", "Ascend"].
         - "pyfunc": supports ["CPU"].
         - "julia": supports ["CPU"].
@@ -233,11 +237,7 @@ class Custom(ops.PrimitiveWithInfer):
         func (Union[function, str]):
 
             - function: If func is of function type, then func should be a Python function which describes the
-              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 pure python function
-              3. An kernel decorated function written by the Hybrid DSL.
+              computation logic of a user defined operator.
 
             - 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".
@@ -251,13 +251,11 @@ class Custom(ops.PrimitiveWithInfer):
 
                  - "xxx.so" file generation:
 
-                   1) GPU Platform: Given user defined "xxx.cu" file (ex. "{path}/add.cu"),
-                   use nvcc command to compile
-                   it.(ex. :code:`nvcc --shared -Xcompiler -fPIC -o add.so add.cu`)
+                   1) GPU Platform: Given user defined "xxx.cu" file (ex. "{path}/add.cu"), use nvcc command to compile
+                   it.(ex. "nvcc --shared -Xcompiler -fPIC -o add.so add.cu")
 
-                   2) CPU Platform: Given user defined "xxx.cc" file (ex. "{path}/add.cc"),
-                   use g++/gcc command to
-                   compile it.(ex. :code:`g++ --shared -fPIC -o add.so add.cc`)
+                   2) CPU Platform: Given user defined "xxx.cc" file (ex. "{path}/add.cc"), use g++/gcc command to
+                   compile it.(ex. "g++ --shared -fPIC  -o add.so add.cc")
 
                  - Define a "xxx.cc"/"xxx.cu" file:
 
@@ -305,14 +303,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":
@@ -356,7 +357,7 @@ class Custom(ops.PrimitiveWithInfer):
 
         func_type (str): The implementation type of `func`, should be one of
 
-            [ ``"hybrid"`` , ``"akg"`` , ``"aot"`` , ``"pyfunc"`` , ``"julia"`` ].
+            [ ``"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
@@ -400,50 +401,14 @@ class Custom(ops.PrimitiveWithInfer):
         >>> input_x = Tensor(np.ones([16, 16]).astype(np.float32))
         >>> input_y = Tensor(np.ones([16, 16]).astype(np.float32))
         >>>
-        >>> # Example, func_type = "hybrid"
-        >>> # This is the default func_type in Custom,
-        >>> # and both out_shape and out_dtype can be None(default value).
-        >>> # In this case, the input func must be a function written in the Hybrid DSL
-        >>> # and decorated by @kernel.
-        >>> @kernel
-        ... def add_script(a, b):
-        ...     c = output_tensor(a.shape, a.dtype)
-        ...     for i0 in range(a.shape[0]):
-        ...         for i1 in range(a.shape[1]):
-        ...             c[i0, i1] = a[i0, i1] + b[i0, i1]
-        ...     return c
-        >>>
-        >>> test_op_hybrid = ops.Custom(add_script)
-        >>> output = test_op_hybrid(input_x, input_y)
-        >>> # the result will be a 16 * 16 tensor with all elements 2
-        >>> print(output.shape)
-        (16, 16)
-        >>> # Example, func_type = "aot"
-        >>> def test_aot(x, y, out_shapes, out_types):
-        ...     program = ops.Custom("./reorganize.so:CustomReorganize", out_shapes, out_types, "aot")
-        ...     out = program(x, y)
-        ...     return out
-        >>>
         >>> # Example, func_type = "pyfunc"
-        >>> def func_multi_output(x1, x2):
-        ...     return (x1 + x2), (x1 - x2)
+        >>> def func_pyfunc(x1, x2):
+        ...     return x1 + x2
         >>>
-        >>> test_pyfunc = ops.Custom(func_multi_output, lambda x, _: (x, x), lambda x, _: (x, x), "pyfunc")
+        >>> test_pyfunc = ops.Custom(func_pyfunc, lambda x, _: x, lambda x, _: x, "pyfunc")
         >>> output = test_pyfunc(input_x, input_y)
-        >>>
-        >>> # Example, func_type = "julia"
-        >>> # julia code:
-        >>> # add.jl
-        >>> # module Add
-        >>> # function add(x, y, z)
-        >>> #   z .= x + y
-        >>> #   return z
-        >>> # end
-        >>> # end
-        >>> def test_julia(x, y, out_shapes, out_types):
-        ...     program = ops.Custom("./add.jl:Add:add", out_shapes, out_types, "julia")
-        ...     out = program(x, y)
-        ...     return out
+        >>> print(output.shape)
+        (16, 16)
     """
 
     registered_func = {}
@@ -469,6 +434,7 @@ class Custom(ops.PrimitiveWithInfer):
         self._func_compile_attrs = {}
         self._is_ms_kernel = False
         self.out_shape = out_shape
+        self.out_dtype = out_dtype
 
         self._check_platform()
         self._check_func()
@@ -486,13 +452,17 @@ class Custom(ops.PrimitiveWithInfer):
 
         if self.out_shape is None and self.func_type == "aot":
             self.add_prim_attr("cpp_infer_shape", True)
-        self.out_dtype = out_dtype
+        if self.out_dtype is None and self.func_type == "aot":
+            self.add_prim_attr("cpp_infer_type", True)
+        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)
@@ -508,10 +478,10 @@ 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
+
+        self.enable_pyboost = (context.get_context("device_target") == "Ascend" and self.func_type == "aot")
+        if self.enable_pyboost:
+            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)
 
@@ -554,10 +524,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 {}"
@@ -1123,9 +1098,188 @@ class Custom(ops.PrimitiveWithInfer):
 
     def __call__(self, *args):
         if self.enable_pyboost:
-            return _convert_stub(pyboost_custom_ext(self.custom_pyboost, [args]))
+            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
+        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().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