mindspore 2.5.0__cp39-cp39-win_amd64.whl → 2.6.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops_generate/{pyboost_op_cpp_code_generator.py → pyboost/pyboost_op_cpp_code_generator.py}

@@ -21,13 +21,15 @@ for different devices (Ascend, CPU, GPU) and manages residual files associated w
 import os
 import re
 
-from pyboost_utils import is_cube, AclnnUtils, get_return_type, merge_strings_by_chunk_size, is_op_multi_output
-import template
-import gen_constants as K
-from gen_utils import save_file
-from op_proto import OpProto
-from op_template_parser import OpTemplateParser
-from base_generator import BaseGenerator
+import common.template as template
+import common.gen_constants as K
+from common.gen_utils import save_file
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+
+from .pyboost_utils import is_cube, AclnnUtils, get_return_type, merge_strings_by_chunk_size, is_op_multi_output, \
+    chunk_list
+from .op_template_parser import OpTemplateParser
 
 
 class PyboostCommonOpHeaderGenerator(BaseGenerator):
@@ -66,7 +68,7 @@ class PyboostCommonOpHeaderGenerator(BaseGenerator):
                                                                                      call_args=call_args_with_type,
                                                                                      return_type=cpp_func_return,
                                                                                      output_is_tuple=output_is_tuple)
-            save_path = os.path.join(work_path, f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
+            save_path = os.path.join(work_path, f"{K.MS_PYBOOST_BASE_PATH}/auto_generate/")
             file_name = f"{op_proto.op_name}.h"
             save_file(save_path, file_name, pyboost_op_header_str)
 
@@ -93,9 +95,11 @@ class PyboostOpHeaderGenerator(BaseGenerator):
                          "gpu": template.PYBOOST_GPU_OP_HEADER_TEMPLATE,
                          "cpu": template.PYBOOST_CPU_OP_HEADER_TEMPLATE}
         if device not in template_dict:
-            raise ValueError(f"Device must be ascend, gpu, or cpu, {device} is not supported")
+            raise ValueError(
+                f"Device must be ascend, gpu, or cpu, {device} is not supported")
         self.PYBOOST_OP_HEADER_TEMPLATE = template_dict[device]
         self.code_generate_path = f"{K.MS_OPS_KERNEL_PATH}/{device}/pyboost/auto_generate/"
+        self.hccl_code_generate_path = "mindspore/ccsrc/plugin/device/ascend/kernel/hccl/pyboost/auto_generate/"
         self.device = device
 
     def generate(self, work_path, op_protos):
@@ -114,6 +118,7 @@ class PyboostOpHeaderGenerator(BaseGenerator):
                 continue
             if getattr(op_proto.op_dispatch, self.device) == 'None':
                 continue
+            is_ascend_comm_op = op_proto.op_dispatch.is_comm_op and self.device == 'ascend'
             op_parser = OpTemplateParser(op_proto)
             op_name_str = op_proto.op_class.name
             call_args_with_type = op_parser.parse_call_args_with_types()
@@ -125,7 +130,7 @@ class PyboostOpHeaderGenerator(BaseGenerator):
                                                                      call_args_with_type=call_args_with_type,
                                                                      return_type=cpp_func_return)
 
-            save_path = os.path.join(work_path, self.code_generate_path)
+            save_path = os.path.join(work_path, self.code_generate_path if not is_ascend_comm_op else self.hccl_code_generate_path)
             file_name = f"{op_proto.op_name}.h"
             save_file(save_path, file_name, pyboost_op_str)
 
@@ -164,14 +169,17 @@ class PyboostOpCppGenerator:
             PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
         else:
-            raise ValueError(f"Device must be ascend, gpu, or cpu, {device} is not supported")
+            raise ValueError(
+                f"Device must be ascend, gpu, or cpu, {device} is not supported")
         self.PYBOOST_CUSTOMIZE_CALL_TEMPLATE = PYBOOST_CUSTOMIZE_CALL_TEMPLATE
         self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE = PYBOOST_SINGLE_OP_HEADER_TEMPLATE
         self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = PYBOOST_SINGLE_OP_SOURCE_TEMPLATE
+        self.PYBOOST_SINGLE_HCLL_OP_HEADER_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_HCLL_OP_HEADER_TEMPLATE
         self.gen_path = gen_path
         self.device = device
 
-    def generate_customize_op_cpp_code(self, op_protos, merge_op_header, merge_op_function):
+    def generate_customize_op_cpp_code(self, op_protos, merge_op_header, merge_op_function, merge_op_inc,
+                                       merge_op_hccl_header=None, merge_op_hccl_function=None, merge_op_hccl_inc=None):
         """
         Generate C++ code for PyBoost operations using the provided operation prototypes.
 
@@ -191,6 +199,7 @@ class PyboostOpCppGenerator:
                 continue
             if getattr(op_proto.op_dispatch, self.device) == 'None':
                 continue
+            is_ascend_comm_op = op_proto.op_dispatch.is_comm_op and self.device == 'ascend'
             op_parser = OpTemplateParser(op_proto)
             call_args = op_parser.parse_original_call_args(op_proto.op_args)
             call_args_with_type = op_parser.parse_call_args_with_types()
@@ -205,20 +214,40 @@ class PyboostOpCppGenerator:
             call_impl = self.PYBOOST_CUSTOMIZE_CALL_TEMPLATE.replace(
                 call_args=call_args,
                 return_values=call_func_outputs,
-                customize_func=getattr(op_proto.op_dispatch, self.device) + "Customize",
+                customize_func=getattr(
+                    op_proto.op_dispatch, self.device) + "Customize",
                 check_expression=check_inplace_func,
             )
-            customize_include = \
-                f'#include "{K.MS_OPS_KERNEL_PATH}/{self.device}/pyboost/customize/{operator_name.lower()}.h"\n'
+            if is_ascend_comm_op and ((merge_op_hccl_header is None) or (merge_op_hccl_function is None)):
+                raise ValueError(f"merge_op_hccl_header and merge_op_hccl_function must be provided for comm op {operator_name}")
+
+            if is_ascend_comm_op:
+                customize_include = \
+                    f'#include "mindspore/ccsrc/plugin/device/ascend/kernel/hccl/pyboost/{operator_name.lower()}.h"\n'
+            else:
+                customize_include = \
+                    f'#include "{K.MS_OPS_KERNEL_PATH}/{self.device}/pyboost/customize/{operator_name.lower()}.h"\n'
+
             register_custom = self._get_register_custom_kernel(op_proto)
             cpp_func_return = _generate_cpp_func_return(op_proto)
-            merge_op_header.append(self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE.replace(operator_name=operator_name,
+            if is_ascend_comm_op:
+                merge_op_hccl_header.append(self.PYBOOST_SINGLE_HCLL_OP_HEADER_TEMPLATE.replace(operator_name=operator_name,
                                                                                   customize_include=customize_include))
-            merge_op_function.append(
-                self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE.replace(op_name=op_name_str,
-                                                               call_args_with_type=call_args_with_type,
-                                                               return_type=cpp_func_return, call_impl=call_impl,
-                                                               register_custom_kernel=register_custom))
+                merge_op_hccl_function.append(
+                    self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE.replace(op_name=op_name_str,
+                                                                call_args_with_type=call_args_with_type,
+                                                                return_type=cpp_func_return, call_impl=call_impl,
+                                                                register_custom_kernel=register_custom))
+                merge_op_hccl_inc.append(op_name_str)
+            else:
+                merge_op_header.append(self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE.replace(operator_name=operator_name,
+                                                                                  customize_include=customize_include))
+                merge_op_function.append(
+                    self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE.replace(op_name=op_name_str,
+                                                                call_args_with_type=call_args_with_type,
+                                                                return_type=cpp_func_return, call_impl=call_impl,
+                                                                register_custom_kernel=register_custom))
+                merge_op_inc.append(op_name_str)
 
     def _get_register_custom_kernel(self, op_proto: OpProto):
         """
@@ -237,7 +266,8 @@ class PyboostOpCppGenerator:
         elif self.device == 'gpu':
             register_custom_kernel = f"MS_REG_PYBOOST_GPU_CUSTOM_KERNEL({op_proto.op_class.name});"
         else:
-            raise ValueError(f"Device must be ascend, gpu, or cpu, {self.device} is not supported")
+            raise ValueError(
+                f"Device must be ascend, gpu, or cpu, {self.device} is not supported")
         return register_custom_kernel
 
 
@@ -275,14 +305,15 @@ class PyboostViewOpCppGenerator:
             PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
         else:
-            raise ValueError(f"Device must be ascend, gpu, or cpu, {device} is not supported")
+            raise ValueError(
+                f"Device must be ascend, gpu, or cpu, {device} is not supported")
         self.PYBOOST_VIEW_CALL_TEMPLATE = PYBOOST_VIEW_CALL_TEMPLATE
         self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE = PYBOOST_SINGLE_OP_HEADER_TEMPLATE
         self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = PYBOOST_SINGLE_OP_SOURCE_TEMPLATE
         self.gen_path = gen_path
         self.device = device
 
-    def generate_view_op_cpp_code(self, op_protos, merge_op_header, merge_op_function):
+    def generate_view_op_cpp_code(self, op_protos, merge_op_header, merge_op_function, ascend_merge_op_inc):
         """
         Generate C++ code for view operations in PyBoost.
 
@@ -327,6 +358,7 @@ class PyboostViewOpCppGenerator:
                                                                return_type=cpp_func_return,
                                                                call_impl=call_impl,
                                                                register_custom_kernel=""))
+            ascend_merge_op_inc.append(op_proto.op_class.name)
 
 
 class AclnnOpCppCodeGenerator:
@@ -366,16 +398,18 @@ class AclnnOpCppCodeGenerator:
             PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
         else:
-            raise ValueError(f"Device must be ascend, gpu, or cpu, {device} is not supported")
+            raise ValueError(
+                f"Device must be ascend, gpu, or cpu, {device} is not supported")
         self.PYBOOST_CALL_TEMPLATE = PYBOOST_CALL_TEMPLATE
         self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.Template(
             '#include "kernel/${device}/pyboost/auto_generate/${operator_name}.h"\n'
         )
+
         self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = PYBOOST_SINGLE_OP_SOURCE_TEMPLATE
         self.gen_path = gen_path
         self.device = device
 
-    def generate_aclnn_op_cpp_code(self, op_protos, merge_op_header, merge_op_function):
+    def generate_aclnn_op_cpp_code(self, op_protos, merge_op_header, merge_op_function, ascend_merge_op_inc):
         """
         Generate C++ code for ACLNN operations in PyBoost.
 
@@ -404,10 +438,12 @@ class AclnnOpCppCodeGenerator:
             aclnn_name = AclnnUtils.get_aclnn_interface(op_proto.op_class.name)
 
             call_args_tensor = op_parser.get_call_args_tensor()
-            create_input_address = self._generate_create_input_address(op_parser)
+            create_input_address = self._generate_create_input_address(
+                op_parser)
             malloc_inputs = self._generate_malloc_input(op_parser)
             op_outputs, call_func_outputs = op_parser.generate_pyboost_outputs()
-            get_inputs_kernel_tensors = self._generate_get_inputs_kernel_tensors(op_parser)
+            get_inputs_kernel_tensors = self._generate_get_inputs_kernel_tensors(
+                op_parser)
 
             cube_math_type, get_cube_math_type = '', ''
             if self.device == 'ascend' and is_cube(op_proto.op_class.name):
@@ -456,6 +492,7 @@ class AclnnOpCppCodeGenerator:
                                                                return_type=cpp_func_return,
                                                                call_impl=call_impl,
                                                                register_custom_kernel=''))
+            ascend_merge_op_inc.append(op_proto.op_class.name)
 
     def _generate_tensor_cpu_cast_input_code(self, op_parser: OpTemplateParser):
         """
@@ -573,6 +610,7 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         self.ascend_gen_path = f"{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/auto_generate/"
         self.cpu_gen_path = f"{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/auto_generate/"
         self.gpu_gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
+        self.hccl_gen_path = "mindspore/ccsrc/plugin/device/ascend/kernel/hccl/pyboost/auto_generate/"
 
     def generate(self, work_path, op_protos):
         """
@@ -607,27 +645,52 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         """
         ascend_merge_op_header = []
         ascend_merge_op_function = []
+        hccl_merge_op_header = []
+        hccl_merge_op_function = []
+        ascend_merge_op_inc = []
+        ascend_merge_op_hccl_inc = []
         self.ascend_op_cpp_generator.generate_customize_op_cpp_code(op_protos, ascend_merge_op_header,
-                                                                    ascend_merge_op_function)
+                                                                    ascend_merge_op_function, ascend_merge_op_inc,
+                                                                    hccl_merge_op_header, hccl_merge_op_function,
+                                                                    ascend_merge_op_hccl_inc)
         self.ascend_view_op_cpp_generator.generate_view_op_cpp_code(op_protos, ascend_merge_op_header,
-                                                                    ascend_merge_op_function)
+                                                                    ascend_merge_op_function,
+                                                                    ascend_merge_op_inc)
         self.ascend_aclnn_cpp_generator.generate_aclnn_op_cpp_code(op_protos, ascend_merge_op_header,
-                                                                   ascend_merge_op_function)
+                                                                   ascend_merge_op_function,
+                                                                   ascend_merge_op_inc)
 
-        ascend_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(ascend_merge_op_header, chunk_size=120)
-        ascend_op_function_merge_by_chunk_size = merge_strings_by_chunk_size(ascend_merge_op_function, chunk_size=120)
+        ascend_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(
+            ascend_merge_op_header, chunk_size=120)
+        ascend_op_function_merge_by_chunk_size = merge_strings_by_chunk_size(
+            ascend_merge_op_function, chunk_size=120)
+        op_inc_list = chunk_list(ascend_merge_op_inc, n=120)
 
         new_gen_num = len(ascend_op_header_merge_by_chunk_size)
-        self._delete_residual_merged_ops_files(os.path.join(work_path, self.ascend_gen_path), new_gen_num)
+        self._delete_residual_merged_ops_files(os.path.join(
+            work_path, self.ascend_gen_path), new_gen_num)
 
         for i, op_header, op_function in zip(range(len(ascend_op_header_merge_by_chunk_size)),
                                              ascend_op_header_merge_by_chunk_size,
                                              ascend_op_function_merge_by_chunk_size):
+            ops_inc_head_set = set()
+            for op_name_inc in op_inc_list[i]:
+                ops_inc_head_set.add(template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_name_inc[0].lower()))
+
             ascend_pyboost_op_source = self.PYBOOST_ASCEND_OP_SOURCE_TEMPLATE.replace(
-                merge_op_header=op_header, merge_op_function=op_function)
+                merge_op_header=op_header, merge_op_function=op_function, ops_inc=list(sorted(ops_inc_head_set)))
             save_file(os.path.join(work_path, self.ascend_gen_path), f"pyboost_ascend_ops_{i}.cc",
                       ascend_pyboost_op_source)
 
+        ops_hccl_inc_head_set = set()
+        for op_name_inc in ascend_merge_op_hccl_inc:
+            ops_hccl_inc_head_set.add(template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_name_inc[0].lower()))
+        hccl_pyboost_op_source = self.PYBOOST_ASCEND_OP_SOURCE_TEMPLATE.replace(
+            merge_op_header='\n'.join(hccl_merge_op_header), merge_op_function='\n'.join(hccl_merge_op_function),
+            ops_inc=list(sorted(ops_hccl_inc_head_set)))
+        save_file(os.path.join(work_path, self.hccl_gen_path), f"pyboost_hccl_ops.cc",
+                    hccl_pyboost_op_source)
+
     def _generate_pyboost_cpu_ops(self, work_path, op_protos):
         """
         Generates CPU PyBoost ops functions source files after being merged into specific chunk sizes.
@@ -639,20 +702,31 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         """
         cpu_merge_op_header = []
         cpu_merge_op_function = []
-        self.cpu_op_cpp_generator.generate_customize_op_cpp_code(op_protos, cpu_merge_op_header, cpu_merge_op_function)
-        self.cpu_view_op_cpp_generator.generate_view_op_cpp_code(op_protos, cpu_merge_op_header, cpu_merge_op_function)
-        self.cpu_aclnn_cpp_generator.generate_aclnn_op_cpp_code(op_protos, cpu_merge_op_header, cpu_merge_op_function)
-        cpu_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(cpu_merge_op_header, chunk_size=120)
-        cpu_op_function_merge_by_chunk_size = merge_strings_by_chunk_size(cpu_merge_op_function, chunk_size=120)
+        cpu_merge_op_inc = []
+        self.cpu_op_cpp_generator.generate_customize_op_cpp_code(
+            op_protos, cpu_merge_op_header, cpu_merge_op_function, cpu_merge_op_inc)
+        self.cpu_view_op_cpp_generator.generate_view_op_cpp_code(
+            op_protos, cpu_merge_op_header, cpu_merge_op_function, cpu_merge_op_inc)
+        self.cpu_aclnn_cpp_generator.generate_aclnn_op_cpp_code(
+            op_protos, cpu_merge_op_header, cpu_merge_op_function, cpu_merge_op_inc)
+        cpu_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(
+            cpu_merge_op_header, chunk_size=120)
+        cpu_op_function_merge_by_chunk_size = merge_strings_by_chunk_size(
+            cpu_merge_op_function, chunk_size=120)
+        op_inc_list = chunk_list(cpu_merge_op_inc, n=120)
 
         new_gen_num = len(cpu_op_header_merge_by_chunk_size)
-        self._delete_residual_merged_ops_files(os.path.join(work_path, self.cpu_gen_path), new_gen_num)
+        self._delete_residual_merged_ops_files(
+            os.path.join(work_path, self.cpu_gen_path), new_gen_num)
 
         for i, op_header, op_function in zip(range(len(cpu_op_header_merge_by_chunk_size)),
                                              cpu_op_header_merge_by_chunk_size,
                                              cpu_op_function_merge_by_chunk_size):
+            ops_inc_head_set = set()
+            for op_name_inc in op_inc_list[i]:
+                ops_inc_head_set.add(template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_name_inc[0].lower()))
             cpu_pyboost_op_source = self.PYBOOST_CPU_OP_SOURCE_TEMPLATE.replace(
-                merge_op_header=op_header, merge_op_function=op_function)
+                merge_op_header=op_header, merge_op_function=op_function, ops_inc=list(sorted(ops_inc_head_set)))
             save_file(os.path.join(work_path, self.cpu_gen_path), f"pyboost_cpu_ops_{i}.cc",
                       cpu_pyboost_op_source)
 
@@ -667,20 +741,31 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         """
         gpu_merge_op_header = []
         gpu_merge_op_function = []
-        self.gpu_op_cpp_generator.generate_customize_op_cpp_code(op_protos, gpu_merge_op_header, gpu_merge_op_function)
-        self.gpu_view_op_cpp_generator.generate_view_op_cpp_code(op_protos, gpu_merge_op_header, gpu_merge_op_function)
-        self.gpu_aclnn_cpp_generator.generate_aclnn_op_cpp_code(op_protos, gpu_merge_op_header, gpu_merge_op_function)
-        gpu_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(gpu_merge_op_header, chunk_size=120)
-        gpu_op_function_merge_by_chunk_size = merge_strings_by_chunk_size(gpu_merge_op_function, chunk_size=120)
+        gpu_merge_op_inc = []
+        self.gpu_op_cpp_generator.generate_customize_op_cpp_code(
+            op_protos, gpu_merge_op_header, gpu_merge_op_function, gpu_merge_op_inc)
+        self.gpu_view_op_cpp_generator.generate_view_op_cpp_code(
+            op_protos, gpu_merge_op_header, gpu_merge_op_function, gpu_merge_op_inc)
+        self.gpu_aclnn_cpp_generator.generate_aclnn_op_cpp_code(
+            op_protos, gpu_merge_op_header, gpu_merge_op_function, gpu_merge_op_inc)
+        gpu_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(
+            gpu_merge_op_header, chunk_size=120)
+        gpu_op_function_merge_by_chunk_size = merge_strings_by_chunk_size(
+            gpu_merge_op_function, chunk_size=120)
+        op_inc_list = chunk_list(gpu_merge_op_inc, n=120)
 
         new_gen_num = len(gpu_op_header_merge_by_chunk_size)
-        self._delete_residual_merged_ops_files(os.path.join(work_path, self.gpu_gen_path), new_gen_num)
+        self._delete_residual_merged_ops_files(
+            os.path.join(work_path, self.gpu_gen_path), new_gen_num)
 
         for i, op_header, op_function in zip(range(len(gpu_op_header_merge_by_chunk_size)),
                                              gpu_op_header_merge_by_chunk_size,
                                              gpu_op_function_merge_by_chunk_size):
+            ops_inc_head_set = set()
+            for op_name_inc in op_inc_list[i]:
+                ops_inc_head_set.add(template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_name_inc[0].lower()))
             gpu_pyboost_op_source = self.PYBOOST_GPU_OP_SOURCE_TEMPLATE.replace(
-                merge_op_header=op_header, merge_op_function=op_function)
+                merge_op_header=op_header, merge_op_function=op_function, ops_inc=list(sorted(ops_inc_head_set)))
             save_file(os.path.join(work_path, self.gpu_gen_path), f"pyboost_gpu_ops_{i}.cc",
                       gpu_pyboost_op_source)
 
@@ -703,7 +788,8 @@ class PyboostOpFunctionGenerator(BaseGenerator):
             None
         """
         all_files = os.listdir(files_path)
-        old_pyboost_ops_files = [file for file in all_files if re.match(r'pyboost_.*_ops_.*\.cc', file)]
+        old_pyboost_ops_files = [file for file in all_files if re.match(
+            r'pyboost_.*_ops_.*\.cc', file)]
         old_files_num = len(old_pyboost_ops_files)
         if new_gen_num != old_files_num:
             for file in old_pyboost_ops_files:
@@ -782,7 +868,8 @@ def delete_residual_files(work_path, op_protos):
         all_operator_name.append(op_proto.op_name)
     code_generate_path_list = [f"{K.MS_OPS_KERNEL_PATH}/{device}/pyboost/auto_generate/" for device in
                                ["ascend", "gpu", "cpu"]]
-    code_generate_path_list.append(f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
+    code_generate_path_list.append(
+        f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
     for code_generate_path in code_generate_path_list:
         filter_files = []
         code_generate_path = os.path.join(work_path, code_generate_path)
@@ -843,9 +930,9 @@ class PyboostOpRegisterCppCodeGenerator:
         for op_name in all_op_names:
             factory_str += "template class OpFactory<{0}>;\n".format(op_name)
         for operator_name in all_functional_names:
-            include_str += f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/{operator_name}.h"\n'
+            include_str += f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/{operator_name}.h"\n'
         op_register_file_str = self.PYBOOST_OP_REGISTER_TEMPLATE.replace(op_includes=include_str,
                                                                          op_factory_templates=factory_str)
-        save_path = os.path.join(work_path, f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
+        save_path = os.path.join(work_path, f"{K.MS_PYBOOST_BASE_PATH}/auto_generate/")
         file_name = "op_register.cc"
         save_file(save_path, file_name, op_register_file_str)

mindspore/ops_generate/{pyboost_overload_functions_cpp_generator.py → pyboost/pyboost_overload_functions_cpp_generator.py}

@@ -21,15 +21,17 @@ the necessary header files for the generated functions.
 """
 
 import os
-import template
-from template import Template
-import gen_constants as K
-import pyboost_utils
-import op_api_proto
-from gen_utils import save_file
-from op_proto import OpProto
-from op_template_parser import OpTemplateParser
-from base_generator import BaseGenerator
+
+import common.template as template
+import common.gen_constants as K
+from common.template import Template
+from common.gen_utils import save_file
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+from pyboost import pyboost_utils
+import api.op_api_proto as op_api_proto
+
+from .op_template_parser import OpTemplateParser
 
 
 class PyboostOverloadFunctionsGenerator(BaseGenerator):
@@ -68,7 +70,9 @@ class PyboostOverloadFunctionsGenerator(BaseGenerator):
         self.pyboost_return_template = Template(
             '${arg_handler_processor}\n'
             'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
-            'return mindspore::tensor::ToPython(${pyboost_base_func_name}_Base(${prim_name}, arg_list));\n'
+            'auto res = ${pyboost_base_func_name}_OP(${prim_name}, parse_args.src_types_, ${convert_args});\n'
+            'trace::Capture(parse_args.arg_list_, "${class_name}", &res);\n'
+            'return res;\n'
         )
         self.callback_python_template = Template(
             'MS_LOG(INFO) << "Callback python method: ${py_method}";\n'
@@ -112,8 +116,9 @@ class PyboostOverloadFunctionsGenerator(BaseGenerator):
         """
 
         mint_classes_def_list = []
-
-        _, single_mint_func_data, overload_mint_func_data = op_api_proto.categorize_func_data(mint_func_protos_data)
+        ops_inc_head_set = set()
+        _, single_mint_func_data, overload_mint_func_data, op_class_name_set = op_api_proto.categorize_func_data(
+            mint_func_protos_data)
         single_func_call_body_list, single_cpp_class_name_list = (
             self._get_single_func_call_body_list(single_mint_func_data))
         overload_func_call_body_list, overload_cpp_class_name_list = (
@@ -125,9 +130,11 @@ class PyboostOverloadFunctionsGenerator(BaseGenerator):
         cpp_class_name_list = single_cpp_class_name_list + overload_cpp_class_name_list
         mint_classes_reg_list = (
             self._get_mint_func_reg_list(single_mint_func_data, overload_mint_func_data, cpp_class_name_list))
-
+        for op_class_name in op_class_name_set:
+            ops_inc_head_set.add(template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_class_name[0].lower()))
         pyboost_overload_file_str = (
-            self.PYBOOST_OVERLOAD_FUNCTIONS_TEMPLATE.replace(mint_func_classes_def=mint_classes_def_list,
+            self.PYBOOST_OVERLOAD_FUNCTIONS_TEMPLATE.replace(ops_inc=list(sorted(ops_inc_head_set)),
+                                                             mint_func_classes_def=mint_classes_def_list,
                                                              pybind_register_code=mint_classes_reg_list))
         save_path = os.path.join(work_path, K.PIPELINE_PYBOOST_FUNC_GEN_PATH)
         file_name = "pyboost_overload_functions.cc"
@@ -306,13 +313,15 @@ class PyboostOverloadFunctionsGenerator(BaseGenerator):
         func_proto_device = getattr(func_proto, device)
         if func_proto_device == 'pyboost':
             arg_handler_processor_str = self._get_arg_handler_processor(func_proto.func_name, func_proto.op_proto)
+            convert_args_str = self._get_convert_args_str(func_proto.op_proto)
             op_parser = OpTemplateParser(func_proto.op_proto)
             op_pyboost_func_name = op_parser.get_pyboost_func_name()
             prim_name = f"prim::kPrim{func_proto.op_proto.op_class.name}"
             return self.pyboost_return_template.replace(arg_handler_processor=arg_handler_processor_str,
                                                         class_name=func_proto.op_proto.op_class.name,
                                                         prim_name=prim_name,
-                                                        pyboost_base_func_name=op_pyboost_func_name)
+                                                        pyboost_base_func_name=op_pyboost_func_name,
+                                                        convert_args=convert_args_str)
         if func_proto_device == 'py_method':
             return self.callback_python_template.replace(py_method=func_proto.py_method)
 
@@ -322,6 +331,10 @@ class PyboostOverloadFunctionsGenerator(BaseGenerator):
         op_parser = OpTemplateParser(op_proto)
         return op_parser.get_arg_handler_processor(func_name, op_proto, is_tensor_api=False)
 
+    def _get_convert_args_str(self, op_proto):
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.get_convert_args_str(op_proto, is_tensor_api=False)
+
     def _get_mint_func_reg_list(self, single_mint_func_data, overload_mint_func_data, cpp_class_names):
         """
         Generates the list of pybind definition strings for mint functions.

mindspore/ops_generate/{pyboost_utils.py → pyboost/pyboost_utils.py}

@@ -16,8 +16,9 @@
 
 import os
 import logging
-from gen_utils import safe_load_yaml
-from op_proto import OpProto
+from common.gen_utils import safe_load_yaml
+from common.op_proto import OpProto
+import common.gen_constants as K
 
 
 def is_optional_param(op_arg):
@@ -114,6 +115,54 @@ def get_convert_type_str(dtype: str, optional):
     raise TypeError(f"""Unsupported convert type {dtype} for args.""")
 
 
+def get_input_args_type_str(dtype: str, optional):
+    """
+    Convert type
+    """
+    # add more type here
+    native_type = {
+        'int': 'Int64ImmPtr',
+        'float': 'FP32ImmPtr',
+        'bool': 'BoolImmPtr',
+        'number': 'ScalarPtr',
+        'tuple[int]': 'ValueTuplePtr',
+        'tuple[float]': 'ValueTuplePtr',
+        'tuple[bool]': 'ValueTuplePtr',
+        'tuple[tensor]': 'ValueTuplePtr',
+        'list[int]': 'ValueTuplePtr',
+        'list[float]': 'ValueTuplePtr',
+        'list[bool]': 'ValueTuplePtr',
+        'list[tensor]': 'ValueTuplePtr',
+        'tensor': 'ValuePtr',
+        'str': 'StringImmPtr',
+        'type': 'Int64ImmPtr',
+    }
+    optional_type = {
+        'int': 'std::optional<Int64ImmPtr>',
+        'float': 'std::optional<FP32ImmPtr>',
+        'number': 'std::optional<ScalarPtr>',
+        'tensor': 'std::optional<ValuePtr>',
+        'type': 'std::optional<Int64ImmPtr>',
+        'str': 'std::optional<StringImmPtr>',
+        'tuple[int]': 'std::optional<ValueTuplePtr>',
+        'tuple[float]': 'std::optional<ValueTuplePtr>',
+        'tuple[bool]': 'std::optional<ValueTuplePtr>',
+        'tuple[tensor]': 'std::optional<ValueTuplePtr>',
+        'list[int]': 'std::optional<ValueTuplePtr>',
+        'list[float]': 'std::optional<ValueTuplePtr>',
+        'list[bool]': 'std::optional<ValueTuplePtr>',
+        'list[tensor]': 'std::optional<ValueTuplePtr>',
+    }
+    if optional:
+        if dtype in optional_type:
+            return optional_type[dtype]
+        raise TypeError(f"""Unknown optional type {dtype} for args.""")
+    if dtype in native_type:
+        return native_type[dtype]
+    raise TypeError(f"""Unknown type {dtype} for args.""")
+
+
+
 def get_value_convert_type_str(dtype: str, optional):
     """
     Convert type
@@ -407,12 +456,47 @@ def merge_strings_by_chunk_size(string_list, chunk_size=50):
     return merged_strings
 
 
+def chunk_list(lst, n):
+    """
+    Divide a list into sublists of length 'n'.
+
+    This function takes a list `lst` and an integer `n`, and returns a new list
+    where each element is a sublist of `lst` containing up to `n` elements.
+    If the length of `lst` is not a multiple of `n`, the last sublist will contain
+    fewer than `n` elements.
+
+    Args:
+        lst (list): The original list to be divided.
+        n (int): The number of elements per sublist.
+
+    Returns:
+        list: A list of sublists, where each sublist has up to `n` elements.
+
+    Raises:
+        ValueError: If `n` is less than or equal to 0.
+
+    Example:
+        >>> my_list = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+        >>> chunked_list = chunk_list(my_list, 3)
+        >>> print(chunked_list)
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+
+    Note:
+        This function assumes that `n` is a positive integer. If `n` is not a
+        positive integer, a ValueError is raised.
+    """
+    if n <= 0:
+        raise ValueError("The chunk size 'n' must be a positive integer.")
+
+    return [lst[i:i + n] for i in range(0, len(lst), n)]
+
+
 class AclnnUtils:
     """
     aclnn utils
     """
-    work_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../../../")
-    aclnn_map = safe_load_yaml(os.path.join(work_path, "./mindspore/python/mindspore/ops_generate/aclnn_config.yaml"))
+    aclnn_map = safe_load_yaml(os.path.join(
+        K.WORK_DIR, K.PY_OPS_GEN_PATH, "pyboost/aclnn_config.yaml"))
 
     @staticmethod
     def get_aclnn_interface(class_name):

mindspore/ops_generate/resources/resource_list.py

@@ -0,0 +1,30 @@
+# Copyright 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Resource list."""
+
+from enum import Enum
+
+
+class ResourceType(Enum):
+    OP_PROTO = 0
+    DEPRECATED_OP_PROTO = 1
+    OP_DOC_YAML = 2
+    TENSOR_METHOD_DOC_YAML = 3
+    MINT_FUNC_DOC_YAML = 4
+    TENSOR_METHOD_PROTOS = 5
+    MINT_FUNC_PROTOS = 6
+    ALIAS_API_MAPPING = 7
+    FUNC_OP_PROTO = 8