mindspore 2.6.0rc1__cp310-cp310-win_amd64.whl → 2.7.0rc1__cp310-cp310-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
@@ -39,7 +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
+from ._custom_ops_utils import ExtensionBuilder, CustomCodeGenerator, CustomInfoGenerator
 
 if platform.system() != "Windows":
     import fcntl
@@ -72,13 +73,18 @@ def _get_cache_path():
     """
     cache_path = os.getenv('MS_COMPILER_CACHE_PATH')
     if cache_path is None:
-        cache_path = "./akg_kernel_meta/"
+        cache_path = "./custom_kernel_meta/"
     elif cache_path[-1] != "/":
         cache_path = cache_path + "/"
 
     if not os.path.exists(cache_path):
         os.makedirs(cache_path, exist_ok=True)
 
+    # for distributed case, we create folders separately to avoid conflict
+    if GlobalComm.INITED:
+        cache_path = os.path.join(cache_path, "rank_" + str(get_rank()), "")
+        os.makedirs(cache_path, exist_ok=True)
+
     return cache_path
 
 
@@ -93,10 +99,6 @@ def _compile_aot(file):
         str, the path to the compiled library.
     """
     cache_path = _get_cache_path()
-    # for distributed case, we create folders separately to avoid conflict
-    if GlobalComm.INITED:
-        cache_path = os.path.join(cache_path, "rank_" + str(get_rank()), "")
-        os.makedirs(cache_path, exist_ok=True)
 
     res_path = importlib.util.find_spec("mindspore").origin
     find_pos = res_path.find("__init__.py")
@@ -109,12 +111,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 +150,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)
 
@@ -224,14 +234,13 @@ class Custom(ops.PrimitiveWithInfer):
     <https://www.mindspore.cn/tutorials/en/master/custom_program/op_custom.html>`_ .
 
     .. warning::
-        - This is an experimental API that is subject to change.
+        This is an experimental API that is subject to change.
 
     .. note::
         The supported platforms are determined by the input `func_type`. The supported platforms are as follows:
 
         - "aot": supports ["GPU", "CPU", "Ascend"].
         - "pyfunc": supports ["CPU"].
-        - "julia": supports ["CPU"].
 
     Args:
         func (Union[function, str]):
@@ -240,101 +249,84 @@ class Custom(ops.PrimitiveWithInfer):
               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".
-
-              1. for "aot":
-
-                 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.
-
-                 - "xxx.so" file generation:
-
-                   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. "g++ --shared -fPIC  -o add.so add.cc")
+              This could be used when func_type is "aot".
 
-                 - Define a "xxx.cc"/"xxx.cu" file:
+              for "aot":
 
-                   "aot" is a cross-platform identity. The functions defined in "xxx.cc" or "xxx.cu" share
-                   the same args. Typically, the function should be as:
+              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.
 
-                   .. code-block::
+              - "xxx.so" file generation:
 
-                       int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
-                               void *stream, void *extra)
+                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")
 
-                   Parameters:
+                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")
 
-                   - nparam(int): total number of inputs plus outputs; suppose the operator has 2 inputs and 3 outputs,
-                     then nparam=5
-                   - params(void \*\*): a pointer to the array of inputs and outputs' pointer; the pointer type of
-                     inputs and outputs is void \* ; suppose the operator has 2 inputs and 3 outputs, then the first
-                     input's pointer is params[0] and the second output's pointer is params[3]
-                   - ndims(int \*): a pointer to the array of inputs and outputs' dimension num; suppose params[i] is a
-                     1024x1024 tensor and params[j] is a 77x83x4 tensor, then ndims[i]=2, ndims[j]=3.
-                   - shapes(int64_t \*\*): a pointer to the array of inputs and outputs' shapes(int64_t \*); the ith
-                     input's jth dimension's size is shapes[i][j](0<=j<ndims[i]); suppose params[i] is a 2x3 tensor and
-                     params[j] is a 3x3x4 tensor, then shapes[i][0]=2, shapes[j][2]=4.
-                   - dtypes(const char \*\*): a pointer to the array of inputs and outputs' types(const char \*);
-                     (ex. "float32", "float16", "float", "float64", "int", "int8", "int16", "int32", "int64", "uint",
-                     "uint8", "uint16", "uint32", "uint64", "bool")
-                   - stream(void \*): stream pointer, only used in cuda file
-                   - extra(void \*): used for further extension
+              - Define a "xxx.cc"/"xxx.cu" file:
 
-                   Return Value(int):
+                "aot" is a cross-platform identity. The functions defined in "xxx.cc" or "xxx.cu" share
+                the same args. Typically, the function should be as:
 
-                   - 0: MindSpore will continue to run if this aot kernel is successfully executed
-                   - others: MindSpore will raise exception and exit
+                .. code-block::
 
-                   Examples: see details in tests/st/ops/graph_kernel/custom/aot_test_files/
+                    int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
+                            void *stream, void *extra)
 
-                 - Use it in Custom:
+                Parameters:
 
-                   .. code-block::
+                - nparam(int): total number of inputs plus outputs; suppose the operator has 2 inputs and 3 outputs,
+                  then nparam=5
+                - params(void \*\*): a pointer to the array of inputs and outputs' pointer; the pointer type of
+                  inputs and outputs is void \* ; suppose the operator has 2 inputs and 3 outputs, then the first
+                  input's pointer is params[0] and the second output's pointer is params[3]
+                - ndims(int \*): a pointer to the array of inputs and outputs' dimension num; suppose params[i] is a
+                  1024x1024 tensor and params[j] is a 77x83x4 tensor, then ndims[i]=2, ndims[j]=3.
+                - shapes(int64_t \*\*): a pointer to the array of inputs and outputs' shapes(int64_t \*); the ith
+                  input's jth dimension's size is shapes[i][j](0<=j<ndims[i]); suppose params[i] is a 2x3 tensor and
+                  params[j] is a 3x3x4 tensor, then shapes[i][0]=2, shapes[j][2]=4.
+                - dtypes(const char \*\*): a pointer to the array of inputs and outputs' types(const char \*);
+                  (ex. "float32", "float16", "float", "float64", "int", "int8", "int16", "int32", "int64", "uint",
+                  "uint8", "uint16", "uint32", "uint64", "bool")
+                - stream(void \*): stream pointer, only used in cuda file
+                - extra(void \*): used for further extension
 
-                       Custom(func="{dir_path}/{file_name}:{func_name}",...)
-                       (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32,
-                       "aot"))
+                Return Value(int):
 
-                 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/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 function:
+                - 0: MindSpore will continue to run if this aot kernel is successfully executed
+                - others: MindSpore will raise exception and exit
 
-                 - 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"` .
+                Examples: see details in tests/st/ops/graph_kernel/custom/aot_test_files/
 
-              2. for "julia":
+              - Use it in Custom:
 
-                 Currently "julia" supports CPU(linux only) platform.
-                 For julia use JIT compiler, and julia support c api to call julia code.
-                 The Custom can directly launches user defined "xxx.jl" file as an operator.
-                 Users need to write a "xxx.jl" file which include modules and functions,
-                 and offer the path of the file along with a module name and function name.
+                .. code-block::
 
-                 Examples: see details in tests/st/ops/graph_kernel/custom/julia_test_files/
+                    Custom(func="{dir_path}/{file_name}:{func_name}",...)
+                    (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32,
+                    "aot"))
 
-                 - Use it in Custom:
+              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/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 function:
 
-                   .. code-block::
+              - 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"` .
 
-                       Custom(func="{dir_path}/{file_name}:{module_name}:{func_name}",...)
-                       (ex. Custom(func="./add.jl:Add:add", out_shape=[1], out_dtype=mstype.float32, "julia"))
-
-        out_shape (Union[function, list, tuple]): The output shape infer function or the value of output shape of
-            `func`. Default: ``None`` .
+        out_shape (Union[function, list, tuple], optional): The output shape infer function or the value of output
+            shape of `func`. Default: ``None`` .
 
             If func has single output, then the value of output shape is a list or tuple of int.
 
@@ -344,8 +336,8 @@ class Custom(ops.PrimitiveWithInfer):
             The input can be None only when the func_type input is "hybrid". In this case, the automatic infer
             shape mechanic will be enabled.
 
-        out_dtype (Union[function, :class:`mindspore.dtype`, tuple[:class:`mindspore.dtype`]]): The output data type
-            infer function or the value of output data type of `func`. Default: ``None`` .
+        out_dtype (Union[function, :class:`mindspore.dtype`, tuple[:class:`mindspore.dtype`]], optional): The output
+            data type infer function or the value of output data type of `func`. Default: ``None`` .
 
             If func has single output, then the value of output shape is a `mindspore.dtype`.
 
@@ -355,13 +347,12 @@ class Custom(ops.PrimitiveWithInfer):
             The input can be None only when the func_type input is "hybrid". In this case, the automatic infer
             value mechanic will be enabled.
 
-        func_type (str): The implementation type of `func`, should be one of
-
-            [ ``"aot"`` , ``"pyfunc"`` , ``"julia"`` ].
+        func_type (str, optional): The implementation type of `func`, should be one of
+            [ ``"aot"`` , ``"pyfunc"``]. Default: ``"pyfunc"``.
 
-        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
-            json format of type str or dict. The reg info specifies supported data types and formats of inputs and
+        bprop (function, optional): The back propagation function of `func`. Default: ``None`` .
+        reg_info (Union[str, dict, list, tuple], optional): Represents the registration information(reg info) of `func`
+            with json format of type str or dict. The reg info specifies supported data types and formats of inputs and
             outputs, attributes and target of `func`. Default: ``None`` .
 
             If reg info is a list or tuple, then each item should be with json format of type str or dict, which
@@ -419,11 +410,11 @@ class Custom(ops.PrimitiveWithInfer):
     op_path_in_cache = []  # Save paths for op functions created in the cached.
     custom_aot_warning = True  # Flag to enable warnings about custom aot path white list
 
-    def __init__(self, func, out_shape=None, out_dtype=None, func_type="hybrid", bprop=None, reg_info=None):
+    def __init__(self, func, out_shape=None, out_dtype=None, func_type="pyfunc", bprop=None, reg_info=None):
         super().__init__("Custom")
 
         self.supported_targets = [ASCEND, GPU, CPU]
-        self.supported_func_type = ["hybrid", "akg", "tbe", "aicpu", "aot", "pyfunc", "julia"]
+        self.supported_func_type = ["hybrid", "akg", "tbe", "aicpu", "aot", "pyfunc"]
         self.log_prefix = "For '{}', 'func_type': {}, 'func': {}".format(self.name, func_type, func)
         self.func = func
         self.func_type = func_type
@@ -435,10 +426,13 @@ class Custom(ops.PrimitiveWithInfer):
         self._is_ms_kernel = False
         self.out_shape = out_shape
         self.out_dtype = out_dtype
+        self.reg_info = reg_info
+        self.is_ascend_c = (context.get_context("device_target") == "Ascend" and self.func_type == "aot")
 
         self._check_platform()
         self._check_func()
-        self._update_func_info(reg_info)
+        self._generate_reg_info()
+        self._update_func_info(self.reg_info)
         self.add_prim_attr("func_name", self.func_name)
         self.add_prim_attr("uniq_name", self.uniq_name)
         if self.func_type == HYBRID_TYPE:
@@ -450,25 +444,22 @@ 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)
-        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._set_infer_flag()
+        self._set_multi_output_flag()
 
         self.bprop = bprop
         self.fake_output = False
         self.single_scalar_output = False
-        if not self.out_dtype and not self.func_type == "pyfunc":
-            self.fake_output = True
-        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)
+        if self.func_type == "pyfunc":
+            if not self.out_dtype:
+                self.fake_output = True
+            elif not self.out_shape:
+                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)
 
         # Register info
-        self._register_info(reg_info)
+        self._register_info(self.reg_info)
 
         if func_type == "akg":
             self._set_akg_kernel_type()
@@ -479,11 +470,23 @@ class Custom(ops.PrimitiveWithInfer):
         self.add_prim_attr("func_type", self.func_type)
         self._update_attr()
 
-        self.enable_pyboost = (context.get_context("device_target") == "Ascend" and self.func_type == "aot")
-        if self.enable_pyboost:
+        if self.is_ascend_c:
             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)
+        self._generate_get_workspace_size_func()
+
+    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_multi_output_flag(self):
+        outputs = self.reg_info.get("outputs", []) if self.reg_info else []
+        self.multi_output = len(outputs) > 1 or (len(outputs) == 1 and outputs[0].get("paramType") == "dynamic")
+        self.add_prim_attr("multi_output", self.multi_output)
 
     def __infer__(self, *args):
         if callable(self.out_shape):
@@ -563,22 +566,6 @@ class Custom(ops.PrimitiveWithInfer):
             self.func_type = HYBRID_TYPE
             self._hybrid_func_analyser()
 
-    def _check_julia_func(self):
-        """Check the validity of julia func"""
-        if not isinstance(self.func, str):
-            raise TypeError("{}, 'func' must be of type str, but got {}".format(self.log_prefix, type(self.func)))
-        if self.func.count(':') != 2:
-            raise ValueError("{}, the format of 'func' must be file:module:func".format(self.log_prefix))
-        source_file, module, func = self.func.split(':')
-        with open(source_file, 'r') as f:
-            jl = f.read()
-            if 'module ' + module not in jl:
-                raise Exception("{}, module {} is not found in source file {}!"
-                                .format(self.log_prefix, module, source_file))
-            if 'function ' + func not in jl:
-                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):
@@ -621,8 +608,6 @@ class Custom(ops.PrimitiveWithInfer):
         if self.func_type == "aot":
             self._check_aot_func()
 
-        elif self.func_type == "julia":
-            self._check_julia_func()
         elif self.func_type == HYBRID_TYPE:
             if not hasattr(self.func, MS_KERNEL_FLAG):
                 raise TypeError("{}, 'func' must be a function decorated by kernel".format(self.log_prefix))
@@ -776,6 +761,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
@@ -806,14 +811,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):
@@ -918,7 +924,7 @@ class Custom(ops.PrimitiveWithInfer):
             return reg_info[IMPLY_TYPE]
         # Infer imply_type from func_type
         func_type_to_imply_type = {"hybrid": AKG, "akg": AKG, "tbe": TBE, "aicpu": "AiCPU", "pyfunc": target,
-                                   "julia": target, "aot": "BiSheng" if target == ASCEND else target}
+                                   "aot": "BiSheng" if target == ASCEND else target}
         return func_type_to_imply_type.get(self.func_type, AKG)
 
     def _save_attr(self, reg_info):
@@ -990,17 +996,6 @@ class Custom(ops.PrimitiveWithInfer):
                 self.set_device(GPU)
             elif registered_targets == [CPU]:
                 self.set_device(CPU)
-        elif self.func_type == "julia":
-            self.set_device(CPU)
-            device_target = context.get_context('device_target')
-            if device_target == CPU:
-                pass
-            elif device_target == GPU and registered_targets and registered_targets == [CPU]:
-                logger.warning("{}, only supports CPU platform, but got registered target {}. "
-                               "We will run it on CPU".format(self.log_prefix, registered_targets))
-            else:
-                raise ValueError("{}, only supports CPU platform, but got target {}."
-                                 .format(self.log_prefix, device_target))
 
     def _update_attr(self):
         """Add input_names, attr_names, primitive_target to primitive's attr."""
@@ -1096,8 +1091,90 @@ class Custom(ops.PrimitiveWithInfer):
 
         return infer_shape, infer_dtype, infer_value
 
+    def _generate_reg_info(self):
+        if not self.is_ascend_c:
+            return
+        if self.reg_info is None:
+            func_name, _ = self._split_func()
+            if func_name.startswith("aclnn"):
+                func_name = func_name[len("aclnn"):]
+            reg_info_generator = CustomInfoGenerator(func_name)
+            self.reg_info = reg_info_generator.generate_custom_reg_op()
+
+    def _split_func(self):
+        func_list = self.func.split(":")
+        func_path = ""
+        if len(func_list) == 2:
+            func_path = func_list[0]
+            func_name = func_list[1]
+        else:
+            func_name = self.func
+        return func_name, func_path
+
+    def _generate_get_worspace_size_func_by_types(self, aclnn_api_types):
+        """generate custom GetWorkSpaceSize func by aclnn api types"""
+        if not self.is_ascend_c:
+            return
+
+        input_output_types = []
+        if isinstance(aclnn_api_types, str):
+            params = re.split(r',\s*', aclnn_api_types)
+            for param in params:
+                param = param.replace('const ', '')
+                type_part = re.search(r'^\s*(\w+\s*\*+|\w+)', param).group(1)
+                type_part = type_part.replace(' ', '')
+                input_output_types.append(type_part)
+        elif isinstance(aclnn_api_types, list):
+            input_output_types = aclnn_api_types
+        else:
+            raise RuntimeError(f"Unsupported type: {type(aclnn_api_types)}, support type is list or string.")
+
+        func_name, _ = self._split_func()
+        file_path = os.path.join(_get_cache_path(), func_name, func_name + "_callback.cc")
+
+        file_path = os.path.abspath(file_path)
+        dir_path = os.path.dirname(file_path)
+        os.makedirs(dir_path, exist_ok=True)
+
+        custom_builder = CustomCodeGenerator()
+        callback_func = custom_builder.generate_callback_by_types(func_name, self.reg_info, input_output_types)
+
+        with open(file_path, 'w') as f:
+            f.write(callback_func)
+
+        custom_callback_func_path = _compile_aot(file_path)
+        custom_callback_func = custom_callback_func_path + ":" + func_name
+        self.add_prim_attr("custom_callback_func", custom_callback_func)
+        self.add_prim_attr("custom_inputs_type", input_output_types[:-2])
+
+    def _generate_get_workspace_size_func(self):
+        """generate custom GetWorkSpaceSize func"""
+        if not self.is_ascend_c:
+            return
+        func_name, _ = self._split_func()
+        file_path = os.path.join(_get_cache_path(), func_name, func_name + "_callback.cc")
+
+        file_path = os.path.abspath(file_path)
+        dir_path = os.path.dirname(file_path)
+        os.makedirs(dir_path, exist_ok=True)
+
+        custom_info_generator = CustomInfoGenerator(func_name)
+        api_types = custom_info_generator.get_aclnn_api_types()
+        custom_builder = CustomCodeGenerator()
+        if api_types == []:
+            api_types = custom_builder.get_api_types_by_reg_info(self.reg_info)
+
+        callback_func = custom_builder.generate_callback_by_types(func_name, self.reg_info, api_types)
+        with open(file_path, 'w') as f:
+            f.write(callback_func)
+
+        custom_callback_func_path = _compile_aot(file_path)
+        custom_callback_func = custom_callback_func_path + ":" + func_name
+        self.add_prim_attr("custom_callback_func", custom_callback_func)
+        self.add_prim_attr("custom_inputs_type", api_types[:-2])
+
     def __call__(self, *args):
-        if self.enable_pyboost:
+        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)
@@ -1130,6 +1207,15 @@ class CustomOpBuilder:
         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.
 
+            - build_dir (str, optional): The directory used to generate the operator build files.
+              If this argument is set, the provided path will be used directly.
+              If not set, a subdirectory named after the operator's name will be created under the path specified by
+              the environment variable `MS_COMPILER_CACHE_PATH` (defaulting to "./kernel_meta"), and the files will
+              be placed in this subdirectory. Default: ``None``.
+
+            - enable_atb (bool, optional): Whether to call ATB (Ascend Transformer Boost) operator. If set to ``True``,
+              the `backend` must be ``Ascend`` or left empty. Default: ``False``.
+
     .. 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
@@ -1149,9 +1235,7 @@ class CustomOpBuilder:
         ... )
         >>> 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
@@ -1160,11 +1244,24 @@ class CustomOpBuilder:
         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")
+        self.build_dir = kwargs.get("build_dir")
+        self.enable_atb = kwargs.get("enable_atb", False)
+        self._ms_path = os.path.dirname(os.path.abspath(ms.__file__))
+        if self.enable_atb:
+            if backend is not None and backend != "Ascend":
+                raise ValueError("For 'CustomOpBuilder', when 'enable_atb' is set to True, the 'backend' must be "
+                                 f"'Ascend' (or left implicit), but got '{backend}'")
+            self.backend = "Ascend"
         if self.backend == "Ascend":
-            self.ascend_cann_path = os.getenv("ASCEND_OPP_PATH").split('opp')[0]
+            ascend_opp_path = os.getenv("ASCEND_OPP_PATH")
+            if not ascend_opp_path:
+                raise ValueError("Environment variable 'ASCEND_OPP_PATH' must be set for Ascend backend.")
+            self.ascend_cann_path = ascend_opp_path.split('opp')[0]
+
+            if self.enable_atb:
+                self.atb_home_path = os.getenv("ATB_HOME_PATH")
+                if not self.atb_home_path:
+                    raise ValueError("Environment variable 'ATB_HOME_PATH' must be set when 'enable_atb' is True.")
 
     def get_sources(self):
         """
@@ -1183,29 +1280,31 @@ class CustomOpBuilder:
             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"))
+        include_list.append(self._ms_path)
+        include_list.append(os.path.join(self._ms_path, "include"))
+        include_list.append(os.path.join(self._ms_path, "include", "third_party"))
+        include_list.append(os.path.join(self._ms_path, "include", "third_party", "robin_hood_hashing"))
+        include_list.append(os.path.join(self._ms_path, "include", "third_party", "securec", "include"))
 
         if self.backend == "Ascend":
             include_list.append(os.path.join(self.ascend_cann_path, "include"))
+            if self.enable_atb:
+                include_list.append(os.path.join(self.atb_home_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")
+        ms_inner_path = os.path.join(self._ms_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")
+        include_list.append(os.path.join(ms_inner_path, "core", "include"))
+        include_list.append(os.path.join(ms_inner_path, "core", "mindrt", "include"))
+        include_list.append(os.path.join(ms_inner_path, "core", "mindrt"))
+        include_list.append(os.path.join(ms_inner_path, "ops"))
+        include_list.append(os.path.join(ms_inner_path, "ops", "kernel", "include"))
+        include_list.append(os.path.join(ms_inner_path, "ccsrc"))
+        include_list.append(os.path.join(ms_inner_path, "ccsrc", "include"))
+        include_list.append(os.path.join(ms_inner_path, "ccsrc", "minddata", "mindrecord", "include"))
         return include_list
 
     def get_cflags(self):
@@ -1215,10 +1314,12 @@ class CustomOpBuilder:
         Returns:
             list[str], A list of C++ compiler flags.
         """
-        flags = ['-fstack-protector-all', '-fPIC', '-pie']
-        flags += ['-DENABLE_FAST_HASH_TABLE=1']
+        flags = [f'-DMS_EXTENSION_NAME={self.name}', '-D_GLIBCXX_USE_CXX11_ABI=0', '-DENABLE_FAST_HASH_TABLE=1']
+        flags += ['-std=c++17', '-fstack-protector-all', '-fPIC', '-pie']
         if self.backend == "Ascend":
             flags.append('-DCUSTOM_ASCEND_OP')
+            if self.enable_atb:
+                flags.append('-DCUSTOM_ENABLE_ATB')
         if self.cflags is not None:
             flags.append(self.cflags)
         return flags
@@ -1230,18 +1331,25 @@ class CustomOpBuilder:
         Returns:
             list[str], A list of linker flags.
         """
-        flags = ['-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath', '-s']
+        flags = ['-shared']
+        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')),
+            f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib'))}",
             '-lmindspore_core',
             '-lmindspore_ms_backend',
-            '-lmindspore_pynative'
+            '-lmindspore_pynative',
+            '-lmindspore_extension'
         ]
         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(f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin'))}")
+            flags.append(f"-L{os.path.abspath(os.path.join(self.ascend_cann_path, 'lib64'))}")
             flags.append('-lascendcl')
             flags.append('-l:libmindspore_ascend.so.2')
+            if self.enable_atb:
+                flags.append(f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin', 'ascend'))}")
+                flags.append('-lmindspore_extension_ascend_atb')
+                flags.append(f"-L{os.path.abspath(os.path.join(self.atb_home_path, 'lib'))}")
+                flags.append('-latb')
         if self.ldflags is not None:
             flags.append(self.ldflags)
         return flags
@@ -1256,7 +1364,7 @@ class CustomOpBuilder:
         Returns:
             str, The path to the compiled module.
         """
-        return ExtensionBuilder().build(
+        return ExtensionBuilder(self._get_build_directory()).build(
             module_name=self.name,
             sources=self.get_sources(),
             extra_include_paths=self.get_include_paths(),
@@ -1283,3 +1391,15 @@ class CustomOpBuilder:
         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