mindspore 2.4.10__cp310-none-any.whl → 2.5.0__cp310-none-any.whl

mindspore/ops_generate/tensor_func_reg_cpp_generator.py

@@ -0,0 +1,422 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+This module defines the PyboostInnerPrimGenerator class, which is responsible for generating Python primitive
+wrappers for Pyboost operations. The generator constructs Python function definitions based on operator prototypes,
+generates necessary import statements, and writes the generated content into Python source files.
+
+The primary functionality is to take operator prototypes, extract relevant fields, and create Python function wrappers
+that can be used to call the Pyboost primitive implementations.
+"""
+
+import os
+
+import template
+from template import Template
+import gen_constants as K
+import pyboost_utils
+import op_api_proto
+from gen_utils import save_file
+from base_generator import BaseGenerator
+from op_proto import OpProto
+from op_template_parser import OpTemplateParser
+
+
+class TensorFuncRegCppGenerator(BaseGenerator):
+    """
+    Generates C++ tensor function registration code for different backends (Ascend, CPU, GPU).
+
+    This class is responsible for generating header and implementation files required to register
+    tensor functions, including device-specific dispatchers and function definitions.
+    """
+
+    def __init__(self):
+        self.TENSOR_FUNC_CC_REG = template.TENSOR_FUNC_CC_REG
+        self.TENSOR_FUNC_CALL_BODY = template.TENSOR_FUNC_CALL_BODY
+        self.TENSOR_FUNC_OVERLOAD_CALL_BODY = template.TENSOR_FUNC_OVERLOAD_CALL_BODY
+        self.TENSOR_API_HEADER = template.TENSOR_API_HEADER
+        self.TENSOR_API_SOURCE = template.TENSOR_API_SOURCE
+        self.TENSOR_FUNC_UTILS = template.TENSOR_FUNC_UTILS
+        self.TENSOR_FUNC_UT_BODY = template.TENSOR_FUNC_UT_BODY
+        self.TENSOR_FUNC_UT_OVERLOAD_BODY = template.TENSOR_FUNC_UT_OVERLOAD_BODY
+        self.TENSOR_CPP_METHOD = template.TENSOR_CPP_METHOD
+
+        self.func_def_reg = Template("tensor_class->def(\"${func_name}\", TensorMethod${cpp_func_name});\n")
+        self.single_case_template = Template(
+            'case ${case_id}:\n'
+            '  ${device_dispatcher}\n'
+            '  break;\n'
+        )
+        self.single_case_in_ut_template = Template(
+            'case ${case_id}:\n'
+            '  ${device_dispatcher}\n'
+        )
+        self.device_dispatcher_template = Template(
+            'if (backend == kAscendDevice || backend == kDavinciDevice) {\n'
+            '  ${ascend_dispatcher}\n'
+            '} else if (backend == kCPUDevice) {\n'
+            '  ${cpu_dispatcher}\n'
+            '} else if (backend == kGPUDevice) {\n'
+            '  ${gpu_dispatcher}\n'
+            '} else {\n'
+            '  MS_LOG(ERROR) << "Device target is not supported!";\n'
+            '  return py::none();\n'
+            '}'
+        )
+        self.pyboost_return_template = Template(
+            '${arg_handler_processor}\n'
+            'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
+            'return ToPython(TensorPyboostMethodRegister::'
+            'GetOp(tensor::TensorPyboostMethod::k${class_name}Reg)(arg_list));\n'
+        )
+        self.callback_python_template = Template(
+            'MS_LOG(INFO) << "Callback python method: ${py_method}";\n'
+            'py::function fn = python_adapter::GetPyFn(\"mindspore.ops.tensor_method\", \"${py_method}\");\n'
+            'py::object res = fn(self, *py_args, **py_kwargs);\n'
+            'return res;\n'
+        )
+        self.callback_python_in_ut_template = Template(
+            'MS_LOG(INFO) << "Callback python method in UT: ${py_method}";\n'
+            'fn = python_adapter::GetPyFn(\"mindspore.ops.tensor_method\", \"${py_method}\");\n'
+            'res = fn(self, *py_args, **py_kwargs);\n'
+            'break;\n'
+        )
+        self.header_func_header_template = Template(
+            "py::object TensorMethod${cpp_func_name}"
+            "(const py::object &self, const py::args &py_args, const py::kwargs &py_kwargs);\n"
+        )
+
+    def generate(self, work_path, op_protos, func_protos_data, alias_func_mapping):
+        """
+        Generates C++ header and source files for tensor function registrations.
+
+        Args:
+            work_path (str): The directory where the generated files will be saved.
+            op_protos (list): A list of tensor op prototypes.
+            func_protos_data (dict): Dictionary mapping function names to lists of TensorFuncProto objects.
+            alias_func_mapping (dict): A dictionary mapping function name to its alias function names.
+        """
+
+        all_op_func_data, single_op_func_data, overload_op_func_data = op_api_proto.categorize_func_data(
+            func_protos_data)
+
+        tensor_method_list = self._get_op_enum_name_list(op_protos)
+        func_call_body_list = []
+        self._create_single_op_source_files(single_op_func_data, func_call_body_list)
+        self._create_overload_op_source_files(overload_op_func_data, func_call_body_list)
+        merge_func_call_body = pyboost_utils.merge_strings_by_chunk_size(func_call_body_list)
+
+        for i, func_body_chunk_str in enumerate(merge_func_call_body):
+            tensor_api_source = self.TENSOR_API_SOURCE.replace(tenosr_func_call_body=func_body_chunk_str)
+            save_file(os.path.join(work_path, K.TENSOR_API_PATH), f"tensor_api_{i}.cc",
+                      tensor_api_source)
+
+        func_def_body_list, tensor_cpp_methods_list, tensor_api_declaration_list = self._get_sorted_func_def_body(
+            all_op_func_data, alias_func_mapping)
+        tensor_api_header = self.TENSOR_API_HEADER.replace(tensor_api_declaration_list=tensor_api_declaration_list)
+        save_file(os.path.join(work_path, K.TENSOR_API_PATH), f"tensor_api.h",
+                  tensor_api_header)
+        self._generate_func_name_for_stub_tensor(work_path, tensor_cpp_methods_list)
+        func_cc_reg = self.TENSOR_FUNC_CC_REG.replace(func_def_body=func_def_body_list)
+        tensor_methods = self.TENSOR_FUNC_UTILS.replace(tensor_methods=tensor_method_list)
+
+        save_file(os.path.join(work_path, K.TENSOR_FUNC_REGISTER_PATH), "tensor_func_utils.h", tensor_methods)
+        save_file(os.path.join(work_path, K.TENSOR_FUNC_REGISTER_PATH), "tensor_func_reg.cc", func_cc_reg)
+
+    def _get_op_enum_name_list(self, op_protos):
+        """
+        Extracts operation class names and returns them as a formatted list.
+
+        Args:
+            op_protos (list): A list of operation prototype objects, where each object has an `op_class`
+                              with a `name` attribute.
+
+        Returns:
+            str: A list of formatted strings, where each string is of the form 'k<name>,\n', where <name>
+                  is the class name from the `op_class` attribute.
+
+        """
+        tensor_method_list = ""
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None or not op_proto.op_dispatch.enable:
+                continue
+            class_name = op_proto.op_class.name
+            tensor_method_list += f"k{class_name}Reg,\n"
+        return tensor_method_list
+
+    def _generate_func_name_for_stub_tensor(self, work_path, tensor_cpp_methods_list):
+        """
+        Generates a Python file containing tensor C++ function methods list and saves it to the specified path.
+
+        This function takes a list of C++ tensor methods, formats them into a Python script as a string,
+        and writes this script to a file named `_tensor_cpp_method.py` under the provided working path.
+
+        Args:
+            work_path (str): The base directory where the generated file will be saved.
+            tensor_cpp_methods_list (list): A list of tensor C++ method definitions to be included in the Python file.
+        """
+        tensor_cpp_methods_str = self.TENSOR_CPP_METHOD.replace(
+            tensor_cpp_methods_list_str=str(tensor_cpp_methods_list))
+        save_file(os.path.join(work_path, K.ADD_TENSOR_DOCS_PY_PATH), "_tensor_cpp_method.py", tensor_cpp_methods_str)
+
+    def _get_sorted_func_def_body(self, all_op_func_data, alias_func_mapping):
+        """
+        Generate sorted function definitions and headers for operations.
+
+        This function processes a dictionary of operation function data and an alias mapping,
+        producing two lists: one containing function definition bodies and another containing
+        function header definitions.
+
+        Args:
+            all_op_func_data (dict): A dictionary where keys are function API names (str), and
+                values are lists of function prototypes.
+            alias_func_mapping (dict): A mapping of function names to a list of their alias names.
+
+        Returns:
+            tuple: A tuple containing two lists:
+                - func_def_body_list (list of str): A list of formatted function definition strings.
+                - tensor_cpp_methods_list (list of str): A list of formatted function header strings.
+        """
+        func_def_body_list = []
+        tensor_cpp_methods_list = []
+        tensor_api_declaration_list = ""
+        for func_api_name, func_protos in all_op_func_data.items():
+            cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+            if len(func_protos) == 1:
+                func_proto = func_protos[0]
+                func_name = func_proto.func_name
+                func_def_body_list.append(self.func_def_reg.replace(func_name=func_name, cpp_func_name=cpp_func_name))
+                tensor_cpp_methods_list.append(func_name)
+                tensor_api_declaration_list += self.header_func_header_template.replace(cpp_func_name=cpp_func_name)
+                if func_name in alias_func_mapping:
+                    for alias_func_name in alias_func_mapping[func_name]:
+                        func_def_body_list.append(
+                            self.func_def_reg.replace(func_name=alias_func_name, cpp_func_name=cpp_func_name))
+                        tensor_cpp_methods_list.append(alias_func_name)
+            elif len(func_protos) > 1:
+                func_def_body_list.append(
+                    self.func_def_reg.replace(func_name=func_api_name, cpp_func_name=cpp_func_name))
+                tensor_cpp_methods_list.append(func_api_name)
+                tensor_api_declaration_list += self.header_func_header_template.replace(cpp_func_name=cpp_func_name)
+                if func_api_name in alias_func_mapping:
+                    for alias_func_name in alias_func_mapping[func_api_name]:
+                        func_def_body_list.append(self.func_def_reg.replace(func_name=alias_func_name,
+                                                                            cpp_func_name=cpp_func_name))
+                        tensor_cpp_methods_list.append(alias_func_name)
+        return func_def_body_list, tensor_cpp_methods_list, tensor_api_declaration_list
+
+    def _create_single_op_source_files(self, single_op_func_data, func_call_body_list):
+        """
+        Generates the list of call body strings for single operation functions.
+
+        Args:
+            single_op_func_data (dict): Dictionary of tensor function prototypes with only one definition.
+
+        Returns:
+            list: Updated str list for generating C++ function call bodies.
+        """
+        for func_api_name, func_proto in single_op_func_data.items():
+            func_name = func_proto.func_name
+            cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+            device_dispatcher_str = self._get_device_dispatchers_str(func_proto)
+            signature_str = self._generate_single_signature_str(
+                func_proto.op_proto, func_proto.kw_only_args, func_proto.varargs
+            )
+            op_args = func_proto.op_proto.op_args
+            max_size = len(op_args)
+            self_index = self._get_input_tensor_index(func_proto)
+            ut_body = self.TENSOR_FUNC_UT_BODY.replace(py_method=func_proto.py_method)
+            tensor_func_single_call_body = self.TENSOR_FUNC_CALL_BODY.replace(cpp_func_name=cpp_func_name,
+                                                                              func_name=func_name,
+                                                                              device_dispatcher=device_dispatcher_str,
+                                                                              signatures=signature_str,
+                                                                              max_args=max_size,
+                                                                              self_index=self_index,
+                                                                              ut_body=ut_body)
+            func_call_body_list.append(tensor_func_single_call_body)
+
+    def _create_overload_op_source_files(self, overload_op_func_data, func_call_body_list):
+        """
+        Generates the list of call body strings for overloaded operation functions.
+
+        Args:
+            overload_op_func_data (dict): Dictionary of tensor function prototypes with overloaded definitions.
+
+        Returns:
+            list: Updated str list for generating C++ function call bodies.
+        """
+        for func_api_name, func_protos in overload_op_func_data.items():
+            tensor_func_overload_call_body = self._get_overload_func_call_str(func_api_name, func_protos)
+            func_call_body_list.append(tensor_func_overload_call_body)
+
+    def _get_overload_func_call_str(self, func_api_name, func_protos):
+        """
+        Generates C++ call body string for overloaded tensor functions.
+
+        Args:
+            func_api_name (str): Name of the function API.
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated call body string for the overloaded functions.
+        """
+        signatures_str = self._generate_func_signatures_list_str(func_protos)
+        dispatch_cases = self._get_dispatch_cases(func_protos)
+        ut_dispatch_cases = self._get_ut_dispatch_cases(func_protos)
+        ut_overload_body = self.TENSOR_FUNC_UT_OVERLOAD_BODY.replace(ut_dispatch_cases=ut_dispatch_cases)
+
+        max_size = 0
+        self_index = 0
+        for tensor_proto in func_protos:
+            op_proto = tensor_proto.op_proto
+            op_args = op_proto.op_args
+            max_size = max(len(op_args), max_size)
+            self_index = self._get_input_tensor_index(tensor_proto)
+        cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+        overload_func_call_str = self.TENSOR_FUNC_OVERLOAD_CALL_BODY.replace(cpp_func_name=cpp_func_name,
+                                                                             func_name=func_api_name,
+                                                                             signatures=signatures_str,
+                                                                             dispatch_cases=dispatch_cases,
+                                                                             max_args=max_size,
+                                                                             self_index=self_index,
+                                                                             ut_overload_body=ut_overload_body)
+        return overload_func_call_str
+
+    def _generate_func_signatures_list_str(self, func_protos) -> str:
+        """
+        Generates function signatures as a string from the given prototypes.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated function signatures string.
+        """
+        sig_str = ''
+        first_sig = True
+        for tensor_proto in func_protos:
+            op_proto = tensor_proto.op_proto
+            if not first_sig:
+                sig_str += ',\n'
+            first_sig = False
+            sig_str += self._generate_single_signature_str(op_proto, tensor_proto.kw_only_args, tensor_proto.varargs)
+        return sig_str
+
+    def _generate_single_signature_str(self, op_proto: OpProto, kw_only_args, varargs) -> str:
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.generate_signature_str(kw_only_args, varargs, is_tensor_api=True)
+
+    def _get_input_tensor_index(self, func_proto):
+        """
+        Get index of input.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate dispatch strings for.
+
+        Returns:
+            int: Index of input.
+        """
+        op_name = func_proto.op_proto.op_class.name
+        op_args = func_proto.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
+
+    def _get_dispatch_cases(self, func_protos):
+        """
+        Generates C++ switch-case statements for dispatching tensor function calls.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated switch-case dispatch statements.
+        """
+        dispatch_cases_str = ''
+        for idx, func_proto in enumerate(func_protos):
+            device_dispatcher_str = self._get_device_dispatchers_str(func_proto)
+            dispatch_cases_str += self.single_case_template.replace(case_id=idx,
+                                                                    device_dispatcher=device_dispatcher_str)
+        dispatch_cases_str += 'default:\n'
+        dispatch_cases_str += '  return py::none();'
+        return dispatch_cases_str
+
+    def _get_ut_dispatch_cases(self, func_protos):
+        """
+        Generates C++ switch-case statements for dispatching tensor function calls.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated switch-case dispatch statements.
+        """
+        dispatch_cases_str = ''
+        for idx, func_proto in enumerate(func_protos):
+            device_dispatcher_str = self.callback_python_in_ut_template.replace(py_method=func_proto.py_method)
+            dispatch_cases_str += self.single_case_in_ut_template.replace(case_id=idx,
+                                                                          device_dispatcher=device_dispatcher_str)
+        dispatch_cases_str += 'default:\n'
+        dispatch_cases_str += '  res = py::none();'
+        return dispatch_cases_str
+
+    def _get_device_dispatchers_str(self, func_proto):
+        """
+        Generates device-specific dispatch strings for the given function prototype.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate dispatch strings for.
+
+        Returns:
+            str: Generated device-specific dispatch string.
+        """
+        ascend_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'ascend')
+        cpu_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'cpu')
+        gpu_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'gpu')
+        device_dispatcher_str = self.device_dispatcher_template.replace(ascend_dispatcher=ascend_dispatcher_str,
+                                                                        cpu_dispatcher=cpu_dispatcher_str,
+                                                                        gpu_dispatcher=gpu_dispatcher_str)
+        return device_dispatcher_str
+
+    def _get_single_device_dispatcher_str(self, func_proto, device):
+        """
+        Generates the dispatch string for a specific device.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate the dispatcher for.
+            device (str): Device type ('ascend', 'cpu', 'gpu').
+
+        Returns:
+            str: Generated device dispatcher string.
+        """
+        func_proto_device = getattr(func_proto, device)
+        if func_proto_device == 'pyboost':
+            arg_handler_processor_str = self._get_arg_handler_processor(func_proto.func_name, func_proto.op_proto)
+            return self.pyboost_return_template.replace(arg_handler_processor=arg_handler_processor_str,
+                                                        class_name=func_proto.op_proto.op_class.name)
+
+        if func_proto_device == 'py_method':
+            return self.callback_python_template.replace(py_method=func_proto.py_method)
+
+        raise TypeError("Only support pyboost or python_method.")
+
+    def _get_arg_handler_processor(self, func_name, op_proto):
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.get_arg_handler_processor(func_name, op_proto, is_tensor_api=True)

mindspore/parallel/__init__.py

@@ -19,12 +19,12 @@ from mindspore.parallel.algo_parameter_config import get_algo_parameters, reset_
     set_algo_parameters
 from mindspore.parallel.checkpoint_transform import rank_list_for_transform, transform_checkpoint_by_rank, \
     transform_checkpoints, merge_pipeline_strategys, sync_pipeline_shared_parameters, \
-    load_segmented_checkpoints
+    load_segmented_checkpoints, set_op_strategy_config
 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
+from mindspore.parallel.transform_safetensors import 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",
-           "safetensors_to_ckpt", "ckpt_to_safetensors", "unified_safetensors"]
+           "unified_safetensors", "set_op_strategy_config"]

mindspore/parallel/_auto_parallel_context.py

@@ -187,6 +187,25 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_dump_local_norm()
 
+    def set_dump_local_norm_path(self, dump_local_norm_path):
+        """
+        Set dump local norm path for auto parallel.
+
+        Args:
+            dump_local_norm_path (str): User need to specify the path to save dump files
+                                        if he want to dump local norm.  Default: ''
+
+        Raises:
+            KeyError: When the value of dump_local_norm_path is not a str value.
+        """
+        self.check_context_handle()
+        self._context_handle.set_dump_local_norm_path(dump_local_norm_path)
+
+    def get_dump_local_norm_path(self):
+        """Get dump local norm path."""
+        self.check_context_handle()
+        return self._context_handle.get_dump_local_norm_path()
+
     def set_dump_device_local_norm(self, dump_device_local_norm):
         """
         Set dump device local norm for auto parallel.
@@ -195,7 +214,7 @@ class _AutoParallelContext:
             dump_device_local_norm (bool): User need to specify if he want to dump device local norm.  Default: False
 
         Raises:
-            ValueError: If the dump_device_local_norm in not a bool value.
+            ValueError: If the dump_device_local_norm is not a bool value.
         """
         self.check_context_handle()
         self._context_handle.set_dump_device_local_norm(dump_device_local_norm)
@@ -1089,13 +1108,6 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_optimizer_weight_shard_size()
 
-    def set_ops_strategy_json_config(self, type, path, mode):
-        """
-         Set configuration of saving ops strategy in file .json.
-        """
-        self.check_context_handle()
-        self._context_handle.set_ops_strategy_json_config(type, path, mode)
-
     def set_optimizer_weight_shard_aggregated_save(self, optimizer_weight_shard_aggregated_save):
         """
         Set optimizer_weight_shard_aggregated_save.
@@ -1233,31 +1245,6 @@ class _AutoParallelContext:
             self.set_enable_all_gather_fusion(openstate)
             self.set_enable_reduce_scatter_fusion(openstate)
 
-
-def _set_ops_strategy_json_config(type="SAVE", path="", mode="all"):
-    """
-    Set strategy json configuration.
-
-    Args:
-        type (str): The parameter for choosing save or load .json file.
-        path (str): Path to save or load parallel strategy json.
-        mode (str): The parameter for choosing save all or important operators.
-
-    Raises:
-        KeyError: When type is not 'SAVE' or 'LOAD'.
-        KeyError: When mode is not 'all' or 'principal'.
-    """
-    dir_path = os.path.dirname(path)
-    if dir_path and not os.path.exists(dir_path):
-        os.makedirs(dir_path, mode=0o700, exist_ok=True)
-    check_type = ["SAVE", "LOAD"]
-    check_mode = ["all", "principal"]
-    if type in check_type and mode in check_mode:
-        auto_parallel_context().set_ops_strategy_json_config(type, path, mode)
-    else:
-        raise KeyError("Type must be 'SAVE' or 'LOAD' and mode must be 'all' or 'principal'")
-
-
 _AUTO_PARALLEL_CONTEXT = None
 
 
@@ -1307,6 +1294,7 @@ _set_auto_parallel_context_func_map = {
     "strategy_ckpt_config": auto_parallel_context().set_strategy_ckpt_config,
     "comm_fusion": auto_parallel_context().set_comm_fusion,
     "dump_local_norm": auto_parallel_context().set_dump_local_norm,
+    "dump_local_norm_path": auto_parallel_context().set_dump_local_norm_path,
     "dump_device_local_norm": auto_parallel_context().set_dump_device_local_norm}
 
 _get_auto_parallel_context_func_map = {
@@ -1341,6 +1329,7 @@ _get_auto_parallel_context_func_map = {
     "strategy_ckpt_config": auto_parallel_context().get_strategy_ckpt_config,
     "full_batch_is_set": auto_parallel_context().get_full_batch_is_set,
     "dump_local_norm": auto_parallel_context().get_dump_local_norm,
+    "dump_local_norm_path": auto_parallel_context().get_dump_local_norm_path,
     "dump_device_local_norm": auto_parallel_context().get_dump_device_local_norm}
 
 
@@ -1452,6 +1441,8 @@ def _set_auto_parallel_context(**kwargs):
                     - reducescatter: If communication fusion type is `reducescatter`. The `mode` contains: `auto`
                         and `size`. Config is same as `allgather`.
 
+
+
     Raises:
         ValueError: If input key is not attribute in auto parallel context.
     """

mindspore/parallel/_parallel_serialization.py

@@ -115,11 +115,15 @@ def _check_strategy_file(strategy_filename):
                          f"be empty. Please check whether the 'strategy_filename' is correct.")
 
 
-def _load_protobuf_strategy(strategy_filename):
+def _load_protobuf_strategy(strategy_filename, strategy_set=None):
     """load strategy from protobuf file"""
     parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
     with open(strategy_filename, 'rb') as f:
         pb_content = f.read()
+        if strategy_set is not None:
+            if pb_content in strategy_set:
+                return {}
+            strategy_set.add(pb_content)
     try:
         parallel_strategy_map.ParseFromString(pb_content)
     except BaseException as e:
@@ -188,8 +192,11 @@ def _merge_protobuf_strategy(src_strategy_files, dst_strategy_file):
     """merge protobuf strategy"""
     dst_parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
     merged_stage = []
+    strategy_set = set()
     for src_strategy_file in src_strategy_files:
-        src_parallel_strategy_map = _load_protobuf_strategy(src_strategy_file)
+        src_parallel_strategy_map = _load_protobuf_strategy(src_strategy_file, strategy_set=strategy_set)
+        if not src_parallel_strategy_map:
+            continue
         strategy_items = src_parallel_strategy_map.parallel_strategy_item
         layout_items = src_parallel_strategy_map.parallel_layout_item
         if not strategy_items or not layout_items:
@@ -339,6 +346,9 @@ def _get_device_num_from_strategy(strategy_file=None):
         src_strategy = strategy_file
     strategy_list = _convert_to_list(src_strategy)
     device_mat = list(strategy_list.values())[0][0]
+    if not device_mat:
+        raise ValueError("The parallel strategy file only contains pipeline-parallelism, which is not supported for "
+                         "parallel strategy conversion now.")
     return np.prod(device_mat)
 
 
@@ -559,3 +569,4 @@ def _get_param_list_when_first_dim_sharded(device_arrangement, first_dim_sharded
     start = rank - offset
     param_total_list = list(range(start, start + range_size))
     return param_total_list
+            

mindspore/parallel/_utils.py

@@ -30,7 +30,7 @@ from mindspore.common.seed import get_seed
 from mindspore._c_expression import GraphExecutor_
 from mindspore.parallel._tensor import _load_tensor_by_layout, _load_tensor_shape_by_layout
 
-SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore", "Custom"]
+SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore", "Custom", "Index"]
 
 
 def _get_parallel_mode():
@@ -135,6 +135,9 @@ def _slice_parameter(parameter, phase, layout):
             rank = get_rank()
             new_tensor_shape = _load_tensor_shape_by_layout(parameter, layout, rank)
             parameter.shape = new_tensor_shape
+            if hasattr(parameter.init_mode, "shape") and parameter.init_mode.shape != parameter.shape:
+                parameter.init_mode.shape = new_tensor_shape
+            parameter.sliced = True
     else:
         graph_executor = GraphExecutor_.get_instance()
         new_param = parameter.init_data(layout, set_sliced=True)

mindspore/parallel/algo_parameter_config.py

@@ -370,7 +370,7 @@ def get_algo_parameters(attr_key):
     Examples:
         >>> import mindspore as ms
         >>> ms.get_algo_parameters("fully_use_devices")
-        True
+        False
     """
     if attr_key not in get_algo_parameters_config_func_map:
         raise ValueError("Get context keyword %s is not recognized!" % attr_key)

mindspore/parallel/checkpoint_transform.py

@@ -28,6 +28,7 @@ from mindspore.parallel._parallel_serialization import _rank_list_for_transform_
     _extract_layout_map, _extract_src_dst_layout_map, _parameter_not_in_local_stage, _extract_pipeline_stage_num, \
     _merge_protobuf_strategy, _merge_json_strategy, _extract_src_dst_layout_map_by_src
 from mindspore.parallel.transform_safetensors import _transform_safetensors, _collect_safetensor_files
+from mindspore._c_expression import AutoParallelContext
 
 __all__ = ["merge_pipeline_strategys", "rank_list_for_transform", "transform_checkpoint_by_rank",
            "transform_checkpoints", "sync_pipeline_shared_parameters", "load_segmented_checkpoints"]
@@ -648,3 +649,46 @@ def load_segmented_checkpoints(ckpt_file_dir, net=None, strict_load=False, filte
         parameter_dict.update(ms.load_checkpoint(checkpoint_file, net, strict_load, filter_prefix, dec_key,
                                                  dec_mode, specify_prefix, choice_func))
     return parameter_dict
+
+
+def set_op_strategy_config(mode="SAVE", path=""):
+    """
+    Set strategy json configuration when using sharding propagation.
+
+    .. warning::
+        This is an experimental interface, may be changed or canceled in the future;
+        This interface currently doesn't support saving or loading strategies using layout.
+
+    Note:
+        - It only works when `parallel_mode=ParallelMode.AUTO_PARALLEL` and `search_mode='sharding_propagation'`.
+        - It only supports saving and reloading with the same configuration for the same network. If the network
+          or training hyperparameters are modified after using the `SAVE` mode to save the strategies of operator
+          to the setting json file, which may lead to the failure of using the `LOAD` mode to load operator
+          strategies from json.
+        - When performing distributed training, users can first save the strategy using dryrun on a single device
+          and then load strategy to perform distributed training.
+
+    Args:
+        mode (str): The parameter for choosing save or load .json file. Default value: ``"SAVE"`` .
+        path (str): Path to save or load parallel strategy json, must be an absolute path. Default value: ``""`` .
+
+    Raises:
+        KeyError: When type is not ``"SAVE"`` or ``"LOAD"`` .
+        KeyError: When path does not end in ``".json"`` .
+        KeyError: When path is not an absolute path.
+    """
+    if not os.path.isabs(path):
+        raise KeyError("File path must be an absolute path")
+    _, file_type = os.path.splitext(path)
+    if file_type != ".json":
+        raise KeyError("File type must be .json")
+    dir_path = os.path.dirname(path)
+    if dir_path and not os.path.exists(dir_path):
+        os.makedirs(dir_path, mode=0o700, exist_ok=True)
+    check_mode_type = ["SAVE", "LOAD"]
+    if mode in check_mode_type:
+        if AutoParallelContext.get_instance() is None:
+            raise ValueError("Get AutoParallelContext instance failed!!!")
+        AutoParallelContext.get_instance().set_ops_strategy_json_config(mode, path, "all")
+    else:
+        raise KeyError("Type must be 'SAVE' or 'LOAD'")