mindspore 2.6.0rc1__cp310-cp310-win_amd64.whl → 2.7.0rc1__cp310-cp310-win_amd64.whl

mindspore/ops_generate/pyboost/pyboost_op_cpp_code_generator.py

@@ -23,6 +23,7 @@ import re
 
 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
@@ -32,6 +33,13 @@ from .pyboost_utils import is_cube, AclnnUtils, get_return_type, merge_strings_b
 from .op_template_parser import OpTemplateParser
 
 
+def check_no_basic_int_type(op_args):
+    for arg in op_args:
+        if arg.arg_dtype in ["tuple[int]", "list[int]", "int"]:
+            return False
+    return True
+
+
 class PyboostCommonOpHeaderGenerator(BaseGenerator):
     """
     Generates common C++ headers for PyBoost operations.
@@ -42,6 +50,11 @@ class PyboostCommonOpHeaderGenerator(BaseGenerator):
 
     def __init__(self):
         self.pyboost_op_header_str = template.PYBOOST_BASE_OP_DEFINE_TEMPLATE
+        self.pyboost_basic_type_func_template = Template(
+            'virtual ${return_type} Call(${call_args_with_type}) {' \
+            '\n   MS_EXCEPTION(NotImplementedError) << "Basic type func not implemented";' \
+            '\n};'
+        )
 
     def generate(self, work_path, op_protos):
         """
@@ -59,7 +72,10 @@ class PyboostCommonOpHeaderGenerator(BaseGenerator):
                 continue
             op_parser = OpTemplateParser(op_proto)
             op_name_str = op_proto.op_class.name
-            call_args_with_type = op_parser.parse_call_args_with_types()
+            if op_proto.op_view and not check_no_basic_int_type(op_proto.op_args):
+                call_args_with_type = op_parser.parse_call_args_with_types(basic_type=True)
+            else:
+                call_args_with_type = op_parser.parse_call_args_with_types(basic_type=False)
             cpp_func_return = _generate_cpp_func_return(op_proto)
             output_is_tuple = "bool output_is_tuple() const override { return true; }" \
                 if is_op_multi_output(op_proto.op_returns) else ''
@@ -121,16 +137,77 @@ class PyboostOpHeaderGenerator(BaseGenerator):
             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()
+            if op_proto.op_view and not check_no_basic_int_type(op_proto.op_args):
+                call_args_with_type = op_parser.parse_call_args_with_types(basic_type=True)
+            else:
+                call_args_with_type = op_parser.parse_call_args_with_types(basic_type=False)
             cpp_func_return = _generate_cpp_func_return(op_proto)
-
             pyboost_op_str = self.PYBOOST_OP_HEADER_TEMPLATE.replace(op_name=op_name_str,
                                                                      op_name_upper=op_name_str.upper(),
                                                                      operator_name=op_proto.op_name,
                                                                      call_args_with_type=call_args_with_type,
                                                                      return_type=cpp_func_return)
+            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)
+
+
+class PyboostInternalOpHeaderGenerator(BaseGenerator):
+    """
+    Generates C++ headers for PyBoost internal operations.
 
-            save_path = os.path.join(work_path, self.code_generate_path if not is_ascend_comm_op else self.hccl_code_generate_path)
+    This class generates header files for Ascend and defines the operation functions accordingly.
+    """
+
+    def __init__(self, device):
+        """
+        Initializes the PyboostOpHeaderGenerator with the appropriate templates for the specified device.
+
+        Args:
+            device (str): The target device (ascend, gpu, or cpu), currently only support ascend.
+
+        Raises:
+            ValueError: If the device is not supported.
+        """
+        if device != 'ascend':
+            raise ValueError(
+                f"Currently, only support 'ascend' for internal operations, {device} is not supported.")
+        self.pyboost_internal_op_header_template = template.PYBOOST_ASCEND_INTERNAL_OP_HEADER_TEMPLATE
+        self.code_generate_path = f"{K.MS_OPS_KERNEL_PATH}/{device}/pyboost/internal/auto_generate/"
+        self.device = device
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates header files for the provided operator prototypes based on the device.
+
+        Args:
+            work_path (str): The directory path where the header files will be saved.
+            op_protos (list): A list of operator prototypes containing information about the operators.
+
+        Returns:
+            None
+        """
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None:
+                continue
+            if getattr(op_proto.op_dispatch, self.device) == 'None':
+                continue
+            if getattr(op_proto.op_dispatch, 'internal_op_ascend') == 'None':
+                continue
+            op_parser = OpTemplateParser(op_proto)
+            op_name_str = op_proto.op_class.name
+            call_args_with_type = op_parser.parse_call_args_with_types()
+            cpp_func_return = _generate_cpp_func_return(op_proto)
+
+            pyboost_op_str = self.pyboost_internal_op_header_template.replace(
+                op_name=op_name_str,
+                op_name_upper=op_name_str.upper(),
+                operator_name=op_proto.op_name,
+                call_args_with_type=call_args_with_type,
+                return_type=cpp_func_return)
+
+            save_path = os.path.join(work_path, self.code_generate_path)
             file_name = f"{op_proto.op_name}.h"
             save_file(save_path, file_name, pyboost_op_str)
 
@@ -158,19 +235,24 @@ class PyboostOpCppGenerator:
             PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_HEADER_TEMPLATE
             PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/auto_generate/"
+            self.device_reg_str = "Ascend"
         elif device == 'cpu':
             PYBOOST_CUSTOMIZE_CALL_TEMPLATE = template.PYBOOST_CPU_CUSTOMIZE_CALL_TEMPLATE
             PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.PYBOOST_CPU_SINGLE_OP_HEADER_TEMPLATE
-            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_CPU_SINGLE_OP_SOURCE_TEMPLATE
+            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/auto_generate/"
+            self.device_reg_str = "CPU"
         elif device == 'gpu':
             PYBOOST_CUSTOMIZE_CALL_TEMPLATE = template.PYBOOST_GPU_CUSTOMIZE_CALL_TEMPLATE
             PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_HEADER_TEMPLATE
-            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_SOURCE_TEMPLATE
+            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
+            self.device_reg_str = "GPU"
         else:
             raise ValueError(
                 f"Device must be ascend, gpu, or cpu, {device} is not supported")
+        self.PYBOOST_REG_OP_TEMPLATE = Template('MS_REG_PYBOOST_OP(${device}, ${op_name});' \
+                                                '${register_custom_kernel}')
         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
@@ -201,8 +283,11 @@ class PyboostOpCppGenerator:
                 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()
+            call_args = OpTemplateParser.parse_original_call_args(op_proto.op_args)
+            if op_proto.op_view and not check_no_basic_int_type(op_proto.op_args):
+                call_args_with_type = op_parser.parse_call_args_with_types(True)
+            else:
+                call_args_with_type = op_parser.parse_call_args_with_types()
             _, call_func_outputs = op_parser.generate_pyboost_outputs()
             operator_name = op_proto.op_name
             op_name_str = op_proto.op_class.name
@@ -219,7 +304,8 @@ class PyboostOpCppGenerator:
                 check_expression=check_inplace_func,
             )
             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}")
+                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 = \
@@ -230,23 +316,30 @@ class PyboostOpCppGenerator:
 
             register_custom = self._get_register_custom_kernel(op_proto)
             cpp_func_return = _generate_cpp_func_return(op_proto)
+            op_register = self.PYBOOST_REG_OP_TEMPLATE.replace(op_name=op_name_str,
+                                                               device=self.device_reg_str,
+                                                               register_custom_kernel=register_custom)
             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_hccl_header.append(
+                    self.PYBOOST_SINGLE_HCLL_OP_HEADER_TEMPLATE.replace(operator_name=operator_name,
+                                                                        customize_include=customize_include))
                 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))
+                                                                   call_args_with_type=call_args_with_type,
+                                                                   return_type=cpp_func_return, call_impl=call_impl,
+                                                                   op_register=op_register,
+                                                                   device=self.device_reg_str))
                 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_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))
+                                                                   call_args_with_type=call_args_with_type,
+                                                                   return_type=cpp_func_return, call_impl=call_impl,
+                                                                   op_register=op_register,
+                                                                   device=self.device_reg_str))
                 merge_op_inc.append(op_name_str)
 
     def _get_register_custom_kernel(self, op_proto: OpProto):
@@ -294,19 +387,24 @@ class PyboostViewOpCppGenerator:
             PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_HEADER_TEMPLATE
             PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/auto_generate/"
+            self.device_reg_str = "Ascend"
         elif device == 'cpu':
             PYBOOST_VIEW_CALL_TEMPLATE = template.PYBOOST_CPU_VIEW_CALL_TEMPLATE
             PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.PYBOOST_CPU_SINGLE_OP_HEADER_TEMPLATE
-            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_CPU_SINGLE_OP_SOURCE_TEMPLATE
+            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/auto_generate/"
+            self.device_reg_str = "CPU"
         elif device == 'gpu':
             PYBOOST_VIEW_CALL_TEMPLATE = template.PYBOOST_GPU_VIEW_CALL_TEMPLATE
             PYBOOST_SINGLE_OP_HEADER_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_HEADER_TEMPLATE
-            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_SOURCE_TEMPLATE
+            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
+            self.device_reg_str = "GPU"
         else:
             raise ValueError(
                 f"Device must be ascend, gpu, or cpu, {device} is not supported")
+        self.PYBOOST_REG_OP_TEMPLATE = Template('MS_REG_PYBOOST_OP(${device}, ${op_name});' \
+                                                '${register_custom_kernel}')
         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
@@ -327,6 +425,7 @@ class PyboostViewOpCppGenerator:
             merge_op_header (list): A list to store the generated C++ header code for view operations.
             merge_op_function (list): A list to store the generated C++ source code for view operations.
         """
+        calc_args_temp = Template("{${call_args}}")
         for op_proto in op_protos:
             if op_proto.op_dispatch is None:
                 continue
@@ -339,10 +438,19 @@ class PyboostViewOpCppGenerator:
 
             op_parser = OpTemplateParser(op_proto)
             call_args_tensor = op_parser.get_call_args_tensor()
-            call_args = op_parser.parse_original_call_args(op_proto.op_args)
-            call_args_with_type = op_parser.parse_call_args_with_types()
+            call_args = OpTemplateParser.parse_original_call_args(op_proto.op_args)
+            if op_proto.op_view and not check_no_basic_int_type(op_proto.op_args):
+                call_args_with_type = op_parser.parse_call_args_with_types(True)
+                storage_calc_str = op_proto.op_class.name + "BasicType"
+                calc_func_args_str = call_args
+            else:
+                call_args_with_type = op_parser.parse_call_args_with_types()
+                storage_calc_str = op_proto.op_class.name
+                calc_func_args_str = calc_args_temp.replace(call_args=call_args)
             _, call_func_outputs = op_parser.generate_pyboost_outputs()
             call_impl = self.PYBOOST_VIEW_CALL_TEMPLATE.replace(op_name=op_proto.op_class.name,
+                                                                storage_calc=storage_calc_str,
+                                                                calc_func_args=calc_func_args_str,
                                                                 call_args=call_args,
                                                                 call_tensors=call_args_tensor,
                                                                 return_values=call_func_outputs,
@@ -351,13 +459,16 @@ class PyboostViewOpCppGenerator:
             cpp_func_return = _generate_cpp_func_return(op_proto)
             merge_op_header.append(self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE.replace(operator_name=op_proto.op_name,
                                                                                   customize_include=customize_include))
-
+            op_register = self.PYBOOST_REG_OP_TEMPLATE.replace(op_name=op_proto.op_class.name,
+                                                               device=self.device_reg_str,
+                                                               register_custom_kernel="")
             merge_op_function.append(
                 self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE.replace(op_name=op_proto.op_class.name,
                                                                call_args_with_type=call_args_with_type,
                                                                return_type=cpp_func_return,
                                                                call_impl=call_impl,
-                                                               register_custom_kernel=""))
+                                                               op_register=op_register,
+                                                               device=self.device_reg_str))
             ascend_merge_op_inc.append(op_proto.op_class.name)
 
 
@@ -389,17 +500,22 @@ class AclnnOpCppCodeGenerator:
             PYBOOST_CALL_TEMPLATE = template.PYBOOST_ASCEND_CALL_TEMPLATE
             PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/auto_generate/"
+            self.device_reg_str = "Ascend"
         elif device == 'cpu':
             PYBOOST_CALL_TEMPLATE = template.PYBOOST_CPU_CALL_TEMPLATE
-            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_CPU_SINGLE_OP_SOURCE_TEMPLATE
+            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/auto_generate/"
+            self.device_reg_str = "CPU"
         elif device == 'gpu':
             PYBOOST_CALL_TEMPLATE = template.PYBOOST_GPU_CALL_TEMPLATE
-            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_GPU_SINGLE_OP_SOURCE_TEMPLATE
+            PYBOOST_SINGLE_OP_SOURCE_TEMPLATE = template.PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE
             gen_path = f"{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/auto_generate/"
+            self.device_reg_str = "GPU"
         else:
             raise ValueError(
                 f"Device must be ascend, gpu, or cpu, {device} is not supported")
+        self.PYBOOST_REG_OP_TEMPLATE = Template('MS_REG_PYBOOST_OP(${device}, ${op_name});' \
+                                                '${register_custom_kernel}')
         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'
@@ -438,11 +554,11 @@ 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(
+            create_input_address = AclnnOpCppCodeGenerator._generate_create_input_address(
                 op_parser)
-            malloc_inputs = self._generate_malloc_input(op_parser)
+            malloc_inputs = AclnnOpCppCodeGenerator._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(
+            get_inputs_kernel_tensors = AclnnOpCppCodeGenerator._generate_get_inputs_kernel_tensors(
                 op_parser)
 
             cube_math_type, get_cube_math_type = '', ''
@@ -454,13 +570,16 @@ class AclnnOpCppCodeGenerator:
             real_output = ', ' + op_outputs \
                 if _generate_inplace_process_cpp_code(op_proto) == '' else ''
 
-            cast_input_code, real_call_args_tensor = self._generate_tensor_cpu_cast_input_code(
+            cast_input_code, real_call_args_tensor = AclnnOpCppCodeGenerator._generate_tensor_cpu_cast_input_code(
                 op_parser)
             cpp_func_return = _generate_cpp_func_return(op_proto)
             _, tensor_list_convert, call_args_with_tensor = op_parser.parse_need_malloc_tensors()
             call_args_after_convert, value_tuple_convert, const_number_convert = op_parser.op_args_converter()
-            call_args = op_parser.parse_original_call_args(op_proto.op_args)
-            call_args_with_type = op_parser.parse_call_args_with_types()
+            call_args = OpTemplateParser.parse_original_call_args(op_proto.op_args)
+            if op_proto.op_view and not check_no_basic_int_type(op_proto.op_args):
+                call_args_with_type = op_parser.parse_call_args_with_types(True)
+            else:
+                call_args_with_type = op_parser.parse_call_args_with_types()
             inplace_process = _generate_inplace_process_cpp_code(op_proto)
             call_impl = self.PYBOOST_CALL_TEMPLATE.replace(aclnn_name=aclnn_name,
                                                            call_args=call_args,
@@ -485,16 +604,20 @@ class AclnnOpCppCodeGenerator:
 
             merge_op_header.append(self.PYBOOST_SINGLE_OP_HEADER_TEMPLATE.replace(operator_name=op_proto.op_name,
                                                                                   device=self.device))
-
+            op_register = self.PYBOOST_REG_OP_TEMPLATE.replace(op_name=op_proto.op_class.name,
+                                                               device=self.device_reg_str,
+                                                               register_custom_kernel="")
             merge_op_function.append(
                 self.PYBOOST_SINGLE_OP_SOURCE_TEMPLATE.replace(op_name=op_proto.op_class.name,
                                                                call_args_with_type=call_args_with_type,
                                                                return_type=cpp_func_return,
                                                                call_impl=call_impl,
-                                                               register_custom_kernel=''))
+                                                               op_register=op_register,
+                                                               device=self.device_reg_str))
             ascend_merge_op_inc.append(op_proto.op_class.name)
 
-    def _generate_tensor_cpu_cast_input_code(self, op_parser: OpTemplateParser):
+    @staticmethod
+    def _generate_tensor_cpu_cast_input_code(op_parser: OpTemplateParser):
         """
         Generates the input casting code for CPU tensor operations.
 
@@ -524,18 +647,17 @@ class AclnnOpCppCodeGenerator:
             cast_input = "auto &select_kernel = kernel_attr_pair.second;\n" + cast_input
         return cast_input, real_call_args_tensor
 
-    def _generate_create_input_address(self, op_parser: OpTemplateParser):
+    @staticmethod
+    def _generate_create_input_address(op_parser: OpTemplateParser):
         need_malloc_tensors, _, _ = op_parser.parse_need_malloc_tensors()
         create_input_address = ''
-        args_list = ''
-        for item in need_malloc_tensors:
-            args_list += f'{item}, '
-        args_list = args_list[:-2]
+        args_list = ', '.join(str(item) for item in need_malloc_tensors)
         if args_list:
             create_input_address = f'PyBoostUtils::PrepareOpInputs(device_context_, op->stream_id(), {args_list});\n'
         return create_input_address
 
-    def _generate_malloc_input(self, op_parser: OpTemplateParser):
+    @staticmethod
+    def _generate_malloc_input(op_parser: OpTemplateParser):
         """
         Generates the code for creating input addresses for tensors that need to be allocated.
 
@@ -547,15 +669,13 @@ class AclnnOpCppCodeGenerator:
         """
         need_malloc_tensors, _, _ = op_parser.parse_need_malloc_tensors()
         malloc_inputs = ''
-        args_list = ''
-        for item in need_malloc_tensors:
-            args_list += f'{item}, '
-        args_list = args_list[:-2]
+        args_list = ', '.join(str(item) for item in need_malloc_tensors)
         if args_list:
             malloc_inputs += f'PyBoostUtils::MallocOpInputs(device_context, {args_list});\n'
         return malloc_inputs
 
-    def _generate_get_inputs_kernel_tensors(self, op_parser: OpTemplateParser):
+    @staticmethod
+    def _generate_get_inputs_kernel_tensors(op_parser: OpTemplateParser):
         """
         Generates the code for retrieving input kernel tensors.
 
@@ -567,16 +687,178 @@ class AclnnOpCppCodeGenerator:
         """
         _, _, call_args_with_tensor = op_parser.parse_need_malloc_tensors()
         inputs_kernel_tensors = ''
-        args_list = ''
-        for item in call_args_with_tensor:
-            args_list += f'{item}, '
-        args_list = args_list[:-2]
+        args_list = ', '.join(str(item) for item in call_args_with_tensor)
         if args_list:
             inputs_kernel_tensors += f'const auto &input_address_info = PyBoostUtils::GetAddressInfo(' \
                                      f'device_context, op->stream_id(), op->input_abs(), {args_list});\n'
         return inputs_kernel_tensors
 
 
+class InternalOpCppCodeGenerator:
+    """
+    Generates C++ code files for internal operations in PyBoost.
+    """
+
+    def __init__(self, device):
+        """
+        Initializes the InternalOpCppCodeGenerator with the appropriate templates.
+        """
+        self.device = device
+        self.internal_op_header_template = template.PYBOOST_INTERNAL_OP_HEADER_TEMPLATE
+        self.internal_single_op_header_template = template.PYBOOST_INTERNAL_SINGLE_OP_HEADER_TEMPLATE
+        self.internal_op_source_template = template.PYBOOST_INTERNAL_OP_SOURCE_TEMPLATE
+        self.internal_single_op_source_template = template.PYBOOST_INTERNAL_SINGLE_OP_SOURCE_TEMPLATE
+        self.internal_single_op_customize_source_template = template.PYBOOST_INTERNAL_SINGLE_OP_CUSTOMIZE_TEMPLATE
+        self.customize_inc_template = Template(
+            '#include "{ms_ops_kernel_path}/ascend/pyboost/internal/customize/${operator_name}.h"\n'
+        )
+        self.gen_path = f"{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/internal/auto_generate/"
+
+    def generate_internal_op_cpp_code(self, work_path, op_protos):
+        """
+        Generate internal op cpp code in pyboost.
+        """
+        merge_op_header = []
+        merge_op_function = []
+        ascend_merge_op_inc = []
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None or not op_proto.op_dispatch.enable:
+                continue
+            if getattr(op_proto.op_dispatch, 'internal_op_ascend') == 'None':
+                continue
+            internal_op_ascend = op_proto.op_dispatch.internal_op_ascend
+            op_name = op_proto.op_class.name
+            if internal_op_ascend == 'AutoGen':
+                self.generate_default_call(work_path, op_proto, merge_op_header,
+                                           merge_op_function, ascend_merge_op_inc)
+            elif internal_op_ascend == 'Internal' + op_name + 'AscendCustomize':
+                self.generate_customize_call(work_path, op_proto, merge_op_header,
+                                             merge_op_function, ascend_merge_op_inc)
+
+        if not ascend_merge_op_inc:
+            return
+        ops_inc_head_set = set()
+        for op_name_inc in ascend_merge_op_inc:
+            ops_inc_head_set.add(template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_name_inc[0].lower()))
+
+        internal_op_source_str = self.internal_op_source_template.replace(ops_prim_inc=list(sorted(ops_inc_head_set)),
+                                                                          merge_op_header=merge_op_header,
+                                                                          merge_op_function=merge_op_function)
+        save_path = os.path.join(work_path, self.gen_path)
+        file_name = "pyboost_ascend_internal_ops.cc"
+        save_file(save_path, file_name, internal_op_source_str)
+
+    def generate_default_op_function(self, op_parser, op_proto):
+        """
+        Generate default op call function.
+        """
+        call_args_with_type = op_parser.parse_call_args_with_types()
+        cpp_func_return = _generate_cpp_func_return(op_proto)
+        _, call_func_outputs = op_parser.generate_pyboost_outputs()
+        call_args = OpTemplateParser.parse_original_call_args(op_proto.op_args)
+        call_args_after_convert, value_tuple_convert, const_number_convert = op_parser.op_args_converter()
+        create_input_address, create_output_address = self._create_input_and_output_address(op_parser, op_proto)
+        internal_op_source_str = self.internal_single_op_source_template.replace(
+            op_name=op_proto.op_class.name,
+            operator_name=op_proto.op_name,
+            call_args_with_type=call_args_with_type,
+            internal_call_args=call_args,
+            internal_real_call_args=call_args_after_convert,
+            create_input_address=create_input_address,
+            create_output_address=create_output_address,
+            value_tuple_convert=value_tuple_convert,
+            const_number_convert=const_number_convert,
+            return_type=cpp_func_return,
+            return_values=call_func_outputs)
+        return internal_op_source_str
+
+    def generate_default_call(self, work_path, op_proto, merge_op_header,
+                              merge_op_function, ascend_merge_op_inc):
+        """
+        Generate internal op default call function in pyboost.
+        """
+        op_parser = OpTemplateParser(op_proto)
+        call_args_with_type = op_parser.parse_call_args_with_types()
+        cpp_func_return = _generate_cpp_func_return(op_proto)
+
+        # generate op header
+        internal_op_header_str = self.internal_op_header_template.replace(
+            operator_name=op_proto.op_name,
+            op_name=op_proto.op_class.name,
+            op_name_upper=op_proto.op_class.name.upper(),
+            call_args_with_type=call_args_with_type,
+            return_type=cpp_func_return)
+        save_path = os.path.join(work_path, self.gen_path)
+        save_file(save_path, f"{op_proto.op_name}.h", internal_op_header_str)
+        merge_op_header.append(
+            self.internal_single_op_header_template.replace(
+                operator_name=op_proto.op_name,
+                customize_inc=''))
+
+        # generate op function
+        internal_op_source_str = self.generate_default_op_function(op_parser, op_proto)
+        merge_op_function.append(internal_op_source_str)
+
+        ascend_merge_op_inc.append(op_proto.op_class.name)
+
+    def generate_customize_call(self, work_path, op_proto, merge_op_header,
+                                merge_op_function, ascend_merge_op_inc):
+        """
+        Generate internal op customize call function in pyboost.
+        """
+        op_parser = OpTemplateParser(op_proto)
+        call_args_with_type = op_parser.parse_call_args_with_types()
+        cpp_func_return = _generate_cpp_func_return(op_proto)
+
+        # generate op header
+        internal_op_header_str = self.internal_op_header_template.replace(
+            operator_name=op_proto.op_name,
+            op_name=op_proto.op_class.name,
+            op_name_upper=op_proto.op_class.name.upper(),
+            call_args_with_type=call_args_with_type,
+            return_type=cpp_func_return)
+        save_path = os.path.join(work_path, self.gen_path)
+        save_file(save_path, f"{op_proto.op_name}.h", internal_op_header_str)
+        self.customize_inc_template.replace(
+            ms_ops_kernel_path=K.MS_OPS_KERNEL_PATH,
+            operator_name=op_proto.op_name)
+        merge_op_header.append(
+            self.internal_single_op_header_template.replace(
+                operator_name=op_proto.op_name,
+                customize_inc=self.customize_inc_template))
+
+        # generate op function
+        _, call_func_outputs = op_parser.generate_pyboost_outputs()
+        call_args = OpTemplateParser.parse_original_call_args(op_proto.op_args)
+        internal_op_source_str = self.internal_single_op_customize_source_template.replace(
+            op_name=op_proto.op_class.name,
+            call_args=call_args,
+            call_args_with_type=call_args_with_type,
+            return_type=cpp_func_return,
+            return_values=call_func_outputs)
+        merge_op_function.append(internal_op_source_str)
+        ascend_merge_op_inc.append(op_proto.op_class.name)
+
+    @staticmethod
+    def _create_input_and_output_address(op_parser: OpTemplateParser, op_proto):
+        """
+        Create input and output address.
+        """
+        need_malloc_tensors, _, _ = op_parser.parse_need_malloc_tensors()
+        create_input_address = ''
+        create_output_address = ''
+        args_list = ''.join(f'{arg}, ' for arg in need_malloc_tensors)
+        args_list = args_list[:-2]
+        if args_list:
+            create_input_address = f'PyBoostUtils::PrepareOpInputs(device_context_, op->stream_id(), {args_list});\n'
+        if op_proto.op_args_signature and op_proto.op_args_signature.rw_write:
+            create_output_address = ''
+        else:
+            create_output_address = 'PyBoostUtils::PrepareOpOutputs(device_context_, op->stream_id(), outputs_);\n'
+
+        return create_input_address, create_output_address
+
+
 class PyboostOpFunctionGenerator(BaseGenerator):
     """
     Generates C++ source files for ACLNN operations in PyBoost.
@@ -595,6 +877,7 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         self.ascend_op_cpp_generator = PyboostOpCppGenerator('ascend')
         self.ascend_view_op_cpp_generator = PyboostViewOpCppGenerator('ascend')
         self.ascend_aclnn_cpp_generator = AclnnOpCppCodeGenerator('ascend')
+        self.ascend_internal_op_cpp_generator = InternalOpCppCodeGenerator('ascend')
 
         self.cpu_op_cpp_generator = PyboostOpCppGenerator('cpu')
         self.cpu_view_op_cpp_generator = PyboostViewOpCppGenerator('cpu')
@@ -659,6 +942,7 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         self.ascend_aclnn_cpp_generator.generate_aclnn_op_cpp_code(op_protos, ascend_merge_op_header,
                                                                    ascend_merge_op_function,
                                                                    ascend_merge_op_inc)
+        self.ascend_internal_op_cpp_generator.generate_internal_op_cpp_code(work_path, op_protos)
 
         ascend_op_header_merge_by_chunk_size = merge_strings_by_chunk_size(
             ascend_merge_op_header, chunk_size=120)
@@ -688,8 +972,8 @@ class PyboostOpFunctionGenerator(BaseGenerator):
         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)
+        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):
         """
@@ -788,8 +1072,7 @@ 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:
@@ -810,8 +1093,8 @@ def _generate_cpp_func_return(op_proto):
     """
     returns_type = []
     type_convert_to_base = {
-        'std::vector<mindspore::tensor::TensorPtr>': 'std::vector<mindspore::tensor::BaseTensorPtr>',
-        'mindspore::tensor::TensorPtr': 'mindspore::tensor::BaseTensorPtr'
+        'std::vector<mindspore::tensor::TensorPtr>': 'std::vector<mindspore::tensor::TensorPtr>',
+        'mindspore::tensor::TensorPtr': 'mindspore::tensor::TensorPtr'
     }
     for return_obj in op_proto.op_returns:
         temp_return = get_return_type(return_obj.arg_dtype)
@@ -866,8 +1149,8 @@ def delete_residual_files(work_path, op_protos):
     all_operator_name = []
     for op_proto in 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"]]
+    devices = ["ascend", "gpu", "cpu"]
+    code_generate_path_list = [f"{K.MS_OPS_KERNEL_PATH}/{device}/pyboost/auto_generate/" for device in devices]
     code_generate_path_list.append(
         f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
     for code_generate_path in code_generate_path_list:
@@ -916,19 +1199,23 @@ class PyboostOpRegisterCppCodeGenerator:
             None
         """
         all_op_names = []
+        internal_op_names = []
         all_functional_names = []
         for op_proto in op_protos:
             if op_proto.op_dispatch is None:
                 continue
-            functional_name = op_proto.op_name
             op_name_str = op_proto.op_class.name
+            if getattr(op_proto.op_dispatch, 'internal_op_ascend') != 'None':
+                internal_op_names.append(op_name_str)
             all_op_names.append(op_name_str)
-            all_functional_names.append(functional_name)
+            all_functional_names.append(op_proto.op_name)
 
         include_str = ''
         factory_str = ''
         for op_name in all_op_names:
             factory_str += "template class OpFactory<{0}>;\n".format(op_name)
+        for op_name in internal_op_names:
+            factory_str += "template class InternalOpFactory<{0}>;\n".format(op_name)
         for operator_name in all_functional_names:
             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,