mindspore 2.2.14__cp38-cp38-manylinux1_x86_64.whl → 2.3.0rc2__cp38-cp38-manylinux1_x86_64.whl

mindspore/.commit_id

@@ -1 +1 @@
-__commit_id__ = '[sha1]:4f30cd5f,[branch]:(HEAD,origin/r2.2,r2.2)'
+__commit_id__ = '[sha1]:c6a1400a,[branch]:(HEAD,origin/master,origin/HEAD,master)'

mindspore/__init__.py

@@ -17,7 +17,7 @@ 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
@@ -30,13 +30,13 @@ 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, \
-    load_segmented_checkpoints
-from mindspore.rewrite import SymbolTree, ScopedValue, Node, NodeType, TreeNodeHelper
+    Layout, sync_pipeline_shared_parameters, parameter_broadcast, load_segmented_checkpoints
+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.ops_generate import arg_dtype_cast, arg_handler
 
 __all__ = ["run_check"]
 __all__.extend(__version__)

mindspore/_akg/akg/composite/build_module.py

@@ -30,6 +30,31 @@ from .construct_args import ConstructType, ConstructKey
 from .construct_args import get_construct_args, get_tune_construct_args, \
     should_enable_attr, get_stmt_for_tune, add_attrs_in_segment_infos
 from utils.util import get_ascend_type
+
+matmul_keys = [
+    "m0", # [max(16,m)%m0==0, m0%16==0?]
+    "k0", # [max(16,m)%m0==0,]
+    "n0", # [max(16,m)%m0==0,]
+    "swizzlCount",#[1-blockDim],
+    "swizzlDirect"#[0,1] nz or zn format
+]
+
+matmul_calc_keys = [
+    "mLoop",#=max(16,m)/m0
+    "kLoop",#=max(16,m)/m0
+    "nLoop",#=max(16,m)/m0
+    "coreLoop",#=mLoop/nLoop/kLoop
+]
+
+pa_keys = [
+    "headSplit",
+]
+
+kv_keys = [
+    "num_token_tile",
+    "n_burst"
+]
+
 def generate_trait(desc):
     """
     generate trait of kernel description
@@ -44,7 +69,7 @@ def generate_trait(desc):
         input_idx.sort()
         input_idx_str = ''.join(str(i) for i in input_idx)
         op_trait = op['name'] + input_idx_str
-        if op['name'] == "MatMul":
+        if  op['name'].find("MatMul") != -1:
             for attr in op['attr']:
                 if attr['name'] == "transpose_a":
                     transpose_a = str(int(attr['value']))
@@ -189,6 +214,8 @@ def _update_target_info(desc_d, attr):
         # auto detect proper gpu device type according to compute capability description
         if target_info.get("compute_capability") == "8.0":
             attr["device_type"] = "a100"
+        elif target_info.get("compute_capability") == "7.0":
+            attr["device_type"] = "v100"
     elif process == "cpu":
         if target_info.get("feature"):
             attr["feature"] = target_info.get("feature")
@@ -653,9 +680,16 @@ def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
             attr["pragma_modshift"] = True
     segment_tree, segment_infos = get_construct_args(desc_s_in, attr, post_funcs)
 
+    if desc_d_in.get("backend"):
+        backend_func = akg.tvm.get_global_func("cce.set_backend")
+        backend_func(desc_d_in["backend"])
+
     if desc_d_in.get("enable_cce_lib"):
         attr["enable_cce_lib"] = True
-        return _build_to_module_ascend_lib(desc_s_in, kernel_name)
+        repository = None
+        if os.getenv('MS_GRAPH_KERNEL_TILING'):
+            repository = read_repo_file(str(os.getenv('MS_GRAPH_KERNEL_TILING')))
+        return _build_to_module_ascend_lib(desc_s_in, kernel_name, repository)
 
     poly = True
     res = _cpp_build(attr, process, poly, segment_tree, segment_infos)
@@ -672,10 +706,83 @@ def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
         is_success = build_tbe_codegen(kernel_name, stmt_json, args_json, attr, ascend_type)
         if not is_success:
             raise TypeError("npu_inference codegen failed.")
+        akg.tvm.get_global_func("build_host_cce")(res[1], res[2], kernel_name)
         return kernel_name
     return res
 
-def _build_to_module_ascend_lib(desc_s_in, kernel_name):
+def _build_to_module_ascend_lib(desc_s_in, kernel_name, repository=None):
+    
+    def _convert_dim_to_attr(dim, tiling_info, compute):
+        f = akg.tvm.get_global_func("cce.current_product_conf_core")
+        coreNum = f("Core_num")
+        if coreNum == 0:
+            raise RuntimeError("Get the cce product value is None")
+        if "MatMul" in compute:
+            batch = tiling_info.get("batch_size")
+            m = tiling_info.get("M")
+            n = tiling_info.get("N")
+            k = tiling_info.get("K")
+            matmul_dict = dict()
+            dim_list = dim.split(" ")
+            if len(dim_list) < 4:
+                return {}
+            cnt =  0
+            for i in range(2, len(dim_list), 4):
+                key = matmul_keys[cnt]
+                value = dim_list[i]
+                matmul_dict[key] = int(value)
+                cnt += 1
+            mLoop = (max(16, m) - 1) // matmul_dict.get("m0") + 1
+            nLoop = (max(16, n) - 1) // matmul_dict.get("n0") + 1
+            kLoop = (max(16, k) - 1) // matmul_dict.get("k0") + 1
+            coreLoop = batch * mLoop * nLoop
+            blockDim = min(coreLoop, coreNum)
+            matmul_dict["mLoop"] = mLoop
+            matmul_dict["kLoop"] = kLoop
+            matmul_dict["nLoop"] = nLoop
+            matmul_dict["coreLoop"] = coreLoop
+            matmul_dict["blockDim"] = blockDim
+            return matmul_dict
+        elif "PagedAttention" in compute or "PagedAttentionMask" in compute:
+            num_heads = tiling_info.get("num_heads")
+            num_tokens = tiling_info.get("num_tokens")
+            pa_dict = dict()
+            dim_list = dim.split(" ")
+            if len(dim_list) < 4:
+                return {}
+            cnt =  0
+            for i in range(2, len(dim_list), 4):
+                key = pa_keys[cnt]
+                value = dim_list[i]
+                pa_dict[key] = int(value)
+                cnt += 1
+            head_split = pa_dict.get("headSplit")  # Tuning<<---- [1, num_heads]
+            loopLen = (num_heads + head_split - 1) // head_split
+            block = min(loopLen * num_tokens, coreNum)
+            pa_dict["blockDim"] = block
+            return pa_dict
+        elif "ReshapeAndCache" in compute:
+            num_tokens = tiling_info.get("num_tokens")
+
+            kv_dict = dict()
+            dim_list = dim.split(" ")
+            cnt =  0
+            for i in range(2, len(dim_list), 4):
+                key = kv_keys[cnt]
+                value = dim_list[i]
+                kv_dict[key] = int(value)
+                cnt += 1
+            if "num_token_tile" in kv_dict:
+                block_dim = str((int(num_tokens) - 1) // kv_dict.get("num_token_tile") + 1)
+            else:
+                block_dim = num_tokens
+            
+            kv_dict["block_dim"] = block_dim
+            kv_dict["kernel_version"] = 2
+            return kv_dict
+        else:
+            return {}
+    
     def _get_all_shape(shapes):
         shape_split = shapes.split(".")
         shape_list = []
@@ -701,10 +808,7 @@ def _build_to_module_ascend_lib(desc_s_in, kernel_name):
             tensor_B = shape_list[1]
 
             tensor_A_split = tensor_A.split("_")
-            if len(tensor_A_split) > 2:
-                batch_size = int(tensor_A.split("_")[0])
-            else:
-                batch_size = 1
+
             if trans_a == "1":
                 M = int(tensor_A_split[-1])
                 K = int(tensor_A_split[-2])
@@ -712,13 +816,30 @@ def _build_to_module_ascend_lib(desc_s_in, kernel_name):
                 M = int(tensor_A_split[-2])
                 K = int(tensor_A_split[-1])
 
+            tensor_B_split = tensor_B.split("_")
             if trans_b == "1":
-                N = int(tensor_B.split("_")[-2])
+                N = int(tensor_B_split[-2])
             else:
-                N = int(tensor_B.split("_")[-1])
-            tensor_A_type = str(dtype.split("-")[0])
+                N = int(tensor_B_split[-1])
+            batch_size = 1
+            if len(tensor_A_split) > 2 and len(tensor_B_split) == 2:
+                # some bmm in onnx will have multi batch such as [b1,b0,m,k]*[k,n]
+                batch_sizes = [int(i) for i in tensor_A.split("_")[:-2]]
+                # change bmm with shape [b1,b0,m,k]*[k,n] to matmul with shape [b1*b0*m,k]*[k,n]
+                for b in batch_sizes:
+                    M = M * b
+                batch_size = 1
+            elif len(tensor_A_split) == len(tensor_B_split) and len(tensor_B_split) > 2:
+                # common bmm [b1,b0,m,k]*[b1,b0,k,n]
+                batch_sizes = [int(i) for i in tensor_A.split("_")[:-2]]
+                for b in batch_sizes:
+                    batch_size = batch_size * b
+            op_type = "MatMul"
+            if bias_flag:
+                op_type = "MatMulMix"
+            tensor_a_type = str(dtype.split("-")[0])
             tiling_info = {"batch_size":batch_size, "M": M, "N": N, "K": K, "trans_a": int(trans_a), "trans_b": int(trans_b),
-                           "tensor_A_type": tensor_A_type, "bias_flag": bias_flag, "op_type": "MatMul"}
+                           "tensor_A_type": tensor_a_type, "bias_flag": bias_flag, "op_type": op_type}
         elif "PagedAttention" in compute or "PagedAttentionMask" in compute:
             shape_list = _get_all_shape(shape)
             query = shape_list[0]
@@ -752,6 +873,24 @@ def _build_to_module_ascend_lib(desc_s_in, kernel_name):
 
             tiling_info = {"num_tokens": num_tokens, "num_heads": num_heads, "head_size": head_size, 
                            "op_type": "ReshapeAndCache"}
+        elif "Transpose" in compute:
+            shape_list = _get_all_shape(shape)
+            input_shape = shape_list[0]
+
+            axis_1 = int(input_shape.split("_")[0])
+            axis_2 = int(input_shape.split("_")[1])
+            axis_3 = int(input_shape.split("_")[2])
+
+            tiling_info = {"op_type": "AddReshapeTranspose", "axis_1":axis_1, "axis_2": axis_2, "axis_3":axis_3}
+        
+        tiling_info["use_repo"] = 0
+        if  repository is not None:
+            repo_attr = get_attr_from_dict([compute, shape, dtype, 'metadata', 'attrs', 'dim'], repository, {})
+            if len(repo_attr) > 0:
+                repo_attr = _convert_dim_to_attr(repo_attr, tiling_info, compute)
+                tiling_info = merge_attrs(tiling_info, repo_attr)
+                tiling_info["use_repo"] = 1
+        tiling_info["arch"] = json.loads(desc_s_in)["target_info"]["arch"]
         return tiling_info
 
     func = tvm.get_global_func("build_cce_lib")
@@ -833,6 +972,11 @@ def get_tiling_space(kernel_desc, level=1, attr=None):
     attr['help_tiling'] = level
     attr['tuning'] = 'on'
     desc_d = json.loads(kernel_desc)
+    from akg.ms.info_version_adapt import InfoVersionAdapt
+    info_adapter = InfoVersionAdapt(desc_d)
+    if not info_adapter.run():
+        raise RuntimeError(info_adapter.msg)
+    kernel_desc = _set_backend(desc_d)
     backend = desc_d['process']
     all_ops = set(op['name'] for op in desc_d['op_desc'])
     if backend == "cuda":

mindspore/_akg/akg/config/repository.json

@@ -2635,5 +2635,43 @@
             }
         }
         }
+    },
+    "1.StridedSliceV21.StridedSliceV22.FastGeLU1.Mul13.4": {
+        "2_4096_2560.2_4096_1280": {
+            "float16-": {
+                "metadata": {
+                    "attrs": {
+                        "dim": "0 0 8 8 0 1 1280 1280"
+                    }
+                }
+            }
+        },
+        "2_1024_5120.2_1024_2560": {
+            "float16-": {
+                "metadata": {
+                    "attrs": {
+                        "dim": "0 0 4 4 0 1 2560 2560"
+                    }
+                }
+            }
+        },
+        "2_256_10240.2_256_5120": {
+            "float16-": {
+                "metadata": {
+                    "attrs": {
+                        "dim": "0 0 2 2 0 1 5120 5120"
+                    }
+                }
+            }
+        },
+        "2_64_10240.2_64_5120": {
+            "float16-": {
+                "metadata": {
+                    "attrs": {
+                        "dim": "0 0 1 1 0 1 5120 5120"
+                    }
+                }
+            }
+        }
     }
 }

mindspore/_akg/akg/ms/info_version_adapt.py

@@ -38,6 +38,30 @@ this op named "value" is guaranteed, which contains the value of this input.
 import logging
 
 
+def type_id_to_type_string(type_id):
+    type_dict = {
+        30: "bool",
+        32: "int8",
+        33: "int16",
+        34: "int32",
+        35: "int64",
+        37: "uint8",
+        38: "uint16",
+        39: "uint32",
+        40: "uint64",
+        42: "float16",
+        43: "float32",
+        44: "float64",
+        45: "bfloat16",
+        48: "complex64",
+        49: "complex128",
+    }
+    if type_id not in type_dict:
+        logging.exception("Type_id is invalid, please check. Type_id: %s", str(type_id))
+    return type_dict[type_id]
+
+
+
 def convert_input_to_attr(kernel_info:dict):
     op_info = {
         "Reshape": [(1, "shape")],
@@ -47,7 +71,10 @@ def convert_input_to_attr(kernel_info:dict):
         "Transpose": [(1, "perm")],
         "ExpandDims": [(1, "axis")],
         "Tile": [(1, "multiples")],
+        "BroadcastTo": [(1, "shape")],
+        "Cast": [(1, "dst_type", type_id_to_type_string)],
         "StridedSlice": [(3, "strides"), (2, "end"), (1, "begin")],
+        "StridedSliceV2": [(4, "strides"),(3, "axes"), (2, "end"), (1, "begin")],
         "OneHot": [(1, "depth")],
         "Gather": [(2, "axis")],
         "UnsortedSegmentSum": [(2, "num_segments")],
@@ -74,6 +101,8 @@ def convert_input_to_attr(kernel_info:dict):
                 input_desc_i = op["input_desc"].pop(input_index)
                 input_value = input_desc_i[0]["value"]
                 input_dtype = "listInt"
+                if len(input_info) > 2:
+                    input_value = input_info[2](input_value)
                 if op_name not in int_input_required_ops and isinstance(input_value, int):
                     input_value = [input_value]
                 if isinstance(input_value, int):

mindspore/_akg/akg/tvm/contrib/nvcc.py

@@ -80,9 +80,12 @@ def compile_cuda(code,
     """
     # supported gcc version for cuda code generation
     gcc_version = get_gcc_version()
-    if gcc_version is None or (7,3,0) > gcc_version or gcc_version > (11,1,0):
+    if gcc_version is None or (7, 3, 0) > gcc_version:
         raise ValueError("gcc version is {}, not in range [7.3.0, 11.1.0]".format(
             gcc_version))
+    if gcc_version > (11, 1, 0):
+        print("[WARNING] gcc version is {}, which beyonds our maintenance range [7.3.0, 11.1.0]".format(
+            gcc_version))
 
     arch_exception = ["sm_00", None]
     temp = util.tempdir()

mindspore/_akg/akg/utils/ascend_profilier/path_manager.py

@@ -6,9 +6,10 @@ class PathManager:
     @classmethod
     def get_cann_path(cls, profiler_path: str) -> str:
         sub_dirs = os.listdir(os.path.realpath(profiler_path))
+        sub_dirs.sort(reverse=True)
         for sub_dir in sub_dirs:
             sub_path = os.path.join(profiler_path, sub_dir)
-            if os.path.isdir(sub_path) and re.match(r"^PROF_\d+_\d+_[a-zA-Z]+", sub_dir):
+            if os.path.isdir(sub_path) and (re.match(r"^PROF_\d+_\d+_\d+[a-zA-Z]+", sub_dir) or re.match(r"^PROF_\d+_\d+[a-zA-Z]+", sub_dir)):
                 return sub_path
         return ""
 

mindspore/_akg/akg/utils/composite_op_helper.py

@@ -573,8 +573,10 @@ def gen_json_data(op_desc, with_compute=True, input_for_mod=None):
     commands = []
     # Parse input_desc
     _gen_input_data(desc, infos, input_for_mod, commands)
-    # Parse op_desc
-    _gen_op_compute(desc, commands)
+    if with_compute:
+        # Parse op_desc
+        _gen_op_compute(desc, commands)
+
     # Parse output_desc
     _gen_output_data(desc, infos, input_for_mod, output_indexes, commands)
 

mindspore/_akg/akg/utils/dump_ascend_meta.py

@@ -32,10 +32,10 @@ def set_ascend910b(code, core_type, title_dict):
         title_dict["coreType"] = "MIX"
         title_dict["intercoreSync"] = 1
         title_dict["taskRation"] = "1:2"
-    elif core_type == "AIC":
+    elif core_type == "AiCore":
         title_dict["coreType"] = "AiCore"
         title_dict["magic"] = "RT_DEV_BINARY_MAGIC_ELF_AICUBE"
-    elif core_type == "AIV":
+    elif core_type == "VectorCore":
         title_dict["coreType"] = "VectorCore"
         title_dict["magic"] = "RT_DEV_BINARY_MAGIC_ELF_AIVEC"
 

mindspore/_akg/akg/utils/gen_random.py

@@ -61,16 +61,10 @@ def func(size_, miu_=0, sigma_=8, seed_=None):
 
 @func_time_required
 def random_gaussian(size, miu=0, sigma=8, epsilon=0, seed=None):
-    """Generate random array with absolution value obeys gaussian distribution."""
-    random_data_disk_path = None
-    if os.environ.get("RANDOM_DATA_DISK_PATH") is not None:
-        random_data_disk_path = os.environ.get("RANDOM_DATA_DISK_PATH") + "/random_data_%s_%s.bin" % (
-            str(miu), str(sigma))
-
-    if random_data_disk_path is None or (not os.path.exists(random_data_disk_path)):
+    def random_gaussian_without_bin():
         if sigma <= 0:
             sys.stderr.write("Error: Expect positive sigmal for gaussian distribution. but get %f\n" % sigma)
-            sys.exit(1)
+            return []
         size_c = 1
         for x in size:
             size_c = size_c * x
@@ -97,8 +91,20 @@ def random_gaussian(size, miu=0, sigma=8, epsilon=0, seed=None):
         ret = ret.flatten()
         return ret[:size_c].reshape(size) + epsilon
 
+    """Generate random array with absolution value obeys gaussian distribution."""
+    random_data_disk_path = None
+    if os.environ.get("RANDOM_DATA_DISK_PATH") is not None:
+        random_data_disk_path = os.environ.get("RANDOM_DATA_DISK_PATH") + "/random_data_%s_%s.bin" % (
+            str(miu), str(sigma))
+
+    if random_data_disk_path is None or (not os.path.exists(random_data_disk_path)):
+        return random_gaussian_without_bin()
+
     data_len = functools.reduce(lambda x, y: x * y, size)
     data_pool = np.fromfile(random_data_disk_path)
+    if len(data_pool) == 0:
+        return random_gaussian_without_bin()
+
     if data_len % len(data_pool) != 0:
         copy_num = (data_len // len(data_pool)) + 1
     else:

mindspore/_akg/akg/utils/op_dsl.py

@@ -18,6 +18,7 @@ import logging
 import inspect
 from copy import deepcopy
 import numpy as np
+import scipy as sp
 
 
 def get_attr(attr_desc, attr_type):
@@ -34,6 +35,11 @@ def get_input(desc):
     value = desc.get('value', None)
     return value if value is not None else desc['tensor_name']
 
+def strided_slice_v2_np(data, begin, end, strides, axes):
+    """numpy implementation of strided_slice_v2"""
+    for i, axis_i in enumerate(axes):
+        data = np.take(data, np.arange(begin[i], end[i], strides[i]), axis=axis_i)
+    return data
 
 def gather_nd_np(data, indices):
     """numpy implementation of gather_nd"""
@@ -383,6 +389,8 @@ def convert_fracal_shape(ori_shape, fractal):
         return ori_shape[:-4] + (ori_shape[-2] * ori_shape[-3], ori_shape[-1] * ori_shape[-4])
     if fractal == "zZ":
         return ori_shape[:-4] + (ori_shape[-4] * ori_shape[-2], ori_shape[-3] * ori_shape[-1])
+    if fractal =="ND":
+        return ori_shape
 
 
 def strided_slice_str(inputs, output, attr):
@@ -1085,6 +1093,9 @@ op_dsl = {
                                                get_input(inputs[1][0])),
     "Erf": lambda inputs, output, attr: "%s = sp.special.erf(%s)" %
                                         (output[0]['tensor_name'], get_input(inputs[0][0])),
+    "StridedSliceV2": lambda inputs, output, attr: "%s = strided_slice_v2_np(%s, %s, %s, %s, %s)" %
+                                        (output[0]['tensor_name'], get_input(inputs[0][0]), get_attr(attr, "begin"),
+                                        get_attr(attr, "end"), get_attr(attr, "strides"), get_attr(attr, "axes")),
     "TensorScatterAdd": lambda inputs, output, attr: "%s = tensor_scatter_add_np(%s, %s, %s)" %
                                                      (output[0]['tensor_name'], get_input(inputs[0][0]),
                                                       get_input(inputs[1][0]),

mindspore/_akg/akg/utils/tbe_codegen_utils.py

@@ -17,6 +17,7 @@ import os
 import logging
 from akg.utils.util import write_code
 import json
+from akg.global_configs import get_kernel_meta_path
 
 logging.getLogger().setLevel(logging.INFO)
 
@@ -34,15 +35,14 @@ def set_workspace_for_json(json_path, workspace_dict):
 
 
 def copy_to_akg_kernel_meta(kernel_name, postfixs, workspace_dict=None):
-    akg_kernel_mate_str = os.getenv("KERNEL_META_DIR",default="akg_kernel_meta")
     source = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', './'))
     import shutil
-    target = source + "/" + akg_kernel_mate_str + "/" + kernel_name
+    target_dir = get_kernel_meta_path()
+    target = target_dir + kernel_name
     source = source + "/" + "kernel_meta/" + kernel_name
     if source == target:
         return True
-    akg_kernel_mate_str = os.path.abspath(akg_kernel_mate_str)
-    create_directory(akg_kernel_mate_str)
+    create_directory(target_dir)
     for postfix in postfixs:
         if os.path.exists(source + postfix):
             try:
@@ -94,11 +94,21 @@ def build_npu_for_akg(kernel_name,
                       simple_mode=False):
     import tbe
     from tbe.tvm.tir import transform
-    from tbe.tvm.driver.cce_build_module import _count_time, generate_cce_code
+    from tbe.tvm.driver.cce_build_module import _count_time
     from tbe.common.buildcfg import set_current_build_config
     from tbe.common.buildcfg.buildcfg_mapping import dynamic_shape, disable_vectorize, tik, enable_const_fold, \
         dynamic_tik, instrument_bound_checkers, tbe_workspace_size_list_length
-    
+    # try to find generate_cce_code function, since it may have different name in different tbe.
+    import importlib
+    cce_build_module = importlib.import_module("tbe.tvm.driver.cce_build_module")
+    generate_cce_code_function = None
+    generate_cce_code_function_name = {"_generate_cce_code","generate_cce_code"}
+    for func_name in generate_cce_code_function_name:
+        if hasattr(cce_build_module, func_name):
+            generate_cce_code_function = getattr(cce_build_module, func_name)
+    if generate_cce_code_function is None:
+       raise ValueError("Can not find generate cce code function.")
+
     set_current_build_config(tbe_workspace_size_list_length,
                              tbe.tvm.runtime.cce_runtime.tbe_workspace_size_list_length())
 
@@ -167,7 +177,7 @@ def build_npu_for_akg(kernel_name,
 
     _count_time(mod)
     mod = transform.SplitCoreCode()(mod)
-    generate_cce_code(mod, "cce", None)
+    generate_cce_code_function(mod, "cce", None)
 
 
 def build_tbe_codegen(kernel_name, stmt_json, arg_json, attr, ascend_type=None):
@@ -194,7 +204,7 @@ def build_tbe_codegen(kernel_name, stmt_json, arg_json, attr, ascend_type=None):
     cfg[auto_unroll_max_extent] = 0
     cfg[unroll_explicit] = True
     cfg[dynamic_shape] = False
-    cfg[enable_multicore_sync_with_atomic] = True
+    cfg[enable_multicore_sync_with_atomic] = False
     cfg[kernel_meta_parent_dir] = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', './'))
     if ascend_type is None:
         ascend_type = "Ascend910"