mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.1__cp310-cp310-win_amd64.whl

mindspore/ops_generate/gen_pyboost_func.py

@@ -28,6 +28,7 @@ import template
 from template import CppTemplate
 from op_proto import OpProto
 from gen_utils import check_change_and_replace_file, py_licence_str, write_file
+import gen_constants as K
 
 
 @dataclass
@@ -47,7 +48,7 @@ def generate_pyboost_base_op_header_code(work_path, op_name_str, operator_name,
                                                                              op_name_upper=op_name_str.upper(),
                                                                              call_args=call_args_with_type,
                                                                              return_type=cpp_func_return)
-    op_header_dir_path = os.path.join(work_path, "mindspore/ccsrc/kernel/pyboost/auto_generate/")
+    op_header_dir_path = os.path.join(work_path, f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
     pathlib.Path(op_header_dir_path).mkdir(parents=True, exist_ok=True)
     tmp_op_file_path = os.path.join(op_header_dir_path, "tmp_" + operator_name + ".h")
     dst_op_file_path = os.path.join(op_header_dir_path, operator_name + ".h")
@@ -187,35 +188,28 @@ def generate_pyboost_op_source_code(work_path, op_proto, template_paths, convert
                                         return_values=converter.call_func_outputs,
                                         customize_func=op_proto.ascend + "Customize",
                                         )
-            customize_include = "#include \"plugin/device/ascend/kernel/pyboost/customize/{}.h\"".format(
-                operator_name.lower())
+            customize_include = f'#include "{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/customize/{operator_name.lower()}.h"'
         elif is_cpu and op_proto.cpu != 'default':
             call_impl = cus_tpl.replace(call_args=converter.call_args,
                                         return_values=converter.call_func_outputs,
                                         customize_func=op_proto.cpu + "Customize",
                                         )
-            customize_include = "#include \"plugin/device/cpu/kernel/pyboost/customize/{}.h\"".format(
-                operator_name.lower())
+            customize_include = f'#include "{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/customize/{operator_name.lower()}.h"'
             register_custom_kernel = "MS_REG_PYBOOST_CPU_CUSTOM_KERNEL({});".format(op_name_str)
         elif is_gpu and op_proto.gpu != 'default':
             call_impl = cus_tpl.replace(call_args=converter.call_args,
                                         return_values=converter.call_func_outputs,
                                         customize_func=op_proto.gpu + "Customize",
                                         )
-            customize_include = "#include \"plugin/device/gpu/kernel/pyboost/customize/{}.h\"".format(
-                operator_name.lower())
+            customize_include = f'#include "{K.MS_OPS_KERNEL_PATH}/gpu/pyboost/customize/{operator_name.lower()}.h"'
             register_custom_kernel = "MS_REG_PYBOOST_GPU_CUSTOM_KERNEL({});".format(op_name_str)
         elif op_proto.is_view:
-            set_output_abs = "SetOutputAbstract();"
-            if converter.call_func_outputs == "outputs_":
-                set_output_abs = "SetOutputTupleAbstract();"
             call_impl = view_tpl.replace(op_name=op_proto.class_name,
                                          call_args=converter.call_args,
                                          call_tensors=call_args_tensor,
                                          return_values=converter.call_func_outputs,
-                                         input=converter.call_args[0],
-                                         set_output_abs=set_output_abs)
-            customize_include = "#include \"mindspore/core/ops/view/{}_strides_calc.h\"".format(proto_operator_name)
+                                         input=converter.call_args[0])
+            customize_include = f'#include "{K.MS_OPS_VIEW_PATH}/{proto_operator_name}_strides_calc.h"'
         else:
             cast_input_code, real_call_args_tensor = generate_tensor_cpu_cast_input_code(
                 converter.call_args_with_tensor, call_args_tensor)
@@ -226,7 +220,7 @@ def generate_pyboost_op_source_code(work_path, op_proto, template_paths, convert
             aclnn_name = AclnnUtils.get_aclnn_interface(op_name_str)
             if converter.inplace_process != '':
                 real_output = ''
-            customize_include = '#include "ops/auto_generate/gen_ops_primitive.h"'
+            customize_include = f'#include "{K.MS_OP_DEF_AUTO_GENERATE_PATH}/gen_ops_primitive.h"'
 
             call_impl = call_tpl.replace(aclnn_name=aclnn_name,
                                          call_args=converter.call_args,
@@ -270,10 +264,10 @@ def generate_pyboost_op_register_source_code(work_path, all_ops, all_operator_na
     for op_name in all_ops:
         factory_str += "template class OpFactory<{0}>;\n".format(op_name)
     for operator_name in all_operator_names:
-        include_str += "#include \"kernel/pyboost/auto_generate/{0}.h\"\n".format(operator_name)
+        include_str += f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/{operator_name}.h"\n'
     op_register_file_str = template.PYBOOST_OP_REGISTER_TEMPLATE.replace(op_includes=include_str,
                                                                          op_factory_templates=factory_str)
-    op_register_dir_path = os.path.join(work_path, "mindspore/ccsrc/kernel/pyboost/auto_generate/")
+    op_register_dir_path = os.path.join(work_path, f"{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/")
     pathlib.Path(op_register_dir_path).mkdir(parents=True, exist_ok=True)
     tmp_op_register_file_path = os.path.join(op_register_dir_path, "tmp_" + "op_register.cc")
     dst_op_register_file_path = os.path.join(op_register_dir_path, "op_register.cc")
@@ -353,19 +347,22 @@ def generate_pyboost_outputs(op_proto):
     return op_outputs, call_outputs
 
 
-def generate_ops_header_files(work_path, yaml_data):
+def generate_ops_header_files(work_path, yaml_data, extra_ops=None):
     """
     :param work_path:
     :param yaml_data:
+    :prama extra_ops(default:None):
     :return: void
     """
     extern_str = ''
-    extern_template = CppTemplate("MS_EXPORT extern OpDef g${op_name};\n")
+    extern_template = CppTemplate("OPS_API extern OpDef g${op_name};\n")
     for operator_name, operator_data in yaml_data.items():
         op_proto = OpProto.load_from_yaml(operator_name, operator_data)
         extern_str += extern_template.replace(op_name=op_proto.class_name)
+    for class_name in extra_ops:
+        extern_str += extern_template.replace(op_name=class_name)
     ops_header_file = template.GEN_OPS_DEF_HEADER_TEMPLATE.replace(extern_variable=extern_str)
-    dir_path = os.path.join(work_path, "mindspore/core/ops/auto_generate")
+    dir_path = os.path.join(work_path, K.MS_OP_DEF_AUTO_GENERATE_PATH)
     pathlib.Path(dir_path).mkdir(parents=True, exist_ok=True)
     dst_file_path = os.path.join(dir_path, "gen_ops_def.h")
     tmp_file_path = os.path.join(dir_path, "tmp_gen_ops_def.h")
@@ -396,10 +393,10 @@ def get_convert_tensor_template():
     Get convert tensor template
     """
     convert_to_tensor_template = CppTemplate(
-        'auto ${output} = PyNativeAlgo::Common::ConvertStubNodeToTensor(${input}, ${need_contiguous}, '\
+        'auto ${output} = PyNativeAlgo::Common::ConvertStubNodeToTensor(${input}, ${need_contiguous}, ' \
         'op_run_info->requires_grad);\n')
     convert_to_tensor_list_template = CppTemplate(
-        'auto ${output} = PyNativeAlgo::Common::ConvertStubNodeToValueTuple(${input}, ${need_contiguous}, '\
+        'auto ${output} = PyNativeAlgo::Common::ConvertStubNodeToValueTuple(${input}, ${need_contiguous}, ' \
         'op_run_info->requires_grad);\n')
     return convert_to_tensor_template, convert_to_tensor_list_template
 
@@ -411,7 +408,8 @@ def generate_pyboost_functions(work_path, yaml_data):
     pyboost_func_str = ''
     pyboost_func_pybind_def = ''
     pyboost_func_include_headers_str = ''
-    pyboost_func_include_header_template = CppTemplate("#include \"kernel/pyboost/auto_generate/${operator_name}.h\"\n")
+    pyboost_func_include_header_template = CppTemplate(
+        f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/${{operator_name}}.h"\n')
     for operator_name, operator_data in yaml_data.items():
         op_proto = OpProto.load_from_yaml(operator_name, operator_data)
         if not op_proto.is_dispatch:
@@ -498,15 +496,18 @@ def generate_pyboost_functions(work_path, yaml_data):
             call_args_str.append(call_arg)
             cast_args_str.append(cast_arg)
         type_num, same_type = gen_signature_same_type_table(op_proto.indexes, operator_data)
-        pyboost_func_str += template.PYBOOST_FUNCTION_TEMPLATE.replace(func_name=op_proto.pyboost_function_name,
-                                                                       op_def_name=op_def_name_str, same_type=same_type,
-                                                                       type_num=type_num, parser_body=parser_body_str,
-                                                                       op_name=op_name_str,
-                                                                       convert_stub=convert_stub_str,
-                                                                       optional_to_value=optional_to_value_str,
-                                                                       call_args=call_args_str, grad_args=grad_args_str,
-                                                                       cast_args=cast_args_str, op_args=op_args_str,
-                                                                       class_name=op_proto.class_name)
+
+        funcion_tpl = template.PYBOOS_COMM_FUNCTION_TEMPLATE \
+            if op_proto.is_comm_op else template.PYBOOST_FUNCTION_TEMPLATE
+        pyboost_func_str += funcion_tpl.replace(func_name=op_proto.pyboost_function_name,
+                                                op_def_name=op_def_name_str, same_type=same_type,
+                                                type_num=type_num, parser_body=parser_body_str,
+                                                op_name=op_name_str,
+                                                convert_stub=convert_stub_str,
+                                                optional_to_value=optional_to_value_str,
+                                                call_args=call_args_str, grad_args=grad_args_str,
+                                                cast_args=cast_args_str, op_args=op_args_str,
+                                                class_name=op_proto.class_name)
         pyboost_func_str = pyboost_func_str + template.NEW_LINE + template.NEW_LINE
         pyboost_func_pybind_def += template.REGISTER_DEFINE_TEMPLATE.replace(
             pyboost_op_name=get_pyboost_name(op_proto.operator_name),
@@ -531,7 +532,7 @@ def convert_value_type(op_proto: OpProto) -> str:
     :return: str
     """
     convert_template = CppTemplate(
-        "auto $arg_name = ValueConverter::${convert_func}(op_runner_info->inputs, $arg_index);\n")
+        "auto convert_$arg_name = ValueConverter::${convert_func}(op_runner_info->inputs[$arg_index]);\n")
     parser_func_str = ''
     for index, arg in enumerate(op_proto.op_args):
         is_optional = is_optional_param(arg)
@@ -541,9 +542,26 @@ def convert_value_type(op_proto: OpProto) -> str:
     return parser_func_str
 
 
-def contiguous_tensor_value(op_proto: OpProto) -> str:
+def convert_native_value_type(op_proto: OpProto) -> str:
+    """
+    Generate native convert func
+    :param op_proto:
+    :return: str
+    """
+    convert_template = CppTemplate(
+        "auto convert_$arg_name = runtime::ValueConverter::${convert_func}(ConvertNode2Value($arg_name));\n")
+    parser_func_str = ''
+    for arg in op_proto.op_args:
+        is_optional = is_optional_param(arg)
+        convert_type_str = get_value_convert_type_str(arg.arg_dtype, is_optional)
+        parser_func_str += convert_template.replace(arg_name=arg.arg_name, convert_func=convert_type_str)
+    return parser_func_str
+
+
+def contiguous_tensor_value(op_proto: OpProto, device_target: str) -> str:
     """
     Generate parser func
+    :param device_target:
     :param op_proto:
     :return: str
     """
@@ -551,13 +569,13 @@ def contiguous_tensor_value(op_proto: OpProto) -> str:
     if op_proto.is_view:
         return ''
     contiguous_template = CppTemplate(
-        "$arg_name = ValueConverter::ContiguousTensorValue(op_runner_info, $arg_name);\n")
+        "convert_$arg_name = runtime::ValueConverter::ContiguousTensorValue($device_target, convert_$arg_name);\n")
     contiguous_func_str = ''
     need_contiguous_dtype = {'tensor', 'tuple[tensor]'}
     for arg in op_proto.op_args:
         if arg.arg_dtype not in need_contiguous_dtype:
             continue
-        contiguous_func_str += contiguous_template.replace(arg_name=arg.arg_name)
+        contiguous_func_str += contiguous_template.replace(arg_name=arg.arg_name, device_target=device_target)
     return contiguous_func_str
 
 
@@ -568,7 +586,8 @@ def generate_pyboost_grad_functions(work_path, yaml_data):
     pyboost_func_str = ''
     pyboost_func_reg_def = ''
     pyboost_func_include_headers_str = ''
-    pyboost_func_include_header_template = CppTemplate("#include \"kernel/pyboost/auto_generate/${operator_name}.h\"\n")
+    pyboost_func_include_header_template = CppTemplate(
+        f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/${{operator_name}}.h"\n')
     for operator_name, operator_data in yaml_data.items():
         if not is_pyboost_enable(operator_data):
             continue
@@ -579,11 +598,12 @@ def generate_pyboost_grad_functions(work_path, yaml_data):
         op_name_str = op_proto.class_name
         op_args_str = [op_arg.arg_name for op_arg in op_proto.op_args]
         convert_value_type_str = convert_value_type(op_proto)
-        convert_value_type_str += contiguous_tensor_value(op_proto)
+        device_target = "op_runner_info->device_target"
+        convert_value_type_str += contiguous_tensor_value(op_proto, device_target)
 
         call_args_str = []
         for op_arg in op_proto.op_args:
-            call_arg = op_arg.arg_name
+            call_arg = 'convert_' + op_arg.arg_name
             call_args_str.append(call_arg)
         pyboost_func_str += template.PYBOOST_GRAD_FUNCTION_TEMPLATE.replace(func_name=op_proto.pyboost_function_name,
                                                                             op_name=op_name_str,
@@ -609,6 +629,70 @@ def generate_pyboost_grad_functions(work_path, yaml_data):
     check_change_and_replace_file(dst_file_path, tmp_file_path)
 
 
+def generate_pyboost_native_grad_functions(work_path, yaml_data):
+    """
+    Generate pyboost native grad  functions file from yaml.
+    """
+    pyboost_func_str = ''
+    pyboost_func_include_headers_str = ''
+    native_function_headers_str = ''
+    native_include_header_template = CppTemplate(
+        f'#include "{K.MS_COMMON_PYBOOST_KERNEL_PATH}/auto_generate/${{operator_name}}.h"\n')
+    native_function_header_template = CppTemplate("static NodePtr $func_name(${call_args_with_type});\n")
+    native_function_sigle_output_template = "const auto &output_value = op->outputs()[0];\n"
+    native_function_multi_output_template = template.MULTI_OUTPUT_TEMPLATE
+    for operator_name, operator_data in yaml_data.items():
+        if not is_pyboost_enable(operator_data):
+            continue
+        op_proto = OpProto.load_from_yaml(operator_name, operator_data)
+        if not op_proto.is_dispatch:
+            continue
+        operator_name = op_proto.operator_name
+        op_name_str = op_proto.class_name
+        op_args_str = [op_arg.arg_name for op_arg in op_proto.op_args]
+        convert_value_type_str = convert_native_value_type(op_proto)
+        convert_value_type_str += contiguous_tensor_value(op_proto, "device_target_")
+        first_var_name = op_proto.op_args[0].arg_name
+        call_args_str = []
+        call_args_with_type = []
+        output_expr = native_function_sigle_output_template
+        if op_proto.is_multi_output:
+            output_expr = native_function_multi_output_template
+        for op_arg in op_proto.op_args:
+            call_arg = 'convert_' + op_arg.arg_name
+            call_args_str.append(call_arg)
+            call_args_with_type.append('const NodePtr &' + op_arg.arg_name)
+        pyboost_func_str += \
+            template.PYBOOST_NATIVE_GRAD_FUNCTION_TEMPLATE.replace(func_name=op_name_str,
+                                                                   op_name=op_name_str,
+                                                                   op_args=op_args_str,
+                                                                   convert_body=convert_value_type_str,
+                                                                   call_args=call_args_str,
+                                                                   call_args_with_type=call_args_with_type,
+                                                                   first_var_name=first_var_name,
+                                                                   output_expr=output_expr)
+        pyboost_func_str = pyboost_func_str + template.NEW_LINE
+        pyboost_func_include_headers_str += native_include_header_template.replace(operator_name=operator_name)
+        func_header = native_function_header_template.replace(func_name=op_name_str,
+                                                              call_args_with_type=call_args_with_type)
+        native_function_headers_str += func_header
+    native_grad_func_file = \
+        template.PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE.replace(include_op_header=pyboost_func_include_headers_str,
+                                                                function_body=pyboost_func_str)
+    native_grad_func_header_file = template.PYBOOST_NATIVE_GRAD_FUNCTIONS_HEADER_TEMPLATE.replace(
+        native_grad_func_def=native_function_headers_str)
+    dir_path = os.path.join(work_path, "mindspore/ccsrc/pipeline/pynative/grad/function/auto_generate")
+    pathlib.Path(dir_path).mkdir(parents=True, exist_ok=True)
+    tmp_file_path = os.path.join(dir_path, "tmp_pyboost_native_grad_functions.cc")
+    dst_file_path = os.path.join(dir_path, "pyboost_native_grad_functions.cc")
+    write_file(tmp_file_path, native_grad_func_file)
+    check_change_and_replace_file(dst_file_path, tmp_file_path)
+    tmp_file_path = os.path.join(dir_path, "tmp_pyboost_native_grad_functions.h")
+    dst_file_path = os.path.join(dir_path, "pyboost_native_grad_functions.h")
+    write_file(tmp_file_path, native_grad_func_header_file)
+    check_change_and_replace_file(dst_file_path, tmp_file_path)
+
+
 def generate_inplace_process_cpp_code(op_proto):
     """ generate_ref_process_cpp_code """
     inplace_process = f'// RefOps update output by input tensor\n'
@@ -640,9 +724,8 @@ def get_auto_generate_template():
                                template.PYBOOST_CPU_CUSTOMIZE_CALL_TEMPLATE]
     op_view_template_path = [template.PYBOOST_ASCEND_VIEW_CALL_TEMPLATE, template.PYBOOST_GPU_VIEW_CALL_TEMPLATE,
                              template.PYBOOST_CPU_VIEW_CALL_TEMPLATE]
-    code_generate_path = ["mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/auto_generate/",
-                          "mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/auto_generate/",
-                          "mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/auto_generate/"]
+    code_generate_path = [f"{K.MS_OPS_KERNEL_PATH}/{device}/pyboost/auto_generate/" for device in
+                          ["ascend", "gpu", "cpu"]]
     return TemplatePaths(op_header_template_path, op_call_template_path, op_source_template_path,
                          op_custom_template_path,
                          op_view_template_path, code_generate_path)
@@ -692,7 +775,6 @@ class OpTemplateConverter:
             call_args_with_types.append("const " + type_name + " &" + arg_name)
         return call_args_with_types
 
-
     @staticmethod
     def parse_need_malloc_tensors(op_args, call_args):
         """
@@ -716,7 +798,6 @@ class OpTemplateConverter:
                 call_args_with_tensor.append(call_arg)
         return need_malloc_tensors, tensor_list_convert, call_args_with_tensor
 
-
     @staticmethod
     def parse_original_call_args(op_args):
         """
@@ -760,7 +841,7 @@ def delete_residual_files(work_path, all_operator_name, code_generate_path_list)
     """
     Delete residual files.
     """
-    code_generate_path_list.append("mindspore/ccsrc/kernel/pyboost/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:
         all_files_name = []
         code_generate_path = os.path.join(work_path, code_generate_path)
@@ -839,7 +920,8 @@ def gen_pyboost_inner_prim(work_path, op_yaml_data):
             arg_handler = arg_info.get('arg_handler')
             processed_arg = arg_name
             if arg_handler is not None and arg_handler != 'dtype_to_type_id':
-                process_func += f"""converted_{arg_name} = {arg_handler}({arg_name})\n"""
+                process_func += \
+                    f"""converted_{arg_name} = {arg_handler}('{operator_name}', '{arg_name}', {arg_name})\n"""
                 processed_arg = 'converted_' + arg_name
             input_args.append(arg_name)
             processed_args.append(processed_arg)
@@ -952,17 +1034,19 @@ def gen_signature_same_type_table(args_map, operator_data):
     return type_num, signature_table
 
 
-def gen_pyboost_code(work_path, ops_yaml_data, doc_yaml_data):
+def gen_pyboost_code(work_path, ops_yaml_data, doc_yaml_data, extra_ops):
     """ gen_pyboost_code """
     # generate pyboost inner prim
     gen_pyboost_inner_prim(work_path, ops_yaml_data)
     # generate pyboost py func
     gen_pyboost_py_func(work_path, ops_yaml_data, doc_yaml_data)
     # generate ops header file
-    generate_ops_header_files(work_path, ops_yaml_data)
+    generate_ops_header_files(work_path, ops_yaml_data, extra_ops)
     # generate pyboost functions
     generate_pyboost_functions(work_path, ops_yaml_data)
     # generate pyboost grad functions
     generate_pyboost_grad_functions(work_path, ops_yaml_data)
+    # generate pyboost native grad functions
+    generate_pyboost_native_grad_functions(work_path, ops_yaml_data)
     # generate pyboost backend cpp code
     generate_pyboost_op_cpp_code(work_path, ops_yaml_data)

mindspore/ops_generate/op_proto.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """Op Proto."""
-from pyboost_utils import convert_python_func_name_to_c
+from pyboost_utils import convert_python_func_name_to_c, is_op_multi_output
 
 
 class Arg:
@@ -44,7 +44,9 @@ class OpProto:
                  gpu,
                  ascend,
                  prim_init,
-                 is_dispatch):
+                 is_dispatch,
+                 is_multi_output,
+                 is_comm_op):
         self.operator_name = operator_name
         self.class_name = class_name
         self.op_args = op_args
@@ -58,6 +60,8 @@ class OpProto:
         self.ascend = ascend
         self.prim_init = prim_init
         self.is_dispatch = is_dispatch
+        self.is_multi_output = is_multi_output
+        self.is_comm_op = is_comm_op
 
     @staticmethod
     def get_device_special_name(dispatch, gpu, cpu, ascend):
@@ -109,6 +113,7 @@ class OpProto:
 
         is_pyboost = False
         is_dispatch = False
+        is_comm_op = False
         gpu = default_str
         cpu = default_str
         ascend = default_str
@@ -117,6 +122,7 @@ class OpProto:
             is_dispatch = True
             is_pyboost = yaml[dispatch_key].get('enable')
             gpu, cpu, ascend = OpProto.get_device_special_name(yaml[dispatch_key], gpu, cpu, ascend)
+            is_comm_op = yaml[dispatch_key].get('is_comm_op')
         return_dict = yaml['returns']
         class_name = convert_python_func_name_to_c(op_name)
         class_key = 'class'
@@ -130,9 +136,10 @@ class OpProto:
             dtype = return_dict[return_name]['dtype']
             arg = Arg(return_name, dtype, type_cast=[], inplace=inplace)
             return_args.append(arg)
+        is_multi_output = is_op_multi_output(return_args)
         is_view = False
         if 'view' in yaml.keys():
             is_view = True
         op_proto = OpProto(op_name, op_args, return_args, class_name,
-                           is_pyboost, is_view, cpu, gpu, ascend, prim_init, is_dispatch)
+                           is_pyboost, is_view, cpu, gpu, ascend, prim_init, is_dispatch, is_multi_output, is_comm_op)
         return op_proto

mindspore/ops_generate/pyboost_utils.py

@@ -44,6 +44,19 @@ def is_list(op_arg):
     return False
 
 
+def is_op_multi_output(args):
+    """
+    is multi output
+    :param args:
+    :return: bool
+    """
+    if len(args) > 1:
+        return True
+    if len(args) == 1 and is_tensor_list(args[0]):
+        return True
+    return False
+
+
 def get_index(index: int):
     """
     get index
@@ -339,7 +352,7 @@ class AclnnUtils:
     """
     aclnn utils
     """
-    work_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../../")
+    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"))
 
     @staticmethod

mindspore/ops_generate/template.py

@@ -15,6 +15,7 @@
 """Template."""
 import re
 import os
+import gen_constants as K
 
 
 class CppTemplate:
@@ -79,14 +80,14 @@ class CppTemplate:
 
 
 NEW_LINE = "\n"
-WORK_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../../")
+WORK_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../../../")
 
 PYTHON_PRIM_TEMPLATE = CppTemplate("""
 
 class _Pyboost${class_name}Prim(${class_name}Prim_):
     def __call__(self, ${input_args}):
         ${process_func}
-        return _convert_stub(super().__call__(${input_args}))
+        return _convert_stub(super().__call__(${processed_args}))
 
 
 ${func_impl_name}_impl = _Pyboost${class_name}Prim()
@@ -119,6 +120,9 @@ REGISTER_PYBOOST_GRAD_TEMPLATE = CppTemplate("${register_func}")
 PYBOOST_FUNCTION_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH, './mindspore/ccsrc/pipeline/pynative/op_function/template/pyboost_function.tpl'))
 
+PYBOOS_COMM_FUNCTION_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/pipeline/pynative/op_function/template/pyboost_comm_function.tpl'))
+
 PYBOOST_HEADER_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH, './mindspore/ccsrc/pipeline/pynative/op_function/template/pyboost_function_header.tpl'))
 
@@ -126,83 +130,93 @@ PYBOOST_GRAD_FUNCTION_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH, './mindspore/ccsrc/runtime/pynative/op_function/template/pyboost_grad_function.tpl'))
 
 PYBOOST_GRAD_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/runtime/pynative/op_function/template/pyboost_grad_function_header.tpl'))
+
+PYBOOST_NATIVE_GRAD_FUNCTION_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/pipeline/pynative/grad/function/template/native_grad_function.tpl'))
+
+PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/runtime/pynative/op_function/template/pyboost_grad_function_header.tpl'))
+                 './mindspore/ccsrc/pipeline/pynative/grad/function/template/pyboost_native_grad_functions.tpl'))
+
+PYBOOST_NATIVE_GRAD_FUNCTIONS_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/pipeline/pynative/grad/function/template/pyboost_native_grad_functions_header.tpl'))
 
 GEN_OPS_DEF_HEADER_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH, './mindspore/python/mindspore/ops_generate/gen_ops_def_header.tpl'))
 
 PYBOOST_BASE_OP_DEFINE_TEMPLATE = CppTemplate.load_from_file(
-    os.path.join(WORK_PATH, './mindspore/ccsrc/kernel/pyboost/template/pyboost_op_header.tpl'))
+    os.path.join(WORK_PATH, f'./{K.MS_COMMON_PYBOOST_KERNEL_PATH}/template/pyboost_op_header.tpl'))
 
 PYBOOST_OP_REGISTER_TEMPLATE = CppTemplate.load_from_file(
-    os.path.join(WORK_PATH, './mindspore/ccsrc/kernel/pyboost/template/pyboost_op_register.tpl'))
+    os.path.join(WORK_PATH, f'./{K.MS_COMMON_PYBOOST_KERNEL_PATH}/template/pyboost_op_register.tpl'))
 
 # Ascend op generate
 PYBOOST_ASCEND_OP_HEADER_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template/pyboost_aclnn_header_template.tpl'))
+                 f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template/pyboost_aclnn_header_template.tpl'))
 
 PYBOOST_ASCEND_OP_SOURCE_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template/pyboost_aclnn_source_template.tpl'))
+                 f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template/pyboost_aclnn_source_template.tpl'))
 
 PYBOOST_ASCEND_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template/pyboost_ascend_call_template.tpl'))
+                 f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template/pyboost_ascend_call_template.tpl'))
 
 PYBOOST_ASCEND_VIEW_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/kernel/pyboost/template/'
+                 f'./{K.MS_COMMON_PYBOOST_KERNEL_PATH}/template/'
                  'pyboost_view_template.tpl'))
 
 PYBOOST_ASCEND_CUSTOMIZE_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template'
+                 f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template'
                  '/pyboost_ascend_customize_call_template.tpl'))
 
 # GPU op generate
 PYBOOST_GPU_OP_HEADER_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template/pyboost_gpu_header_template.tpl'))
+                 './mindspore/ops/kernel/gpu/pyboost/template/pyboost_gpu_header_template.tpl'))
 
 PYBOOST_GPU_OP_SOURCE_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template/pyboost_gpu_source_template.tpl'))
+                 './mindspore/ops/kernel/gpu/pyboost/template/pyboost_gpu_source_template.tpl'))
 
 PYBOOST_GPU_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template/pyboost_gpu_call_template.tpl'))
+                 './mindspore/ops/kernel/gpu/pyboost/template/pyboost_gpu_call_template.tpl'))
 
 PYBOOST_GPU_VIEW_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/kernel/pyboost/template/pyboost_view_template.tpl'))
+                 f'./{K.MS_COMMON_PYBOOST_KERNEL_PATH}/template/pyboost_view_template.tpl'))
 
 PYBOOST_GPU_CUSTOMIZE_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template'
+                 './mindspore/ops/kernel/gpu/pyboost/template'
                  '/pyboost_gpu_customize_call_template.tpl'))
 
 # CPU op generate
 PYBOOST_CPU_OP_HEADER_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template/pyboost_cpu_header_template.tpl'))
+                 f'./{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/template/pyboost_cpu_header_template.tpl'))
 
 PYBOOST_CPU_OP_SOURCE_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template/pyboost_cpu_source_template.tpl'))
+                 f'./{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/template/pyboost_cpu_source_template.tpl'))
 
 PYBOOST_CPU_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template/pyboost_cpu_call_template.tpl'))
+                 f'./{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/template/pyboost_cpu_call_template.tpl'))
 
 PYBOOST_CPU_VIEW_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/kernel/pyboost/template/pyboost_view_template.tpl'))
+                 f'./{K.MS_COMMON_PYBOOST_KERNEL_PATH}/template/pyboost_view_template.tpl'))
 
 PYBOOST_CPU_CUSTOMIZE_CALL_TEMPLATE = CppTemplate.load_from_file(
     os.path.join(WORK_PATH,
-                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template'
+                 f'./{K.MS_OPS_KERNEL_PATH}/cpu/pyboost/template'
                  '/pyboost_cpu_customize_call_template.tpl'))
 
 PYBOOST_PY_FUNC_IMPORT_HEADEAR = CppTemplate(
@@ -224,7 +238,7 @@ OpDef g${class_name} = {
     ${input_args}
   },
   /* .returns_ = */ {
-    ${return_args} 
+    ${return_args}
   },
   /*.signatures_ =*/ {
     ${signatures}
@@ -236,4 +250,12 @@ OpDef g${class_name} = {
   /*.enable_dispatch_ =*/${enable_dispatch},
   /*.is_view_ =*/${is_view},
 };
+REGISTER_PRIMITIVE_OP_DEF(${class_name}, &g${class_name});
 """)
+
+MULTI_OUTPUT_TEMPLATE = """
+ValuePtrList values;
+(void)std::transform(op->outputs().begin(), op->outputs().end(), std::back_inserter(values),
+                   [](const auto &value){ return value;});
+auto output_value = std::make_shared<ValueTuple>(values);
+"""

mindspore/parallel/__init__.py

@@ -22,7 +22,9 @@ from mindspore.parallel.checkpoint_transform import rank_list_for_transform, tra
     load_segmented_checkpoints
 from mindspore.parallel.parameter_broadcast import parameter_broadcast
 from mindspore.parallel.shard import shard, Layout
+from mindspore.parallel.transform_safetensors import safetensors_to_ckpt, ckpt_to_safetensors, unified_safetensors
 
 __all__ = ["set_algo_parameters", "reset_algo_parameters", "get_algo_parameters", "rank_list_for_transform",
            "transform_checkpoint_by_rank", "transform_checkpoints", "merge_pipeline_strategys", "shard",
-           "sync_pipeline_shared_parameters", "Layout", "parameter_broadcast", "load_segmented_checkpoints"]
+           "sync_pipeline_shared_parameters", "Layout", "parameter_broadcast", "load_segmented_checkpoints",
+           "safetensors_to_ckpt", "ckpt_to_safetensors", "unified_safetensors"]