mindspore 2.2.11__cp39-cp39-macosx_11_0_arm64.whl → 2.4.0__cp39-cp39-macosx_11_0_arm64.whl

mindspore/.commit_id

@@ -1 +1 @@
-__commit_id__ = ''[sha1]:8c390933,[branch]:(HEAD,origin/r2.2,r2.2)''
+__commit_id__ = ''[sha1]:8c86f33f,[branch]:(HEAD,origin/master,origin/HEAD,master)''

mindspore/__init__.py

@@ -17,25 +17,29 @@ from __future__ import absolute_import
 
 from mindspore.run_check import run_check
 from mindspore import common, dataset, mindrecord, train, log, amp
-from mindspore import profiler, communication, numpy, parallel
+from mindspore import profiler, communication, numpy, parallel, hal
 from mindspore.common import *
 from mindspore.mindrecord import *
-from mindspore.ops import _op_impl, grad, value_and_grad, vjp, jvp, jacfwd, jacrev, vmap, get_grad, constexpr
+from mindspore.ops import _op_impl, grad, value_and_grad, vjp, jvp, jacfwd, jacrev, vmap, get_grad, constexpr, reshard
 from mindspore.train import *
 from mindspore.log import *
+from mindspore.utils import *
 from mindspore.context import GRAPH_MODE, PYNATIVE_MODE, set_context, get_context, set_auto_parallel_context, \
     get_auto_parallel_context, reset_auto_parallel_context, ParallelMode, set_ps_context, \
     get_ps_context, reset_ps_context, set_offload_context, get_offload_context, STRICT, COMPATIBLE, LAX
 from mindspore.version import __version__
 from mindspore.profiler import Profiler, EnvProfiler
 from mindspore.parallel import set_algo_parameters, get_algo_parameters, reset_algo_parameters, \
-    rank_list_for_transform, transform_checkpoint_by_rank, transform_checkpoints, merge_pipeline_strategys, shard
-from mindspore.rewrite import SymbolTree, ScopedValue, Node, NodeType, TreeNodeHelper
+    rank_list_for_transform, transform_checkpoint_by_rank, transform_checkpoints, merge_pipeline_strategys, shard, \
+    Layout, sync_pipeline_shared_parameters, parameter_broadcast, load_segmented_checkpoints, \
+    safetensors_to_ckpt, ckpt_to_safetensors, unified_safetensors
+from mindspore.rewrite import SymbolTree, ScopedValue, Node, NodeType
 from mindspore.safeguard import obfuscate_ckpt, load_obf_params_into_net
 from mindspore._check_jit_forbidden_api import get_obj_module_and_name_info, is_jit_forbidden_module, \
     is_invalid_or_jit_forbidden_method
 from mindspore import _install_custom
-
+from mindspore import mint
+from mindspore.ops_generate import arg_dtype_cast, arg_handler
 
 __all__ = ["run_check"]
 __all__.extend(__version__)

mindspore/_checkparam.py

@@ -18,6 +18,7 @@ from __future__ import absolute_import
 import re
 import inspect
 import math
+from types import FunctionType, MethodType
 from functools import reduce, wraps
 from itertools import repeat
 from collections.abc import Iterable
@@ -28,7 +29,6 @@ from mindspore import log as logger
 from mindspore.common import dtype as mstype
 from mindspore._c_expression import Tensor as Tensor_
 
-
 EQ = 1  # ==
 NE = 2  # !=
 LT = 3  # <
@@ -147,7 +147,7 @@ def _check_3d_int_or_tuple(arg_name, arg_value, prim_name, allow_five=False, ret
             ret = (1, 1, arg_value, arg_value, arg_value) if ret_five else (arg_value, arg_value, arg_value)
         elif len(arg_value) == 3:
             ret = (1, 1, arg_value[0], arg_value[1], arg_value[2]) if ret_five else arg_value
-        else: # case: len(arg_value) == 5
+        else:  # case: len(arg_value) == 5
             ret = arg_value if ret_five else (arg_value[2], arg_value[3], arg_value[4])
 
         return ret
@@ -239,6 +239,7 @@ def check_is_number(arg_value, arg_type, arg_name=None, prim_name=None):
         else:
             raise TypeError(f"{prim_name} type of {arg_name} must be '{arg_type.__name__}', " \
                             f"but got '{type(arg_value).__name__}'.")
+
     _check_param()
     return arg_value
 
@@ -264,6 +265,7 @@ def check_number_range(arg_value, lower_limit, upper_limit, rel, value_type, arg
             rel_str = _format_str_two_value(lower_limit, upper_limit, rel)
             raise ValueError(f"{prim_name} {arg_name} must be in range of {rel_str}, " \
                              f"but got {arg_value} with type '{type(arg_value).__name__}'.")
+
     _check_param()
     return arg_value
 
@@ -273,6 +275,7 @@ def check(arg_name, arg_value, value_name, value, rel=EQ, prim_name=None, excp_c
     Method for judging relation between two int values or list/tuple made up of ints.
     This method is not suitable for judging relation between floats, since it does not consider float error.
     """
+
     def _check():
         if not _check_binary_rel(arg_value, value, rel):
             rel_str = _format_str_one_value(f'{value_name}: {value}', rel)
@@ -474,20 +477,24 @@ def check_non_negative_float(arg_value, arg_name=None, prim_name=None):
 
 def check_number(arg_name, arg_value, value, rel, prim_name):
     """Number value judgment."""
+
     def _check():
         if not _check_binary_rel(arg_value, value, rel):
             rel_str = _format_str_one_value(value, rel)
             raise ValueError(f'For \'{prim_name}\', the argument \'{arg_name}\' ' \
                              f'must {rel_str}, but got {arg_value}.')
+
     _check()
     return arg_value
 
 
 def check_isinstance(arg_name, arg_value, classes):
     """Check arg isinstance of classes"""
+
     def _check():
         if not isinstance(arg_value, classes):
             raise ValueError(f'The parameter \'{arg_name}\' must be isinstance of {classes}, but got {arg_value}.')
+
     _check()
     return arg_value
 
@@ -506,6 +513,7 @@ def check_bool(arg_value, arg_name=None, prim_name=None):
     def _check():
         if not isinstance(arg_value, bool):
             raise TypeError(f"{prim_name} {arg_name} must be a bool, but got {type(arg_value).__name__}.")
+
     _check()
     return arg_value
 
@@ -546,6 +554,7 @@ def check_string(arg_value, valid_values, arg_name=None, prim_name=None):
         if not (isinstance(arg_value, str) and arg_value in valid_values):
             raise ValueError(f"{msg_prefix} '{arg_name}' must be str and must be in '{valid_values}'," \
                              f" but got '{arg_value}'.")
+
     _check()
     return arg_value
 
@@ -625,10 +634,12 @@ def check_subclass(arg_name, type_, template_types, prim_name, addition_error_in
 
 def check_valid_input(arg_name, arg_value, prim_name):
     """Checks valid value."""
+
     def _check():
         if arg_value is None:
             raise ValueError(f"For \'{prim_name}\', the argument '{arg_name}'" \
                              f"can not be None, but got {arg_value}.")
+
     _check()
     return arg_value
 
@@ -785,6 +796,7 @@ def check_astype_dtype(dtype):
 
 def check_transpose_axis(axes, ndim):
     """Check the axis argument for tensor.transpose"""
+
     def _check_dim():
         # if multiple arguments provided, it must be `ndim` number of ints
         if len(axes) != ndim:
@@ -792,7 +804,7 @@ def check_transpose_axis(axes, ndim):
                              f"but got {len(axes)} in the number of axes.")
 
     if not axes or (len(axes) == 1 and axes[0] is None):
-        return tuple(range(ndim-1, -1, -1))
+        return tuple(range(ndim - 1, -1, -1))
 
     if len(axes) == 1:
         perm = axes[0]
@@ -911,6 +923,7 @@ def prepare_shape_for_squeeze(shape, axes):
 
 def check_axis_in_range(axis, ndim):
     """Checks axes are with the bounds of ndim"""
+
     def _check():
         if not isinstance(axis, int):
             raise TypeError(f'The axes must be integers, but got {type(axis)}')
@@ -927,6 +940,7 @@ def check_axis_valid(axes, ndim):
     Checks axes are valid given ndim, and returns axes that can be passed
     to the built-in operator (non-negative, int or tuple)
     """
+
     def _check_range(axes):
         for axis in axes:
             check_axis_in_range(axis, ndim)
@@ -976,16 +990,17 @@ def infer_out_shape(*shapes):
     """
     Returns shape of output after broadcasting. Raises ValueError if shapes cannot be broadcast.
     """
+
     def _check(items, max_size, shapes):
         for item in items:
             if item not in (1, max_size):
                 raise ValueError(f'For Tensor, the dimension on each axis must be 1 or the max value on the axis' \
                                  f'to support broadcasting, but got shapes {shapes,}')
+
     shape_out = ()
     max_len = max([len(it) for it in shapes])
     for i in range(max_len):
-        items = [it[i-(max_len-len(it))] if i - (max_len - len(it))
-                 >= 0 else 1 for it in shapes]
+        items = [it[i - (max_len - len(it))] if i - (max_len - len(it)) >= 0 else 1 for it in shapes]
         max_size = 0 if 0 in items else max(items)
         _check(items, max_size, shapes)
         shape_out = shape_out + (max_size,)
@@ -1014,6 +1029,7 @@ def check_axis_type(axis, type_int=True, type_tuple=True, type_list=True):
 
 def check_and_canonicalize_axes(axes, ndim):
     """Check whether the types and values of input axes are valid."""
+
     def _check(axes, ax, ndim):
         if not isinstance(ax, int):
             raise TypeError(f"Each axis should be integer, but got {type(ax)} in {axes}.")
@@ -1090,8 +1106,8 @@ def check_csr_tensor_shape(indptr_shp, indices_shp, values_shp, csr_shp):
                          f"{len(csr_shp)}")
     if values_shp[1:] != csr_shp[2:]:
         raise ValueError(f"CSRTensor's shape[2: ] must be equal to value's shape[1: ]," \
-                         f"but CSRTensor's shape[2: ] got: {csr_shp[2: ]} and value's shape[1: ]" \
-                         f"got: {values_shp[1: ]}")
+                         f"but CSRTensor's shape[2: ] got: {csr_shp[2:]} and value's shape[1: ]" \
+                         f"got: {values_shp[1:]}")
 
 
 def check_csr_tensor_dtype(indptr_dtype, indices_dtype):
@@ -1144,19 +1160,8 @@ def check_coo_tensor_dtype(indices_dtype):
                         f"{indices_dtype}.")
 
 
-def check_dynamic_shape(dyn_elem, actual_input, i):
-    """Check the consistency of dynamic shape tensors and actual input tensors."""
-    if dyn_elem.dtype != actual_input.dtype:
-        raise TypeError(f"The data type of '{i}'th args in actual input tensors should be '{dyn_elem.dtype}', " \
-                        f"but got '{actual_input.dtype}'.")
-    if dyn_elem.ndim != actual_input.ndim:
-        raise ValueError(f"The dimension of '{i}'th args in actual input tensors should be '{dyn_elem.ndim}', " \
-                         f"but got '{actual_input.ndim}'.")
-    check_dyn_shape_value_equal(i, dyn_elem.shape, actual_input.shape)
-
-
 def check_element_type_of_iterable(arg_name, arg_value, valid_types, prim_name=None):
-    """Check type of the element of a iterabel object, execpt dict."""
+    """Check type of the element of a iterabel object, except dict."""
     check_value_type(arg_name, arg_value, [list, tuple], prim_name)
     type_names = [t.__name__ if hasattr(t, '__name__') else str(t) for t in valid_types]
     num_types = len(valid_types)
@@ -1196,12 +1201,61 @@ def check_size_and_element_type_of_tuple(arg_name, arg_value, expect_size, expec
     check_element_type_of_iterable('arg_name', arg_value, [expect_element_type], prim_name)
 
 
-def check_dyn_shape_value_equal(index, dyn_shape, actual_shape):
-    """Check the consistency of dynamic shape and actual input shape."""
-    for i, x in enumerate(dyn_shape):
-        if x not in (-1, actual_shape[i]):
-            raise ValueError(f"The {i}th shape value of `{index}`th actual input args should be `{x}`, but got " \
-                             f"`{actual_shape[i]}`.")
+def _check_symbol(dyn_input, net_input, index, symbolic_shape_data):
+    """Check symbolic shape values."""
+    actual_shape = net_input.shape
+    for i, sym in enumerate(dyn_input.symbolic_shape):
+        # the Symbol is converted to dict
+        if not isinstance(sym, dict):
+            continue
+        # the value of symbols with same "id" should be equal.
+        if "id" in sym:
+            sym_id = sym["id"]
+            k_idval = "unique_id_value_map"
+            if k_idval not in symbolic_shape_data:
+                symbolic_shape_data[k_idval] = {}
+            unique_id_value = symbolic_shape_data[k_idval]
+            if sym_id not in unique_id_value:
+                unique_id_value[sym_id] = actual_shape[i]
+            elif unique_id_value[sym_id] != actual_shape[i]:
+                raise ValueError(
+                    f"The {i + 1}th shape value of {index + 1}th actual input args is a unique symbol, all values must "
+                    f"be the same. The previous value is {unique_id_value[sym_id]}, but the current value is "
+                    f"{actual_shape[i]}. Actual shape: {actual_shape}, axis: {i}.")
+        # check the value in range [min, max].
+        if "min" in sym and actual_shape[i] < sym["min"]:
+            raise ValueError(
+                f"The {i + 1}th shape value of {index + 1}th actual input args must be greater than or equal to the "
+                f"'min' value '{sym['min']}' of `Symbol`, but got '{actual_shape[i]}'.  Actual shape: {actual_shape}, "
+                f"axis: {i}.")
+        if "max" in sym and actual_shape[i] > sym["max"]:
+            raise ValueError(
+                f"The {i + 1}th shape value of {index + 1}th actual input args must be less than or equal to the "
+                f"'max' value '{sym['max']}' of `Symbol`, but got '{actual_shape[i]}'. Actual shape: {actual_shape}, "
+                f"axis: {i}.")
+        # check the shape item that satisfies the "divisor * N + remainder, N >= 1".
+        d = sym.get("divisor", 1)
+        r = sym.get("remainder", 0)
+        if actual_shape[i] < d or actual_shape[i] % d != r:
+            raise ValueError(
+                f"The {i + 1}th shape value of {index + 1}th actual input args must be match the 'divisor'(d) and "
+                f"'remainder'(r) of `Symbol`. The value should be 'd * N + r' for 'N > 0', got d={d} and r={r}, but "
+                f"actual shape value is '{actual_shape[i]}'. Actual shape: {actual_shape}, axis: {i}")
+
+
+def check_symbolic_shape(dynamic_inputs, actual_inputs):
+    """Check the symboic shape"""
+    symbolic_shape_data = {}
+
+    def run_check(dyn_inputs, net_inputs):
+        """the real checking function"""
+        for index, (dyn_input, net_input) in enumerate(zip(dyn_inputs, net_inputs)):
+            if isinstance(dyn_input, (tuple, list)):
+                run_check(dyn_input, net_input)
+            elif hasattr(dyn_input, "symbolic_shape"):
+                _check_symbol(dyn_input, net_input, index, symbolic_shape_data)
+
+    run_check(dynamic_inputs, actual_inputs)
 
 
 def check_input_format(input_param):
@@ -1317,4 +1371,49 @@ def args_type_check(*type_args, **type_kwargs):
     return type_check
 
 
+def check_hook_fn(hook_type, hook_fn):
+    """Check hook fn"""
+    if context.get_context("mode") != context.PYNATIVE_MODE:
+        logger.warning(f"'{hook_type}' function is only supported in pynative mode, you can use "
+                       f"context.set_context to set pynative mode.")
+        return False
+
+    if not isinstance(hook_fn, (FunctionType, MethodType)):
+        raise TypeError(f"When using 'hook_type(hook_fn)', the type of 'hook_fn' must be python "
+                        f"function, but got {type(hook_fn)}.")
+
+    if hook_fn.__code__.co_name == "staging_specialize":
+        raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@jit' is not supported.")
+
+    tensor_hook_func_args_num = 1
+    pre_hook_func_args_num = 2
+    forward_hook_and_backward_hook_func_args_num = 3
+    # Real args number, exclude class method self param
+    hook_fn_args_num = len(inspect.signature(hook_fn).parameters)
+
+    if hook_type == "register_hook" and hook_fn_args_num != tensor_hook_func_args_num:
+        raise TypeError(f"Tensor hook function {hook_fn.__name__} arg num should be {tensor_hook_func_args_num}, but "
+                        f"got {hook_fn_args_num}")
+
+    if hook_type == "register_forward_pre_hook" and hook_fn_args_num != pre_hook_func_args_num:
+        raise TypeError(f"forward_pre_hook function {hook_fn.__name__} args num should be {pre_hook_func_args_num}, "
+                        f"but got {hook_fn_args_num}")
+
+    if (hook_type == "register_forward_hook" and
+            hook_fn_args_num != forward_hook_and_backward_hook_func_args_num):
+        raise TypeError(f"forward_hook function {hook_fn.__name__} args num should be "
+                        f"{forward_hook_and_backward_hook_func_args_num}, but got {hook_fn_args_num}")
+
+    if hook_type == "register_backward_pre_hook" and hook_fn_args_num != pre_hook_func_args_num:
+        raise TypeError(f"backward_pre_hook function {hook_fn.__name__} args num should be {pre_hook_func_args_num},"
+                        f" but got {hook_fn_args_num}")
+
+    if (hook_type == "register_backward_hook" and
+            hook_fn_args_num != forward_hook_and_backward_hook_func_args_num):
+        raise TypeError(f"backward_hook function {hook_fn.__name__} args num should be "
+                        f"{forward_hook_and_backward_hook_func_args_num}, but got {hook_fn_args_num}")
+
+    return True
+
+
 _set_record = {}

mindspore/_extends/builtin_operations.py

@@ -216,7 +216,8 @@ def mixed_precision_cast(dst_type, x):
     """Implement `mixed_precision_cast`."""
 
     def cast_inner(data):
-        if isinstance(data, Tensor) and data.dtype in (mstype.float32, mstype.float16, mstype.float64):
+        if isinstance(data, Tensor) and data.dtype in (mstype.float32, mstype.float16,
+                                                       mstype.float64, mstype.bfloat16):
             return F.cast(data, dst_type)
         return data
 

mindspore/_extends/graph_kernel/model/graph_parallel.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
+# Copyright 2021-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.
@@ -290,12 +290,22 @@ def block_parallel_estimate(graphs):
     return ParalGain(fusion_type, max_weight, sum_weight - max_weight, blocks, type_info)
 
 
+def block_parallel_estimate_dvm(graphs):
+    """estimate block parallel gain when on Ascend platform"""
+    max_ops = 96
+    total_ops = 0
+    bottleneck = 1000
+    gain = 1000
+    for g in graphs:
+        total_ops += len(g.ops)
+    if total_ops > max_ops:
+        gain = 0
+        return ParalGain("block_fusion", bottleneck, gain, list(0 for _ in graphs), None)
+    return ParalGain("block_fusion", bottleneck, gain, list(1 for _ in graphs), None)
+
+
 def parallel_estimate(graphs, target):
     """Estimate parallel gain"""
     if target == "aicore":
-        fusion_type = "block_fusion"
-        type_info = None
-        fake_estimate = 1000
-        fake_blocks = list(1 for g in graphs)
-        return ParalGain(fusion_type, fake_estimate, fake_estimate, fake_blocks, type_info)
+        return block_parallel_estimate_dvm(graphs)
     return block_parallel_estimate(graphs)

mindspore/_extends/parallel_compile/akg_compiler/akg_process.py

@@ -21,7 +21,7 @@ from multiprocessing import Pool, cpu_count
 from mindspore import log as logger
 from mindspore._extends.parallel_compile.akg_compiler.get_file_path import get_akg_path
 from mindspore._extends.parallel_compile.akg_compiler.util import get_ascend_compile_dirs, create_compile_dirs, \
-    get_log_level, update_attr, select_best, print_compile_log, check_tbe_support, get_kernel_meta_parent_dir
+    get_log_level, update_attr, select_best, print_compile_log, get_kernel_meta_parent_dir
 
 
 def _compile_akg_task_default(json_strs, attrs, func):
@@ -66,7 +66,6 @@ def _compile_akg_task_ascend(json_strs, attrs):
     compile_dirs = get_ascend_compile_dirs(attrs)
     kernel_meta_dir = compile_dirs.get("kernel_meta_dir")
     akg_compile_dir = compile_dirs.get("akg_compile_dir")
-    tbe_compile_dir = compile_dirs.get("tbe_compile_dir")
     composite_graph_dir = compile_dirs.get("composite_graph_dir")
     attrs = update_attr(attrs, {"dump_composite_graph": composite_graph_dir, "optimize_for_tbe": True})
     for json_str in json_strs:
@@ -79,24 +78,12 @@ def _compile_akg_task_ascend(json_strs, attrs):
         info_path = os.path.join(kernel_meta_dir, op_name + ".info")
         if not os.path.isfile(info_path):
             raise FileNotFoundError("Can not compile non-existing file \"{}\"".format(info_path))
-
         # Compile json str with AKG
         _compile_subprocess(compiler, akg_compile_dir, info_path, "AKG", attrs, compile_log, log_level)
 
-        # Load composite optimized json str and compile it with TBE
-        composite_graph_path = os.path.join(composite_graph_dir, op_name + ".info")
-        if not os.path.isfile(composite_graph_path):
-            composite_graph_path = info_path
-        with open(composite_graph_path, 'r') as f:
-            composite_graph = f.read()
-        if "buffer_stitch" not in json_desc and "parallel_fusion" not in json_desc and \
-                check_tbe_support(json.loads(composite_graph)):
-            _compile_subprocess(compiler, tbe_compile_dir, composite_graph_path, "TBE", attrs, compile_log, log_level)
-
         print_compile_log(compile_log)
         # Select best compile result
-        res = select_best([os.path.join(akg_compile_dir, "akg_kernel_meta"), os.path.join(
-            tbe_compile_dir, "kernel_meta")], kernel_meta_dir, op_name)
+        res = select_best([os.path.join(akg_compile_dir, "kernel_meta")], kernel_meta_dir, op_name)
         if not res:
             if log_level == "ERROR":
                 raise ValueError("Compile error, json str: {}! build attrs: {}".format(json_str, attrs))
@@ -116,7 +103,7 @@ def create_akg_parallel_process(process_num, wait_time, platform):
 def _is_input_shape_dynamic(desc_d):
     input_lists = desc_d.get("input_desc", [])
     if input_lists is None:
-        return True
+        return False
     for input_desc in input_lists:
         shape = input_desc[0].get("shape", ())
         if -1 in shape or -2 in shape:

mindspore/_extends/parallel_compile/akg_compiler/build_tbe_kernel.py

@@ -41,11 +41,22 @@ ENABLE_VECTOR_2X = "enable_vector_2x"
 ENABLE_GROUP_INPLACE = "enable_group_inplace"
 
 
-def initialize(kernel_meta_parent_dir):
+def get_arch_name(op_json):
+    """Get the arch name from op json."""
+    target_info = op_json.get("target_info")
+    if isinstance(target_info, dict):
+        arch_name = target_info.get("arch")
+        if arch_name:
+            return arch_name
+    return "Ascend910A"
+
+
+def initialize(kernel_meta_parent_dir, arch_name):
     """Initialize the TBE compile environment."""
     os.environ["CONTEXT_MODELCOMPILING"] = "TRUE"
+    core_type = "VectorCore" if arch_name != "Ascend910A" else ""
     # socVersion, coreType, coreNum, l1Fusion, l2Mode, l2Fusion
-    soc_info = ["Ascend910A", "", "", "false", "2", "false",
+    soc_info = [arch_name, core_type, "", "false", "2", "false",
                 {"op_impl_mode": "",
                  "op_debug_level": "3",
                  "op_impl_mode_list": [],
@@ -83,7 +94,7 @@ def add_new_shape(names, shapes, new_shapes, inputs):
         if shapes[i] == new_shapes[i]:
             continue
         if name not in inputs:
-            raise RuntimeError("Can not support reshape on output tensor {}".format(name))
+            continue
         if NEW_SHAPE not in inputs[name]:
             inputs[name][NEW_SHAPE] = new_shapes[i]
         elif new_shapes[i] != inputs[name][NEW_SHAPE]:
@@ -512,6 +523,7 @@ def build(json_str, kernel_meta_parent_dir):
 
 def build_tbe_kernel(json_str, kernel_meta_parent_dir):
     """Build TBE kernel."""
-    initialize(kernel_meta_parent_dir)
+    arch_name = get_arch_name(json.loads(json_str))
+    initialize(kernel_meta_parent_dir, arch_name)
     with build_config(kernel_meta_parent_dir=kernel_meta_parent_dir):
         build(json_str, kernel_meta_parent_dir)

mindspore/_extends/parallel_compile/akg_compiler/compiler.py

@@ -39,6 +39,7 @@ def run_compiler(op_json, compile_backend="AKG", attrs=None, kernel_meta_parent_
         build_tbe_kernel(op_json, kernel_meta_parent_dir)
     else:
         os.environ["MS_COMPILER_CACHE_PATH"] = kernel_meta_parent_dir
+        os.environ["KERNEL_META_DIR"] = "kernel_meta"
         from get_file_path import get_akg_path
         sys.path.insert(0, get_akg_path())
         p = __import__("akg", globals(), locals(), ['ms'], 0)

mindspore/_extends/parallel_compile/akg_compiler/gen_custom_op_files.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+# coding: utf-8
+# Copyright 2023 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.
+"""Used to generate custom akg op files, which will be invoked by GE."""
+
+import os
+import sys
+import json
+import shutil
+from itertools import product
+
+SUPPORTED_INPUT_NUM = [1, 2, 3, 4, 5, 6, 7]
+SUPPORTED_OUTPUT_NUM = [1, 2, 3, 4, 5]
+SUPPORTED_DEVICE_ARCH = ["ascend910", "ascend910b"]
+VALUE_ALL = "all"
+VALUE = "value"
+NAME = "name"
+SHAPE = "shape"
+REQUIRED = "required"
+PARAM_TYPE = "paramType"
+
+
+def gen_ops_info():
+    """Generate the custom akg op registration information."""
+    ops_info = {}
+    # supported_io_num: [(1, 1), (1, 2), ...], list of (input_num, output_num)
+    supported_io_num = list(product(SUPPORTED_INPUT_NUM, SUPPORTED_OUTPUT_NUM))
+    for input_num, output_num in supported_io_num:
+        op_info = {"attr": {"list": "info_path"},
+                   "attr_info_path": {PARAM_TYPE: REQUIRED, "type": "str", VALUE: VALUE_ALL},
+                   "opFile": {VALUE: "custom"},
+                   "opInterface": {VALUE: "custom_train"},
+                   "dynamicFormat": {"flag": "true"}}
+        for i in range(input_num):
+            op_info["input" + str(i)] = {NAME: "x" + str(i),
+                                         PARAM_TYPE: REQUIRED,
+                                         SHAPE: VALUE_ALL}
+        for i in range(output_num):
+            op_info["output" + str(i)] = {NAME: "y" + str(i),
+                                          PARAM_TYPE: REQUIRED,
+                                          SHAPE: VALUE_ALL}
+        op_type = "Fused_x{}_y{}".format(input_num, output_num)
+        ops_info[op_type] = op_info
+    return ops_info
+
+
+def gen_custom_op_files(config_dir, dsl_dir):
+    """Copy custom akg op registration information file to config_dir, and copy python dsl file to dsl_dir."""
+    if not os.path.isdir(config_dir):
+        os.makedirs(config_dir, exist_ok=True)
+    config_dir = os.path.realpath(config_dir)
+    if not os.path.isdir(dsl_dir):
+        os.makedirs(dsl_dir, exist_ok=True)
+    dsl_dir = os.path.realpath(dsl_dir)
+    cur_path = os.path.split(os.path.realpath(__file__))[0]
+
+    # generate custom akg op registration information
+    ops_info = gen_ops_info()
+    for device_arch in SUPPORTED_DEVICE_ARCH:
+        sub_dir = os.path.join(config_dir, device_arch)
+        if not os.path.isdir(sub_dir):
+            os.makedirs(sub_dir)
+        ops_info_path = os.path.join(sub_dir, "aic-{}-ops-info.json".format(device_arch))
+        if os.path.isfile(ops_info_path):
+            with open(ops_info_path, 'r') as f:
+                info_des = json.loads(f.read())
+                ops_info.update(info_des)
+        with os.fdopen(os.open(ops_info_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o660), 'w') as f:
+            f.write(json.dumps(ops_info, indent=4))
+
+    # custom akg op dsl file
+    custom_py = os.path.join(cur_path, "../../../../../lite/tools/kernel_builder/ascend/akg/custom.py")
+    if not os.path.isfile(custom_py):
+        raise RuntimeError("custom.py path is invalid: {}".format(custom_py))
+    shutil.copy(custom_py, dsl_dir)
+    shutil.copy(os.path.join(cur_path, "compiler.py"), dsl_dir)
+    shutil.copy(os.path.join(cur_path, "build_tbe_kernel.py"), dsl_dir)
+    shutil.copy(os.path.join(cur_path, "tbe_topi.py"), dsl_dir)
+    shutil.copy(os.path.join(cur_path, "get_file_path.py"), dsl_dir)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) == 3:
+        gen_custom_op_files(sys.argv[1], sys.argv[2])

mindspore/_extends/parallel_compile/akg_compiler/tbe_topi.py

@@ -114,7 +114,8 @@ def _add(x0, x1, attrs=None):
         return tbe.dsl.vadds(x1, x0)
     x0, x1 = _broadcast(x0, x1)
     from impl.add import add_compute
-    return add_compute(x0, x1, None, kernel_name=attrs["fusion_op_name"])
+    output_desc = {"dtype": x1.dtype, "shape": x1.shape}
+    return add_compute(x0, x1, output_desc, kernel_name=attrs["fusion_op_name"])
 
 
 @reg_op("Asin", pattern=OpPattern.ELEMWISE)

mindspore/_extends/parallel_compile/akg_compiler/util.py

@@ -20,6 +20,7 @@ from mindspore import log as logger
 from mindspore._extends.parallel_compile.akg_compiler.tbe_topi import get_op_reg_info
 
 O_SUFFIX = ".o"
+SO_SUFFIX = ".so"
 JSON_SUFFIX = ".json"
 
 
@@ -27,7 +28,7 @@ def update_attr(attr, new_attr):
     """Update new_attr to attr."""
     if attr is None:
         attr = {}
-    elif attr is str:
+    if isinstance(attr, str):
         attr = json.loads(attr)
     if isinstance(attr, dict):
         attr.update(new_attr)
@@ -114,7 +115,7 @@ def select_best(src_dirs, dst_dir, op_name):
     max_block_dim = 1
     max_block_dim_idx = -1
     for i, src_dir in enumerate(src_dirs):
-        o_path = os.path.join(src_dir, op_name + O_SUFFIX)
+        o_path = os.path.join(src_dir, op_name + SO_SUFFIX)
         json_path = os.path.join(src_dir, op_name + JSON_SUFFIX)
         if os.path.isfile(o_path) and os.path.isfile(json_path):
             with open(json_path, 'r') as f:
@@ -125,8 +126,10 @@ def select_best(src_dirs, dst_dir, op_name):
                     max_block_dim = json_dict["blockDim"]
     if max_block_dim_idx >= 0:
         o_path = os.path.join(src_dirs[max_block_dim_idx], op_name + O_SUFFIX)
+        so_path = os.path.join(src_dirs[max_block_dim_idx], op_name + SO_SUFFIX)
         json_path = os.path.join(src_dirs[max_block_dim_idx], op_name + JSON_SUFFIX)
         _copy_file(o_path, os.path.join(dst_dir, op_name + O_SUFFIX))
+        _copy_file(so_path, os.path.join(dst_dir, op_name + SO_SUFFIX))
         _copy_file(json_path, os.path.join(dst_dir, op_name + JSON_SUFFIX))
         logger.info("{}, best compile result dir: {}".format(op_name, src_dirs[max_block_dim_idx]))
         return True