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

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

@@ -22,12 +22,14 @@ function signatures, and generates conversion stubs for PyBoost operations.
 
 import re
 
-import pyboost_utils
-from pyboost_utils import get_input_dtype, tuple_input_to_cpp_type, get_return_type, \
-    number_input_to_cpp_type, get_const_number_convert, get_tuple_input_convert, is_optional_param
-import gen_constants as K
-from template import Template
-from op_proto import OpProto
+import common.gen_constants as K
+from common.template import Template
+from common.op_proto import OpProto
+
+from . import pyboost_utils
+from .pyboost_utils import get_input_dtype, tuple_input_to_cpp_type, get_return_type, \
+    number_input_to_cpp_type, get_const_number_convert, get_tuple_input_convert, is_optional_param, \
+    get_input_args_type_str
 
 
 class OpTemplateParser:
@@ -57,7 +59,8 @@ class OpTemplateParser:
         call_args_types = []
         for op_arg in op_args:
             is_optional = is_optional_param(op_arg)
-            call_args_types.append(get_input_dtype(op_arg.arg_dtype, is_optional))
+            call_args_types.append(get_input_dtype(
+                op_arg.arg_dtype, is_optional))
         return call_args_types
 
     def parse_call_args_with_types(self):
@@ -95,7 +98,8 @@ class OpTemplateParser:
                 call_args_with_tensor.append(call_arg)
             elif tuple_input_to_cpp_type(op_arg.arg_dtype) and pyboost_utils.is_tensor_list(op_arg):
                 need_malloc_tensors.append(call_arg + "_vector")
-                tensor_list_convert.append(get_tuple_input_convert(call_arg, op_arg.arg_dtype))
+                tensor_list_convert.append(
+                    get_tuple_input_convert(call_arg, op_arg.arg_dtype))
                 call_args_with_tensor.append(call_arg + "_vector")
             else:
                 call_args_with_tensor.append(call_arg)
@@ -139,14 +143,17 @@ class OpTemplateParser:
         for op_arg, call_arg in zip(self.op_proto.op_args, call_args):
             if number_input_to_cpp_type(op_arg.arg_dtype):
                 call_args_after_convert.append(call_arg + "_imm")
-                const_number_convert.append(get_const_number_convert(call_arg, op_arg))
+                const_number_convert.append(
+                    get_const_number_convert(call_arg, op_arg))
             elif tuple_input_to_cpp_type(op_arg.arg_dtype):
                 call_args_after_convert.append(call_arg + "_vector")
-                value_tuple_convert.append(get_tuple_input_convert(call_arg, op_arg.arg_dtype))
+                value_tuple_convert.append(
+                    get_tuple_input_convert(call_arg, op_arg.arg_dtype))
             else:
                 call_args_after_convert.append(call_arg)
         if const_number_convert:
-            const_number_convert.insert(0, '// Convert ValuePtr to c++ scalar\n')
+            const_number_convert.insert(
+                0, '// Convert ValuePtr to c++ scalar\n')
         if value_tuple_convert:
             value_tuple_convert.insert(0, '// ValueTuple to std::vector\n')
         return call_args_after_convert, value_tuple_convert, const_number_convert
@@ -192,7 +199,8 @@ class OpTemplateParser:
         args_signature = self.op_proto.op_args_signature
         if args_signature is not None:
             dtype_group = args_signature.dtype_group
-            indexes = {arg.arg_name: index for index, arg in enumerate(self.op_proto.op_args)}
+            indexes = {arg.arg_name: index for index,
+                       arg in enumerate(self.op_proto.op_args)}
             if dtype_group is not None:
                 match = re.findall(r'\((.*?)\)', dtype_group)
                 for item in match:
@@ -292,7 +300,8 @@ class OpTemplateParser:
                 op_outputs = 'outputs'
                 call_outputs = 'outputs_'
             else:
-                raise Exception("Not support return type {}".format(returns_type[0]))
+                raise Exception(
+                    "Not support return type {}".format(returns_type[0]))
         elif len(returns_type) > 1:
             outputs_str = ''
             for i in range(len(returns_type)):
@@ -390,43 +399,62 @@ class OpTemplateParser:
             str: Generated argument handler processing code.
         """
         tensor_arg_handler_prt_template = Template(
-            "arg_list[${idx}] = "
-            "(*pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", arg_list[${idx}]))->value();\n"
+            "parse_args.arg_list_[${idx}] = "
+            "py::cast((*pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", \
+parse_args.arg_list_[${idx}]))->value());\n"
+            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
+            "parse_args.dst_types_[${idx}] = ${new_type};\n"
         )
         function_arg_handler_prt_template = Template(
-            "arg_list[${idx}] = "
-            "(*${func_str}(\"${func_name}\", \"${op_arg_name}\", arg_list[${idx}]))->value();\n"
+            "parse_args.arg_list_[${idx}] = "
+            "py::cast((*${func_str}(\"${func_name}\", \"${op_arg_name}\", parse_args.arg_list_[${idx}]))->value());\n"
+            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
+            "parse_args.dst_types_[${idx}] = ${new_type};\n"
         )
         arg_handler_prt_template = (
             tensor_arg_handler_prt_template) if is_tensor_api else function_arg_handler_prt_template
 
         arg_handler_template = Template(
-            "arg_list[${idx}] = "
-            "pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", arg_list[${idx}]);\n"
+            "parse_args.arg_list_[${idx}] = "
+            "py::cast(pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", parse_args.arg_list_[${idx}]));\n"
+            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
+            "parse_args.dst_types_[${idx}] = ${new_type};\n"
         )
         arg_handler_optional_template = Template(
-            'if (!py::isinstance<py::none>(arg_list[${idx}])) {\n'
+            'if (!py::isinstance<py::none>(parse_args.arg_list_[${idx}])) {\n'
             '  ${arg_handler_str}\n'
             '}\n'
         )
-
+        arg_handler_type_map = {"to_2d_paddings": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                "dtype_to_type_id": "ops::OP_DTYPE::DT_INT",
+                                "to_kernel_size": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                "to_strides": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                "str_to_enum": "ops::OP_DTYPE::DT_INT",
+                                "to_pair": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                "to_dilations": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                "to_output_padding": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                "to_rates": "ops::OP_DTYPE::DT_TUPLE_INT"}
         arg_handler_processor = []
         op_args = op_proto.op_args
         for idx, op_arg in enumerate(op_args):
             arg_handler = op_arg.arg_handler
-            func_str = ''.join(word.capitalize() for word in arg_handler.split('_'))
+            func_str = ''.join(word.capitalize()
+                               for word in arg_handler.split('_'))
             if arg_handler:
                 op_arg_name = op_arg.arg_name
+                new_type = arg_handler_type_map[arg_handler]
                 if func_str in ("StrToEnum", "DtypeToTypeId"):
                     arg_handler_str = arg_handler_prt_template.replace(func_str=func_str,
                                                                        func_name=func_name,
                                                                        op_arg_name=op_arg_name,
-                                                                       idx=idx)
+                                                                       idx=idx,
+                                                                       new_type=new_type)
                 else:
                     arg_handler_str = arg_handler_template.replace(func_str=func_str,
                                                                    func_name=func_name,
                                                                    op_arg_name=op_arg_name,
-                                                                   idx=idx)
+                                                                   idx=idx,
+                                                                   new_type=new_type)
 
                 if op_arg.default == "None":
                     arg_handler_str = arg_handler_optional_template.replace(idx=idx,
@@ -434,3 +462,56 @@ class OpTemplateParser:
                 arg_handler_processor.append(arg_handler_str)
 
         return arg_handler_processor
+
+    def get_input_tensor_index(self, op_proto):
+        """
+        Get index of input.
+
+        Args:
+            op_proto (OpProto): Function prototype to generate dispatch strings for.
+
+        Returns:
+            int: Index of input.
+        """
+        op_name = op_proto.op_class.name
+        op_args = op_proto.op_args
+        if op_name in K.INPUT_NAME_MAP:
+            self_index = [i for i in range(
+                len(op_args)) if op_args[i].arg_name == K.INPUT_NAME_MAP[op_name]]
+        else:
+            self_index = [i for i in range(
+                len(op_args)) if op_args[i].arg_name in K.INPUT_ARGS_NAME]
+        if len(self_index) != 1:
+            raise ValueError(
+                f'There must be only one field named \'input\'. But got {len(self_index)} in {op_name}')
+        return self_index[0]
+
+    def get_convert_args_str(self, op_proto, is_tensor_api):
+        """
+        Generates argument convert processing code for the given function prototype.
+
+        Args:
+            op_proto (OpProto): Operator prototype instance to generate argument processing for.
+
+        Returns:
+            str: Generated argument convert processing code.
+        """
+        self_index = 0
+        if is_tensor_api:
+            self_index = self.get_input_tensor_index(op_proto)
+        convert_args_str = ""
+        for idx, op_arg in enumerate(op_proto.op_args):
+            if is_tensor_api:
+                if self_index == idx:
+                    convert_args_str += "input_tensor, "
+                    continue
+            is_optional = is_optional_param(op_arg)
+            arg_convert_template = Template("parse_args.ConvertOptional<${des_type}>(${index}), ") if is_optional \
+                                   else Template("parse_args.Convert<${des_type}>(${index}), ")
+            if op_arg.is_type_id:
+                arg_type_str = get_input_args_type_str('type', False)
+            else:
+                arg_type_str = get_input_args_type_str(op_arg.arg_dtype, False)
+            convert_args_str += arg_convert_template.replace(index=idx,
+                                                             des_type=arg_type_str[:-3])
+        return convert_args_str[:-2]

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

@@ -22,15 +22,16 @@ the necessary header files for the generated functions.
 
 import os
 
-import pyboost_utils
-from pyboost_utils import get_convert_type_str, is_optional_param, is_op_multi_output
-import template
-from template 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.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
+from pyboost.pyboost_utils import get_convert_type_str, is_optional_param, is_op_multi_output, get_input_args_type_str, is_tensor_list
+
+from .op_template_parser import OpTemplateParser
 
 
 class PyboostFunctionsGenerator(BaseGenerator):
@@ -45,7 +46,7 @@ class PyboostFunctionsGenerator(BaseGenerator):
     def __init__(self):
         """Initializes the PyboostFunctionsGenerator with the necessary templates."""
         self.pyboost_func_include_header_template = Template(
-            f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/${{operator_name}}.h"\n'
+            f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n'
         )
         self.convert_optional_to_value_template = Template(
             "auto ${output} = PyNativeAlgo::PyBoost::OptionalToValue(${input});\n"
@@ -59,14 +60,17 @@ class PyboostFunctionsGenerator(BaseGenerator):
             'op_run_info->requires_grad);\n'
         )
         self.convert_template = Template("auto $arg_name = converter.${convert_func}(args, $arg_index);\n")
+        self.input_args_template = Template(" const ${arg_type}& ${arg_name},")
         self.PYBOOST_FUNCTION_TEMPLATE = template.PYBOOST_FUNCTION_TEMPLATE
         self.PYBOOST_COMM_FUNCTION_TEMPLATE = template.PYBOOST_COMM_FUNCTION_TEMPLATE
+        self.PYBOOST_FUNCTION_DYNAMIC_OUTPUT_TEMPLATE = template.PYBOOST_FUNCTION_DYNAMIC_OUTPUT_TEMPLATE
         self.REGISTER_DEFINE_TEMPLATE = template.REGISTER_DEFINE_TEMPLATE
         self.REGISTER_TEMPLATE = template.REGISTER_TEMPLATE
         self.PYBOOST_HEADER_TEMPLATE = template.PYBOOST_FUNCTIONS_CC_TEMPLATE
         self.TENSOR_FUNC_CLASS_REG = template.TENSOR_FUNC_CLASS_REG
+        self.OP_DEF_INC_HEAD_TEMPLATE = template.OP_DEF_INC_HEAD_TEMPLATE
 
-    def generate(self, work_path, op_protos, tensor_func_protos_data):
+    def generate(self, work_path, op_protos):
         """
         Generates the C++ PyBoost functions and writes them to the specified files.
 
@@ -77,7 +81,6 @@ class PyboostFunctionsGenerator(BaseGenerator):
         Args:
             work_path (str): The file path where the generated files will be saved.
             op_protos (list): A list of operator prototypes containing information about the operators to be processed.
-            tensor_func_protos_data(dict): A dict of tensor prototypes containing device-related information.
 
         Returns:
             None
@@ -85,6 +88,7 @@ class PyboostFunctionsGenerator(BaseGenerator):
         pyboost_func_str = ''
         pyboost_func_pybind_def = ''
         pyboost_func_include_headers_str = ''
+        ops_inc_head_set = set()
         for op_proto in op_protos:
             if op_proto.op_dispatch is None or not op_proto.op_dispatch.enable:
                 continue
@@ -100,15 +104,16 @@ class PyboostFunctionsGenerator(BaseGenerator):
             grad_args_str = self._get_grad_args_str(op_proto)
             cast_args_str = self._get_cast_to_value_str(op_proto)
             view_arg_str = self._get_first_str(op_proto.op_view, grad_args_str)
+            op_input_args_str = self._get_input_args_str(op_proto)
             view_arg_str = ", " + view_arg_str if view_arg_str else ''
             multi_ouptut_str = 'Multi' if is_op_multi_output(op_proto.op_returns) else ''
-            # communication operators have different func template
-            function_tpl = self.PYBOOST_COMM_FUNCTION_TEMPLATE \
-                if op_proto.op_dispatch.is_comm_op else self.PYBOOST_FUNCTION_TEMPLATE
+            output_num_str = len(op_proto.op_returns)
+            function_tpl = self._get_function_tpl(op_proto)
             pyboost_func_str += function_tpl.replace(func_name=op_pyboost_func_name,
                                                      op_def_name=op_def_name_str,
                                                      type_num=type_num,
                                                      same_type=same_type,
+                                                     input_args=op_input_args_str,
                                                      parser_body=parser_body_str,
                                                      op_name=op_proto.op_class.name,
                                                      class_name=op_proto.op_class.name,
@@ -120,6 +125,7 @@ class PyboostFunctionsGenerator(BaseGenerator):
                                                      cast_args=cast_args_str,
                                                      view_arg=view_arg_str,
                                                      is_multi=multi_ouptut_str,
+                                                     output_num=output_num_str,
                                                      operator_name=op_proto.op_name)
             pyboost_func_str = pyboost_func_str + template.NEW_LINE + template.NEW_LINE
             pyboost_op_name = op_parser.get_pyboost_name()
@@ -130,9 +136,11 @@ class PyboostFunctionsGenerator(BaseGenerator):
                 class_name=op_proto.op_class.name)
             pyboost_func_include_headers_str += self.pyboost_func_include_header_template.replace(
                 operator_name=op_proto.op_name)
+            ops_inc_head_set.add(self.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_proto.op_class.name[0].lower()))
         register_func_str = self.REGISTER_TEMPLATE.replace(register_func=pyboost_func_pybind_def)
         function_class_register = self._get_function_class_register(op_protos)
-        pyboost_func_file = self.PYBOOST_HEADER_TEMPLATE.replace(include_op_header=pyboost_func_include_headers_str,
+        pyboost_func_file = self.PYBOOST_HEADER_TEMPLATE.replace(ops_inc=list(sorted(ops_inc_head_set)),
+                                                                 include_op_header=pyboost_func_include_headers_str,
                                                                  function_body=pyboost_func_str,
                                                                  register_function_body=register_func_str,
                                                                  function_class_register=function_class_register)
@@ -140,7 +148,6 @@ class PyboostFunctionsGenerator(BaseGenerator):
         file_name = "pyboost_functions.cc"
         save_file(save_path, file_name, pyboost_func_file)
 
-
     def _get_cast_args_with_type_str(self, op_proto, cast_args_str):
         args_with_type = []
         for op_arg, cast_args_name in zip(op_proto.op_args, cast_args_str):
@@ -148,7 +155,6 @@ class PyboostFunctionsGenerator(BaseGenerator):
             args_with_type.append("const " + input_dtype + " &" + cast_args_name)
         return list(args_with_type)
 
-
     def _get_function_class_register(self, op_protos) -> str:
         """
         Generates a function class registration string for tensor functions.
@@ -189,10 +195,33 @@ class PyboostFunctionsGenerator(BaseGenerator):
                 convert_type_str = get_convert_type_str('type', is_optional)
             else:
                 convert_type_str = get_convert_type_str(op_arg.arg_dtype, is_optional)
+
             parser_func_str += self.convert_template.replace(arg_name=op_arg.arg_name, convert_func=convert_type_str,
                                                              arg_index=pyboost_utils.get_index(index))
         return parser_func_str
 
+
+    def _get_input_args_str(self, op_proto: OpProto) -> str:
+        """
+        Generates the input arguments list for the pyboost operator.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing the argument information.
+
+        Returns:
+            str: The generated input arguments list as a string.
+        """
+        parser_func_str = ''
+        for _, op_arg in enumerate(op_proto.op_args):
+            is_optional = is_optional_param(op_arg)
+            if op_arg.is_type_id:
+                arg_type_str = get_input_args_type_str('type', is_optional)
+            else:
+                arg_type_str = get_input_args_type_str(op_arg.arg_dtype, is_optional)
+            parser_func_str += self.input_args_template.replace(arg_name=op_arg.arg_name, arg_type=arg_type_str)
+        return parser_func_str[:-1]
+
+
     def _get_convert_stub_str(self, op_proto: OpProto):
         """
         Generates the conversion stub code for the operator's arguments.
@@ -368,3 +397,11 @@ class PyboostFunctionsGenerator(BaseGenerator):
                 arg_str = grad_arg
                 break
         return arg_str
+
+
+    def _get_function_tpl(self, op_proto: OpProto):
+        if len(op_proto.op_returns) == 1 and is_tensor_list(op_proto.op_returns[0]):
+            # op output size is unknown
+            return self.PYBOOST_FUNCTION_DYNAMIC_OUTPUT_TEMPLATE
+        return self.PYBOOST_COMM_FUNCTION_TEMPLATE \
+            if op_proto.op_dispatch.is_comm_op else self.PYBOOST_FUNCTION_TEMPLATE

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

@@ -22,13 +22,15 @@ operation's primitive and arguments. The generated file is saved to the specifie
 
 import os
 
-import template
+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.base_generator import BaseGenerator
+from common.op_proto import OpProto
 
-from template import Template
-import gen_constants as K
-from gen_utils import save_file
-from op_template_parser import OpTemplateParser
-from base_generator import BaseGenerator
+from .op_template_parser import OpTemplateParser
+from .pyboost_utils import is_optional_param, get_input_args_type_str
 
 
 class PyboostFunctionsHeaderGenerator(BaseGenerator):
@@ -41,8 +43,13 @@ class PyboostFunctionsHeaderGenerator(BaseGenerator):
         self.PYBOOST_FUNCTION_HEADER_TEMPLATE = template.PYBOOST_FUNCTION_HEADER_TEMPLATE
 
         self.pyboost_func_template = Template(
-            'py::object ME_EXPORT ${func_name}_Base(const PrimitivePtr &prim, const py::list &args);'
+            'py::object PYNATIVE_EXPORT ${func_name}_Base(const PrimitivePtr &prim, const py::list &args);'
         )
+        self.pyboost_op_func_template = Template(
+            'py::object ME_EXPORT ${func_name}_OP(const PrimitivePtr &prim, '
+            'const std::vector<ops::OP_DTYPE>& source_type, ${input_args});'
+        )
+        self.input_args_template = Template(" const ${arg_type}& ${arg_name},")
 
     def generate(self, work_path, op_protos):
         """
@@ -55,14 +62,39 @@ class PyboostFunctionsHeaderGenerator(BaseGenerator):
         Returns:
             None: The method writes the generated header file to the specified directory.
         """
-        func_list = []
+        prim_func_list = []
+        op_func_list_str = []
         for op_proto in op_protos:
             if op_proto.op_dispatch is None or not op_proto.op_dispatch.enable:
                 continue
             op_parser = OpTemplateParser(op_proto)
             op_pyboost_func_name = op_parser.get_pyboost_func_name()
-            func_list.append(self.pyboost_func_template.replace(func_name=op_pyboost_func_name))
-        pyboost_func_h_str = self.PYBOOST_FUNCTION_HEADER_TEMPLATE.replace(prim_func_list=func_list)
+            op_input_args_str = self._get_input_args_str(op_proto)
+            prim_func_list.append(self.pyboost_func_template.replace(func_name=op_pyboost_func_name))
+            op_func_list_str.append(self.pyboost_op_func_template.replace(func_name=op_pyboost_func_name,
+                                                                          input_args=op_input_args_str))
+        pyboost_func_h_str = self.PYBOOST_FUNCTION_HEADER_TEMPLATE.replace(prim_func_list=prim_func_list,
+                                                                           op_func_list=op_func_list_str)
         save_path = os.path.join(work_path, K.PIPELINE_PYBOOST_FUNC_GEN_PATH)
         file_name = "pyboost_functions.h"
         save_file(save_path, file_name, pyboost_func_h_str)
+
+    def _get_input_args_str(self, op_proto: OpProto) -> str:
+        """
+        Generates the input arguments list for the pyboost operator.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing the argument information.
+
+        Returns:
+            str: The generated input arguments list as a string.
+        """
+        parser_func_str = ''
+        for _, op_arg in enumerate(op_proto.op_args):
+            is_optional = is_optional_param(op_arg)
+            if op_arg.is_type_id:
+                arg_type_str = get_input_args_type_str('type', is_optional)
+            else:
+                arg_type_str = get_input_args_type_str(op_arg.arg_dtype, is_optional)
+            parser_func_str += self.input_args_template.replace(arg_name=op_arg.arg_name, arg_type=arg_type_str)
+        return parser_func_str[:-1]

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

@@ -22,11 +22,11 @@ argument processing and includes appropriate documentation descriptions.
 
 import os
 
-import template
-import gen_constants as K
-from gen_utils import save_file
-from op_proto import OpProto
-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
 
 
 class PyboostFunctionsPyGenerator(BaseGenerator):
@@ -78,7 +78,7 @@ class PyboostFunctionsPyGenerator(BaseGenerator):
                                                             func_args=func_args,
                                                             input_args=input_args,
                                                             func_impl_name=func_impl_name)
-        py_header = K.PY_LICENSE + self.IMPORT_PYBOOST_FUNC_HEADER
+        py_header = template.PY_LICENSE_STR + self.IMPORT_PYBOOST_FUNC_HEADER
         save_file(os.path.join(work_path, K.PY_AUTO_GEN_PATH), "gen_extend_func.py", py_header + gen_py)
 
     def _get_op_description_dict(self, doc_yaml_data):

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

@@ -21,14 +21,15 @@ function definitions, includes necessary headers, and generates the registration
 
 import os
 
-import template
-from template import Template
-from gen_utils import save_file
-import gen_constants as K
-from op_proto import OpProto
-from base_generator import BaseGenerator
-from op_template_parser import OpTemplateParser
-import pyboost_utils
+import common.template as template
+from common.template import Template
+from common.gen_utils import save_file
+import common.gen_constants as K
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+from pyboost import pyboost_utils
+
+from .op_template_parser import OpTemplateParser
 
 
 class PyboostGradFunctionsGenerator(BaseGenerator):
@@ -44,7 +45,7 @@ class PyboostGradFunctionsGenerator(BaseGenerator):
     def __init__(self):
         super().__init__()
         self.pyboost_func_include_header_template = Template(
-            f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/${{operator_name}}.h"\n')
+            f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n')
         self.GEN_OPS_DEF_HEADER_TEMPLATE = template.GEN_OPS_DEF_HEADER_TEMPLATE
         self.contiguous_template = Template(
             "convert_$arg_name = runtime::ValueConverter::ContiguousTensorValue($device_target, convert_$arg_name);\n")
@@ -101,7 +102,7 @@ class PyboostGradFunctionsGenerator(BaseGenerator):
             template.PYBOOST_GRAD_HEADER_TEMPLATE.replace(include_op_header=pyboost_func_include_headers_str,
                                                           function_body=pyboost_func_str,
                                                           register_function_body=register_func_str)
-        save_path = os.path.join(work_path, K.RUNTIME_PYBOOST_FUNC_GEN_PATH)
+        save_path = os.path.join(work_path, K.PYBOOST_GRAD_FUNC_GEN_PATH)
         file_name = "pyboost_grad_functions.cc"
         save_file(save_path, file_name, pyboost_func_file)
 

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

@@ -22,13 +22,14 @@ that can be used to call the Pyboost primitive implementations.
 """
 
 import os
-import template
-import gen_constants as K
-from gen_utils import save_file
+import common.template as template
+import common.gen_constants as K
+from common.gen_utils import save_file
 
-from op_proto import OpProto
-from op_template_parser import OpTemplateParser
-from base_generator import BaseGenerator
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+
+from .op_template_parser import OpTemplateParser
 
 
 class PyboostInnerPrimGenerator(BaseGenerator):
@@ -72,7 +73,7 @@ class PyboostInnerPrimGenerator(BaseGenerator):
             None
         """
         gen_py = ''
-        gen_header = K.PY_LICENSE + self.IMPORT_PYBOOST_PRIM_HEADER
+        gen_header = template.PY_LICENSE_STR + self.IMPORT_PYBOOST_PRIM_HEADER
         for op_proto in op_protos:
             # only process pyboost enabled scenario
             if op_proto.op_dispatch is None:

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

@@ -22,14 +22,14 @@ while the PyboostGradFunctionsHeaderGenerator creates the corresponding function
 
 import os
 
-import pyboost_utils
-from pyboost_utils import is_optional_param
-import template
-from template import Template
-import gen_constants as K
-from gen_utils import save_file
-from op_proto import OpProto
-from base_generator import BaseGenerator
+from pyboost import pyboost_utils
+from pyboost.pyboost_utils import is_optional_param
+import common.template as template
+from common.template import 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
 
 
 class PyboostGradFunctionsCppGenerator(BaseGenerator):
@@ -55,7 +55,7 @@ class PyboostGradFunctionsCppGenerator(BaseGenerator):
         self.native_function_multi_output_template = template.MULTI_OUTPUT_TEMPLATE
         self.native_function_single_output_template = "const auto &output_value = op->outputs()[0];\n"
         self.native_include_header_template = Template(
-            f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/${{operator_name}}.h"\n')
+            f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n')
         self.convert_template = Template(
             "auto convert_$arg_name = runtime::ValueConverter::${convert_func}(ConvertNode2Value($arg_name));\n")
 
@@ -76,6 +76,7 @@ class PyboostGradFunctionsCppGenerator(BaseGenerator):
         """
         pyboost_func_str = ''
         pyboost_func_include_headers_str = ''
+        ops_inc_head_set = set()
         for op_proto in op_protos:
             if op_proto.op_dispatch is None:
                 continue
@@ -103,9 +104,12 @@ class PyboostGradFunctionsCppGenerator(BaseGenerator):
             pyboost_func_str = pyboost_func_str + template.NEW_LINE
             pyboost_func_include_headers_str += (
                 self.native_include_header_template.replace(operator_name=op_proto.op_name))
+            ops_inc_head_set.add(
+                template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_proto.op_class.name[0].lower()))
         native_grad_func_file = \
             self.PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE.replace(include_op_header=pyboost_func_include_headers_str,
-                                                                function_body=pyboost_func_str)
+                                                                function_body=pyboost_func_str,
+                                                                ops_inc=list(sorted(ops_inc_head_set)))
         save_file(os.path.join(work_path, K.PYBOOST_NATIVE_GRAD_FUNC_GEN_PATH),
                   "pyboost_native_grad_functions.cc", native_grad_func_file)