mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/_akg/akg/utils/kernel_exec.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # coding: utf-8
-# Copyright 2019-2022 Huawei Technologies Co., Ltd
+# Copyright 2019-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.
@@ -35,8 +35,6 @@ import numpy as np
 
 import akg
 import akg.tvm
-from akg.tvm import autotvm
-from akg.tvm import rpc
 from akg.tvm import _api_internal
 from akg.build_module import help_tiling_level
 from akg.utils import result_analysis as ra_util
@@ -53,8 +51,6 @@ sh = logging.StreamHandler(sys.stdout)
 logging.getLogger().addHandler(sh)
 logging.getLogger().setLevel(logging.INFO)
 
-rpc_machine = {}
-rpc_lb = {}
 
 PERFORMANCE_TEST_FILE = "PERFORMANCE_TEST_FILE"
 BINDS = "binds"
@@ -173,211 +169,6 @@ def gen_name_kernel(kernel, dtype, shapes):
     return res
 
 
-def load_rpc_server_info(mode):
-    """
-    load rpc server host and port info.
-
-    Args:
-        mode (str): string of runtime choose, can set ca aic and rpc.
-    """
-    env_dic = os.environ
-    if env_dic.get('RPC_HOST') and env_dic.get('RPC_PORT'):
-        return
-
-    if mode == 'rpc_cloud':
-        logging.error("runtime_mode=rpc_cloud must set 1980 host ip and port!")
-        raise Exception("ERROR:runtime_mode=rpc_cloud must set 1980 host ip and port!")
-
-    rpc_server_info_config = env_dic.get('RPC_SERVER_INFO_FILE')
-    if not rpc_server_info_config:
-        logging.error("runtime_mode=rpc must set RPC_SERVER_INFO_FILE for rpc server info config")
-        raise Exception("ERROR:runtime_mode=rpc must set RPC_SERVER_INFO_FILE for rpc server info config")
-
-    # load rpc server host and port info from local file.
-    import json
-    with open(rpc_server_info_config, 'r') as f:
-        info = json.load(f)
-
-    for i in info:
-        rpc_machine[i] = info[i]
-        rpc_lb[i] = 0.0
-    return
-
-
-def dispatch(rank=0):
-    """Function for lock waiting dispatch handle version 1."""
-
-    def _sort_by_value(d):
-        items = list(d.items())
-        random.shuffle(items)
-        items.sort(key=lambda x: x[1])
-        return list(item[0] for item in items)
-
-    for k, v in rpc_lb.items():
-        logging.info("######rpc_lb[%s]=%f", rpc_machine.get(k)[0], v)
-    lb_list = _sort_by_value(rpc_lb)
-    if len(lb_list) > rank:
-        return lb_list[rank]
-    return lb_list[len(lb_list) - 1]
-
-
-def commit(remote, weight):
-    rpc_lb[remote] = weight
-
-
-@func_time_required
-def mod_launch_rpc_worker(mod, args, outputs, host, port, tuning=False):
-    """internal RPC worker, should be called by mod_launch_rpc_thread."""
-    logging.info("%s:====start connect to rpc ip: %s, rpc port: %d ",
-                 datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), host, port)
-    remote = rpc.connect(host, port, session_timeout=300)
-    logging.info("%s:====connect to rpc ip: %s, rpc port: %d finished ",
-                 datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), host, port)
-    uuid_str = uuid.uuid4().hex
-    temp_file_name = "stackvm_%s.o" % uuid_str
-    mod.save(temp_file_name)
-    remote.upload(temp_file_name)
-    remote_mod = remote.load_module(temp_file_name)
-    ctx = remote.cce()
-    arg_list = []
-    for a in args:
-        arg_list.append(akg.tvm.nd.array(a, ctx))
-    start_time = timer()
-    remote_mod(*arg_list)
-    ctx.sync()
-    if os.path.exists(temp_file_name):
-        os.remove(temp_file_name)
-    out_list = []
-    for i in outputs:
-        out = arg_list[len(arg_list) + i if i < 0 else i].asnumpy()
-        out_list.append(out)
-    # this time measure is no accurate now, to be improved soon
-    t = timer() - start_time
-    if not tuning:
-        return out_list[0] if len(out_list) == 1 else tuple(out_list)
-    stat_info = {"run_time": t}
-    return out_list[0] if len(out_list) == 1 else tuple(out_list), stat_info
-
-
-def mod_launch_rpc_thread(mode, mod, args, outputs, results, need_retry, retry, tuning=False):
-    """internal RPC thread, should be called by mod_launch_rpc_multithread."""
-    remoteevb = '0'
-    host = None
-    port = None
-    env_dic = os.environ
-    if env_dic.get('RPC_HOST') and env_dic.get('RPC_PORT'):
-        host = env_dic.get('RPC_HOST')
-        port = int(env_dic.get('RPC_PORT'))
-    else:
-        if mode == 'rpc_cloud':
-            logging.error("runtime_mode=rpc_cloud must set 1980 host ip and port!")
-            raise Exception("ERROR:runtime_mode=rpc_cloud must set 1980 host ip and port!")
-        remoteevb = dispatch(retry)
-        host = rpc_machine.get(remoteevb)[0]
-        port = rpc_machine.get(remoteevb)[1]
-
-    start_time = timer()
-    end_time = 0.0
-    logging.debug("rpc ip: %s, rpc port: %d", host, port)
-    try:
-        out_list = mod_launch_rpc_worker(mod, args, outputs, host, port, tuning=tuning)
-        end_time = timer()
-        t = end_time - start_time
-        if not env_dic.get('RPC_HOST'):
-            commit(remoteevb, 20 if t > 20 else t)
-        logging.info("===this round host is %s time is %f", host, (end_time - start_time))
-        results[retry] = out_list
-    except RuntimeError:
-        need_retry[retry] = True
-        end_time = timer()
-        logging.error("===Failed! this round host is %s time is %f", host, (end_time - start_time))
-        if not env_dic.get('RPC_HOST'):
-            commit(remoteevb, end_time - start_time + 20 * (retry + 1))
-        logging.error("rpc retry error: %d %s", retry, sys.exc_info())
-
-
-def _get_rpc_result(poll_count, threads, thread_index, poll_interval, need_retry, results, retried):
-    """Get rpc run result."""
-    while poll_count > 0:
-        poll_count -= 1
-        # wait for the newly created thread, because it is most likely to complete first
-        threads[thread_index].join(poll_interval)
-        for poll_index in range(thread_index + 1):
-            if not threads[poll_index].is_alive() and not need_retry[poll_index]:
-                return True, results[poll_index]
-            if need_retry[poll_index] and not retried[poll_index]:
-                logging.error("Thread %d exit with error, spawn a new thread immediately", poll_index)
-                poll_count = 0
-                retried[poll_index] = True
-    return False, False
-
-
-def mod_launch_rpc(mode, mod, args, outputs, tuning=False):
-    """
-    launch rpc or rpc_cloud module with retry.
-
-    Note:
-        To minimize waiting time of struggler RPC servers, we wait for a short timeout and spawn
-        a new thread after the timeout.
-        In normal case, RPC would complete before the short timeout, so, only one thread will be created.
-        When the RPC server is slow, we create multiple threads that run concurrently.
-        We wait for the first thread that successfully completes its work and return the result.
-        If a thread fails (an exception is raised), we spawn a new thread to retry.
-        Newly spawned threads will use different RPC servers.
-        We bound the maximum number of threads, i.e. maximum number of retries.
-    """
-    max_num_threads = 5
-
-    import operator
-    arg_filter = filter(lambda x: isinstance(x, np.ndarray), args)
-    arg_tensor = list(arg_filter)
-    tensor_size = reduce(operator.add, (reduce(operator.mul, arg.shape) for arg in arg_tensor))
-    expected_upload_speed = 5e6
-    expected_upload_time = int(tensor_size / expected_upload_speed)
-
-    timeout_before_spawning_new_thread = 200 + expected_upload_time
-    poll_interval = 1
-    thread_timeout = 400 + expected_upload_time * 3
-
-    load_rpc_server_info(mode)
-
-    threads = [None] * max_num_threads
-    results = [None] * max_num_threads
-    need_retry = [None] * max_num_threads
-    retried = [False] * max_num_threads
-    for thread_index in range(max_num_threads):
-        if thread_index > 0:
-            logging.error("Thread %d run for %d seconds, spawn a new thread to retry",
-                          (thread_index - 1), timeout_before_spawning_new_thread)
-        threads[thread_index] = Thread(target=mod_launch_rpc_thread,
-                                       args=(mode, mod, args, outputs, results, need_retry, thread_index, tuning))
-        # daemonize the thread to prevent long running threads from hanging the whole process
-        threads[thread_index].daemon = True
-        threads[thread_index].start()
-        poll_count = timeout_before_spawning_new_thread // poll_interval
-        has_res, res = _get_rpc_result(poll_count, threads, thread_index, poll_interval, need_retry, results, retried)
-        if has_res:
-            return res
-
-    logging.error("All %d threads are created, poll the threads until the first one exits normally, \
-                  or all threads exit abnormally or timeout", max_num_threads)
-    poll_count = thread_timeout // poll_interval
-    for _ in range(poll_count):
-        threads[max_num_threads - 1].join(poll_interval)
-        exit_thread_count = 0
-        for poll_index in range(max_num_threads):
-            if not threads[poll_index].is_alive() and not need_retry[poll_index]:
-                return results[poll_index]
-            if not threads[poll_index].is_alive():
-                exit_thread_count += 1
-            if exit_thread_count == max_num_threads:
-                logging.error("All %d threads exit abnormally", max_num_threads)
-                return None
-
-    logging.error("All %d threads timeout", max_num_threads)
-    return None
-
-
 def profiling_mode_run(kernel_name, args, outputs, tuning, device_id):
     """
     Function for collecting cycle data from device.
@@ -402,9 +193,7 @@ def profiling_mode_run(kernel_name, args, outputs, tuning, device_id):
         logging.error("OOPS, can't correctly parsing cycles!")
     TestUtils.record_cycle(cycle)
     logging.info('=====parsing cycles==============================')
-    if tuning:
-        return output_data, {'run_time': cycle}
-    return output_data
+    return output_data, {'run_time': cycle}
 
 
 def profiling_analyse(device_id, time_before_launch):
@@ -618,6 +407,20 @@ def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None
     if device_id == -1:
         device_id = int(os.environ.get("DEVICE_ID", 0))
 
+    # npu-inference process
+    if isinstance(mod, str):
+        kernel_name = mod
+        run_func = ascend_run
+        run_args = [kernel_name, args, outputs, device_id]
+        if os.environ.get("PROFILING_MODE") == "true":
+            run_func = profiling_mode_run
+            run_args = [kernel_name, args, outputs, tuning, device_id]
+            if os.environ.get("PROFILING_DIR", None) is None:
+                os.environ["PROFILING_DIR"] = "."
+                logging.info("[RUNTIME_WARNING] In profiling mode, while profiling dir is not set!Set to current dir by default.")
+        output = run_func(*run_args)
+        return output
+
     module = mod if mod.type_key == LLVM else mod.imported_modules[0]
     target = module.type_key
     if target == LLVM or target == CUDA:
@@ -635,8 +438,6 @@ def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None
             return output
         ra_util.get_ticks(stat_info)
         return output, stat_info
-    if mode in ('rpc', 'rpc_cloud'):
-        return mod_launch_rpc(mode, mod, args, outputs, tuning)
 
     # The air_cloud is the current default mode and needs to be modified in the future
     if mode == 'air_cloud':
@@ -658,7 +459,7 @@ def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None
         mod(*tvm_array)
         return tvm_array[-1].asnumpy()
 
-    raise ValueError("mode must be aic, rpc, aic_cloud, ca, compile_cloud, compile_mini, cpu, csim, ccesim or cdiff")
+    raise ValueError("mode must be aic, aic_cloud, ca, compile_cloud, compile_mini, cpu, csim, ccesim or cdiff")
 
 
 def _extract_shape_dtype(input_shapes, input_types):
@@ -1004,44 +805,6 @@ def _create_gpu_mod(s, op_var, target, shape_var, kernel_name, attrs, polyhedral
     return mod
 
 
-def _create_gpu_tuning_mod(sch_tmpl, shape_var, kernel_name, attrs, binds):
-    """Create tuning module on gpu."""
-    @autotvm.template
-    def _autotune_template():
-        s = sch_tmpl['schedule'](sch_tmpl['output'])
-        return s, op_var
-
-    # create autotune task
-    task = autotvm.task.create(_autotune_template, args=list(), target='cuda')
-    print("task config: ", task.config_space)
-
-    # set measure_option
-    measure_option = autotvm.measure_option(
-        builder=autotvm.LocalBuilder(),
-        runner=autotvm.LocalRunner(repeat=5, min_repeat_ms=150, timeout=4)
-    )
-
-    # Begin tuning, log records to file `kernel_name.log`
-    tuner = autotvm.tuner.RandomTuner(task)
-    if not os.path.exists(kernel_name + '.log'):
-        tuner.tune(n_trial=len(task.config_space),
-                   measure_option=measure_option,
-                   callbacks=[autotvm.callback.log_to_file(kernel_name + '.log')])
-
-    # query best config
-    dispatch_context = autotvm.apply_history_best(kernel_name + '.log')
-    best_config = dispatch_context.query(task.target, task.workload)
-    print("\nBest config is:")
-    print(best_config)
-
-    # apply best config
-    with autotvm.apply_history_best(kernel_name + '.log'):
-        s, op_var = _autotune_template()
-        mod = akg.build(s, op_var, "cuda", shape_var, name=kernel_name, attrs=attrs,
-                        polyhedral=False, binds=binds)
-    return mod
-
-
 def create_gpu_mod(sch_tmpl, s, op_func, op_var, shape_var, kernel_name, attrs, polyhedral, binds, dump_ir, dump_code,
                    tuning):
     """
@@ -1079,7 +842,7 @@ def create_gpu_mod(sch_tmpl, s, op_func, op_var, shape_var, kernel_name, attrs,
                 s = sch_tmpl['schedule'](sch_tmpl['output'])
                 mod = _create_gpu_mod(s, op_var, "cuda", shape_var, kernel_name, attrs, False, binds, dump_ir)
             else:
-                mod = _create_gpu_tuning_mod(sch_tmpl, shape_var, kernel_name, attrs, binds)
+                raise ValueError("Tuning is not supported.")
     else:
         mod = _create_gpu_mod(s, op_var, target, shape_var, kernel_name, attrs, polyhedral, binds, dump_ir)
     if dump_code:
@@ -1213,6 +976,10 @@ def op_build(op_func, input_shapes, input_types, op_attrs=None, kernel_name="",
         compute_func(s)
         polyhedral = False
 
+    if attrs.get("simple_mode"):
+        attrs.pop("simple_mode")
+        return s, inputs, output, attrs
+
     level = attrs.get("help_tiling") if attrs and "help_tiling" in attrs else None
     if tuning or (level is not None and level > help_tiling_level.get('None')):
         return gen_spaces_dim_key(op_func, args, s, op_var, kernel_name, attrs, polyhedral, tuning, target)
@@ -1231,10 +998,11 @@ def op_build(op_func, input_shapes, input_types, op_attrs=None, kernel_name="",
                             polyhedral=polyhedral, binds=binds)
             source_code = mod.get_source()
     elif target_name == CCE:
-        with akg.build_config(dump_pass_ir=dump_ir):
-            mod = akg.build(s, op_var, target, shape_var, name=kernel_name, attrs=attrs,
-                            polyhedral=polyhedral, binds=binds)
+        mod = npu_op_build(s, op_var, shape_var, kernel_name, binds, attrs, dump_ir, polyhedral)
+        if attrs.get("is_tbe_codegen"):
             source_code = mod.imported_modules[0].get_source()
+        else:
+            return mod
 
     if log_code:
         logging.debug("#################code####################")
@@ -1244,11 +1012,41 @@ def op_build(op_func, input_shapes, input_types, op_attrs=None, kernel_name="",
     return mod
 
 
+def npu_op_build(s, op_var, shape_var, kernel_name="", binds=None, attrs=None,
+                 dump_ir=True, polyhedral=True):
+    if attrs.get("is_tbe_codegen"):
+        # use akg + tbe compile
+        from akg.tvm import build_module
+        from akg.python.akg.utils.tbe_codegen_utils import build_tbe_codegen
+        if attrs is None:
+            attrs = {} 
+        attrs.update({"is_tbe_codegen":True})
+        binds, arg_list = build_module.get_binds(op_var)
+        stmt = akg.lower(s, op_var, shape_params=shape_var, name=kernel_name, binds=binds, attrs=attrs,
+            simple_mode=True, polyhedral=polyhedral, tuning=False, target="cce")
+
+        json_str = akg.tvm.save_json(stmt, "0.8.0")
+
+        args_json = []
+        for buf in enumerate(arg_list):
+            args_json.append(akg.tvm.save_json(buf, "0.8.0"))
+
+        is_success = build_tbe_codegen(kernel_name, json_str, args_json, attrs.get("dynamic", False))
+        if not is_success:
+            raise TypeError("npu_inference codegen failed.")
+        return kernel_name
+    else:
+        # use the whole akg complie
+        with akg.build_config(dump_pass_ir=dump_ir):
+            mod = akg.build(s, op_var, CCE, shape_var, name=kernel_name, attrs=attrs,
+                            polyhedral=polyhedral, binds=binds)
+        return mod
+
 def get_runtime_mode():
     """get runtime mode."""
     env_dic = os.environ
     if not env_dic.get('RUNTIME_MODE'):
-        mode = 'rpc_cloud'
+        mode = 'aic_cloud'
     else:
         mode = env_dic.get('RUNTIME_MODE')
     return mode
@@ -1265,7 +1063,7 @@ def get_profiling_mode():
 def product_is_mini():
     """check whether in mini environment."""
     mode = get_runtime_mode()
-    if mode in ('rpc', 'air', 'aic', 'compile_mini'):
+    if mode in ('air', 'aic', 'compile_mini'):
         return True
     return False
 

mindspore/_akg/akg/utils/op_dsl.py

@@ -906,6 +906,21 @@ def layout_transform_str(inputs, output, attr, op_type):
     return res
 
 
+def fast_gelu_str(inputs, output, attr):
+    x = inputs[0][0]["tensor_name"]
+    x_type = inputs[0][0]["data_type"]
+    res = "abs_x = np.absolute(%s)\n" % x
+    res += "sub1 = np.subtract(%s, abs_x)\n" % x
+    res += "mul1 = np.multiply(sub1, np.%s(0.851))\n" % x_type
+    res += "exp1 = np.exp(mul1)\n"
+    res += "tmp1 = np.multiply(%s, exp1)\n" % x
+    res += "mul2 = np.multiply(abs_x, np.%s(-1.702))\n" % x_type
+    res += "exp2 = np.exp(mul2)\n"
+    res += "tmp2 = np.add(exp2, np.%s(1))\n" % x_type
+    res += "%s = np.divide(tmp1, tmp2)\n" % output[0]["tensor_name"]
+    return res
+
+
 op_dsl = {
     "Custom": lambda inputs, output, attr: custom_str(inputs, output, attr),
     "ReduceSum": lambda inputs, output, attr: reduce_str(inputs, output, attr, "sum"),
@@ -1122,5 +1137,6 @@ op_dsl = {
                                                get_input(inputs[1][0])),
     "Pool2D": lambda inputs, output, attr: pool2d_str(inputs, output, attr),  
     "LayoutTransform": lambda inputs, output, attr: layout_transform_str(inputs, output, attr, "layout_transform"),
-    "Concat": lambda inputs, output, attr: concat_str(inputs, output, attr)
+    "Concat": lambda inputs, output, attr: concat_str(inputs, output, attr),
+    "FastGeLU": lambda inputs, output, attr: fast_gelu_str(inputs, output, attr)
 }

mindspore/_akg/akg/utils/result_analysis.py

@@ -351,16 +351,8 @@ def _collect_inputs(input_desc):
     return inputs
 
 
-def _get_op_attr(op_name, attrs, attr_name):
-    """Get op attr value."""
-    for attr in attrs:
-        if attr["name"] == attr_name:
-            return attr["value"]
-    raise ValueError("Can not find attr '{}' in op {}".format(attr_name, op_name))
-
-
 def precision_analyze(desc: dict, tensors):
-    exclude_op_list = ["Minimum", "Maximum", "Reshape", "ZerosLike", "Tile", "Select", "InplaceAssign", "Greater",
+    exclude_op_list = ["Minimum", "Maximum", "Reshape", "ZerosLike", "Tile", "Select", "Greater",
                        "SelectGT", "SelectLT", "LessEqual", "Less", "EquivFormat", "ExpandDims", "Transpose",
                        "TransData", "BroadcastTo", "Assign"]
     input_tensors = _collect_inputs(desc["input_desc"])
@@ -369,21 +361,9 @@ def precision_analyze(desc: dict, tensors):
     graph = {}
     ops = {}  # recorder the operator that generates the current output
     for op in desc["op_desc"]:
-        if op["name"] == "InplaceAssign":
-            output = IOInfo(op["input_desc"][0][0]["tensor_name"], op["input_desc"][0][0]["data_type"])
-            inputs = IOInfo(op["input_desc"][1][0]["tensor_name"], op["input_desc"][1][0]["data_type"])
-            graph[output] = [inputs]
-            ops[output] = op["name"]
-            fake_output = _get_op_attr(op["name"], op["attr"], "fake_output")
-            if not fake_output:
-                output = IOInfo(op["output_desc"][0]["tensor_name"], op["output_desc"][0]["data_type"])
-                inputs = IOInfo(op["input_desc"][2][0]["tensor_name"], op["input_desc"][2][0]["data_type"])
-                graph[output] = [inputs]
-                ops[output] = op["name"]
-        else:
-            output = IOInfo(op["output_desc"][0]["tensor_name"], op["output_desc"][0]["data_type"])
-            graph[output] = _collect_inputs(op["input_desc"])
-            ops[output] = op["name"]
+        output = IOInfo(op["output_desc"][0]["tensor_name"], op["output_desc"][0]["data_type"])
+        graph[output] = _collect_inputs(op["input_desc"])
+        ops[output] = op["name"]
 
     def _precision_reduce(x: IOInfo):
         if x in input_tensors:

mindspore/_akg/akg/utils/tbe_codegen_utils.py

@@ -0,0 +1,198 @@
+#!/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.
+import os
+import logging
+
+logging.getLogger().setLevel(logging.INFO)
+
+
+def copy_to_akg_kernel_meta(kernel_name, postfixs):
+    akg_kernel_mate_str = "akg_kernel_meta"
+    source = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', './'))
+    import shutil
+    target = source + "/" + akg_kernel_mate_str + "/" + kernel_name
+    source = source + "/" + "kernel_meta/" + kernel_name
+    if not os.path.exists(akg_kernel_mate_str):
+        os.mkdir(akg_kernel_mate_str)
+    for postfix in postfixs:
+        if os.path.exists(source + postfix):
+            try:
+                shutil.move(source + postfix, target + postfix)
+            except IOError as e:
+                logging.error("Unable to move file. {}".format(e))
+            except Exception as e:
+                logging.error("Unexpected error:", e)
+        else:
+            logging.info("Move {} fail, exit.".format(source + postfix))
+            return False
+    return True
+
+
+def clean_env():
+    import gc
+    import sys
+
+    imported_modules = set(sys.modules.keys())
+    for obj_key in imported_modules:
+        if "conda" in obj_key:
+            continue
+        if "akg" in obj_key or "topi" in obj_key or "tvm" in obj_key:
+            del sys.modules[obj_key]
+            try:
+                del globals()[obj_key]
+            except KeyError:
+                pass
+            try:
+                del locals()[obj_key]
+            except KeyError:
+                pass
+
+    gc.collect()
+
+
+def auto_init_soc(ascend_type):
+    from tbe.common.platform import set_current_compile_soc_info
+    set_current_compile_soc_info(ascend_type)
+
+
+def build_npu_for_akg(kernel_name,
+                      stmt=None,
+                      arg_list=None,
+                      is_dynamic=False,
+                      cfg=None,
+                      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.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
+    
+    set_current_build_config(tbe_workspace_size_list_length,
+                             tbe.tvm.runtime.cce_runtime.tbe_workspace_size_list_length())
+
+    if stmt is None or arg_list is None:
+        raise ValueError("No json, exit.")
+
+    func = tbe.tvm.tir.PrimFunc(arg_list, stmt)
+    mod = tbe.tvm.IRModule({kernel_name : func})
+    # _static_lower_phase_0
+    mod = transform.InjectSocVersion()(mod)
+    mod = transform.DeduceOpPlatform()(mod)
+    mod = transform.EmitInsn()(mod)
+
+    # phase 1 _static_lower_phase_emit_insn
+    mod = transform.InjectMultiCoreSync()(mod)
+    mod = transform.SplitCoproc()(mod)
+    mod = transform.SequenceSprInsn()(mod)
+
+    # phase 2
+    mod = transform.TikDoubleBufferSupport()(mod)
+    mod = transform.InjectPipeBuffer()(mod)
+    mod = transform.OptimizeDMA()(mod)
+    mod = transform.SubstituteInstr()(mod)
+    mod = transform.InjectAccessPtrMSG()(mod)
+    mod = transform.InjectPipe()(mod)
+    mod = transform.DeSequenceSprInsn()(mod)
+    mod = transform.CanonicalSimplify()(mod)
+    mod = transform.SetSPROptimizer()(mod)
+    if cfg[enable_const_fold]:
+        mod = transform.ConstantFolding()(mod)
+    if not simple_mode:
+        mod = transform.LoopPartition()(mod)
+    if cfg[disable_vectorize]:
+        mod = transform.SkipVectorize()(mod)
+    else:
+        mod = transform.VectorizeLoop()(mod)
+    mod = transform.InjectVirtualThread()(mod)
+
+    # phase 3 _static_lower_phase_3
+    mod = transform.StorageRewriteCCE()(mod)
+    mod = transform.ReorderProcess()(mod)
+    if cfg[tik] and cfg[dynamic_tik]:
+        mod = transform.TikDynamicShapeAllocMem()(mod)
+    mod = transform.UnrollLoop()(mod)
+
+    mod = transform.AutoFuseBuffer()(mod)
+    mod = transform.SetCacheMode()(mod)
+    mod = transform.Simplify()(mod)
+    mod = transform.GMConflictElimination()(mod)
+    mod = transform.MarkScalarCoreType()(mod)
+
+    # phase 4 _static_lower_phase_4
+    mod = transform.JumpInstructionElimination()(mod)
+    mod = transform.InjectSync()(mod)
+    mod = transform.PackIntrinArgConfig()(mod)
+    mod = transform.RemoveAccessPtrMSG()(mod)
+    mod = transform.Simplify()(mod)
+    mod = transform.GmAddrPrompt()(mod)
+    mod = transform.InsertCheckInvalidAccessOfDDR()(mod)
+    mod = transform.RemoveNoOp()(mod)
+    mod = transform.DeviceMark()(mod)
+    if cfg[instrument_bound_checkers]:
+        mod = transform.InstrumentBoundCheckers()(mod)
+    mod = transform.ConvertFloorDivToTruncDiv()(mod)
+    mod = transform.BuildVirtualCore()(mod)
+
+    _count_time(mod)
+    mod = transform.SplitCoreCode()(mod)
+    generate_cce_code(mod, "cce", None)
+
+
+def build_tbe_codegen(kernel_name, stmt_json, arg_json, ascend_type=None, is_dynamic=False):
+    import sys
+    copy_modules = sys.modules.copy()
+    clean_env()
+
+    print("build_cce_for_akg")
+    import tbe
+    from tbe.common.buildcfg.default_buildcfg import cce_default_static_build_config
+    from tbe.common.buildcfg.ascend import AscendPassContext
+    from tbe.common.buildcfg.buildcfg_mapping import dump_cce_code, save_temp_cce_file, disable_vectorize, \
+        instrument_bound_checkers, partition_const_loop, auto_unroll_max_step, auto_unroll_max_depth, \
+        auto_unroll_max_extent, unroll_explicit, dynamic_shape, enable_multicore_sync_with_atomic, \
+        kernel_meta_parent_dir
+    cfg = cce_default_static_build_config.copy()
+    cfg[dump_cce_code] = False
+    cfg[save_temp_cce_file] = True
+    cfg[disable_vectorize] = False
+    cfg[instrument_bound_checkers] = False
+    cfg[partition_const_loop] = False
+    cfg[auto_unroll_max_step] = 0
+    cfg[auto_unroll_max_depth] = 8
+    cfg[auto_unroll_max_extent] = 0
+    cfg[unroll_explicit] = True
+    cfg[dynamic_shape] = False
+    cfg[enable_multicore_sync_with_atomic] = True
+    cfg[kernel_meta_parent_dir] = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', './'))
+    if ascend_type is None:
+        ascend_type = "Ascend910"
+    auto_init_soc(ascend_type)
+
+    stmt = tbe.tvm.ir.load_json(stmt_json)
+    arg_list = []
+    for buff in arg_json:
+        arg_list.append(tbe.tvm.ir.load_json(buff))
+    with AscendPassContext(config=cfg):
+        build_npu_for_akg(kernel_name,
+                          stmt,
+                          arg_list,
+                          is_dynamic=is_dynamic,
+                          cfg=cfg)
+    postfixs = [".o", ".cce", ".json"]
+    is_success = copy_to_akg_kernel_meta(kernel_name, postfixs)
+    sys.modules = copy_modules
+    return is_success