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

mindspore/ops/operations/custom_ops.py

@@ -40,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
@@ -73,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
 
 
@@ -94,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")
@@ -233,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]):
@@ -249,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:
+              This could be used when func_type is "aot".
 
-                   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")
+              for "aot":
 
-                   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")
+              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.
 
-                 - Define a "xxx.cc"/"xxx.cu" file:
+              - "xxx.so" file generation:
 
-                   "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:
+                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")
 
-                   .. code-block::
+                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")
 
-                       int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
-                               void *stream, void *extra)
+              - Define a "xxx.cc"/"xxx.cu" file:
 
-                   Parameters:
+                "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:
 
-                   - 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
+                .. code-block::
 
-                   Return Value(int):
+                    int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
+                            void *stream, void *extra)
 
-                   - 0: MindSpore will continue to run if this aot kernel is successfully executed
-                   - others: MindSpore will raise exception and exit
+                Parameters:
 
-                   Examples: see details in tests/st/ops/graph_kernel/custom/aot_test_files/
+                - 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
 
-                 - Use it in Custom:
+                Return Value(int):
 
-                   .. code-block::
+                - 0: MindSpore will continue to run if this aot kernel is successfully executed
+                - others: MindSpore will raise exception and exit
 
-                       Custom(func="{dir_path}/{file_name}:{func_name}",...)
-                       (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32,
-                       "aot"))
+                Examples: see details in tests/st/ops/graph_kernel/custom/aot_test_files/
 
-                 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:
+              - Use it in Custom:
 
-                 - 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"` .
+                .. code-block::
 
-              2. for "julia":
+                    Custom(func="{dir_path}/{file_name}:{func_name}",...)
+                    (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32,
+                    "aot"))
 
-                 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.
+              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:
 
-                 Examples: see details in tests/st/ops/graph_kernel/custom/julia_test_files/
+              - 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"` .
 
-                 - Use it in Custom:
-
-                   .. code-block::
-
-                       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.
 
@@ -353,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`.
 
@@ -364,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
+        func_type (str, optional): The implementation type of `func`, should be one of
+            [ ``"aot"`` , ``"pyfunc"``]. Default: ``"pyfunc"``.
 
-            [ ``"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
-            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
@@ -428,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
@@ -444,11 +426,13 @@ class Custom(ops.PrimitiveWithInfer):
         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.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:
@@ -460,23 +444,22 @@ class Custom(ops.PrimitiveWithInfer):
             add_pyfunc(func_id, self.func)
             self.add_prim_attr("fn_id", func_id)
 
-        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._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()
@@ -488,25 +471,22 @@ class Custom(ops.PrimitiveWithInfer):
         self._update_attr()
 
         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)
+        self._generate_get_workspace_size_func()
 
-    def set_infer_flag(self):
+    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 _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):
@@ -586,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):
@@ -644,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))
@@ -962,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):
@@ -1034,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."""
@@ -1140,6 +1091,88 @@ 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.is_ascend_c:
             res = pyboost_custom_ext(self.custom_pyboost, [args])
@@ -1174,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
@@ -1193,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
@@ -1205,11 +1245,23 @@ class CustomOpBuilder:
         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")
+        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):
         """
@@ -1228,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):
@@ -1260,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
@@ -1275,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