mindspore 2.2.0__cp37-cp37m-manylinux1_x86_64.whl → 2.2.10__cp37-cp37m-manylinux1_x86_64.whl

mindspore/.commit_id

@@ -1 +1 @@
-__commit_id__ = '[sha1]:9390851d,[branch]:(HEAD,origin/r2.2,r2.2)'
+__commit_id__ = '[sha1]:93fa4635,[branch]:(HEAD,origin/r2.2.10-1,r2.2.10-1)'

mindspore/_akg/akg/composite/build_module.py

@@ -21,14 +21,14 @@ from collections.abc import Iterable
 import akg
 from akg import tvm
 from tvm.autotvm.env import AutotvmGlobalScope
+from akg.utils.util import parse_workspace_map
 from akg.utils.tbe_codegen_utils import build_tbe_codegen
 from akg.utils.kernel_exec import ReturnType, is_symbolic_tiling
 from .split_stitch import split_stitch_attr
 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
 def generate_trait(desc):
     """
     generate trait of kernel description
@@ -523,17 +523,6 @@ def _build_to_module(desc_s, desc_d, attrs=None, poly=True):
 
     return _cpp_build(attrs, process, poly, segment_tree, segment_infos)
 
-
-def _get_ascend_type(desc):
-    if "target_info" not in desc.keys():
-        return None
-
-    target_info_type = desc["target_info"]
-    if target_info_type.get("arch"):
-        return target_info_type.get("arch")
-    return None
-
-
 def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
     """
     build kernel with compute description in json format
@@ -650,7 +639,7 @@ def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
         ConstructType.NORMAL: _normal_postprocess,
     }
     process = desc_d_in["process"]
-    ascend_type = _get_ascend_type(desc_d_in)
+    ascend_type = get_ascend_type(desc_d_in)
     ascend_type_to_section = {"Ascend910A": "1.6", "Ascend310P3": "1.7",
                               "Ascend910B1": "2.1", "Ascend910B2": "2.2", "Ascend910B3": "2.3", "Ascend910B4": "2.4"}
     if ascend_type is not None:
@@ -659,6 +648,7 @@ def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
         config_func(section)
         if section >= "2.1":
             attr["is_tbe_codegen"] = True
+            attr["pragma_modshift"] = True
     segment_tree, segment_infos = get_construct_args(desc_s_in, attr, post_funcs)
     poly = True
     res = _cpp_build(attr, process, poly, segment_tree, segment_infos)
@@ -668,8 +658,12 @@ def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
         args_json = []
         for buf in res[1]:
             args_json.append(akg.tvm.save_json(buf, "0.8.0"))
+        
+        workspace_dict = parse_workspace_map(res[2])
+        if workspace_dict is not None:
+            attr["workspace"] = workspace_dict
 
-        is_success = build_tbe_codegen(kernel_name, stmt_json, args_json, ascend_type, attr.get("dynamic", False))
+        is_success = build_tbe_codegen(kernel_name, stmt_json, args_json, attr, ascend_type)
         if not is_success:
             raise TypeError("npu_inference codegen failed.")
         return kernel_name

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

@@ -0,0 +1,76 @@
+import os
+import re
+import subprocess
+from enum import Enum
+
+from .file_manager import FileManager
+from .path_manager import PathManager
+
+
+class CANNDataEnum(Enum):
+    OP_SUMMARY = 0
+    NPU_MEMORY = 1
+    MSPROF_TIMELINE = 2
+    STEP_TRACE = 3
+    GE_MEMORY_RECORD = 4
+    GE_OPERATOR_MEMORY = 5
+    L2_CACHE = 6
+    AI_CPU = 7
+    COMMUNICATION = 8
+    MATRIX = 9
+
+
+class CANNFileParser:
+    COMMAND_SUCCESS = 0
+    ACL_TO_NPU = "acl_to_npu"
+    START_FLOW = "s"
+    END_FLOW = "f"
+    SUMMARY = "summary"
+    TIMELINE = "timeline"
+    ANALYZE = "analyze"
+    CANN_DATA_MATCH = {
+        CANNDataEnum.OP_SUMMARY: [r"^op_summary_\d+_\d+\.csv", r"^op_summary_\d+_\d+_\d+\.csv",
+                                  r"^op_summary_\d+_\d+_\d+_\d+\.csv"],
+    }
+
+    def __init__(self, profiler_path: str):
+        self._cann_path = PathManager.get_cann_path(profiler_path)
+        self._file_dict = {}
+        self._file_dispatch()
+
+    def export_cann_profiling(self):
+        if not os.path.isdir(self._cann_path):
+            return
+        self._del_summary_and_timeline_data()
+        completed_process = subprocess.run(["msprof", "--export=on", f"--output={self._cann_path}"],
+                                           capture_output=True)
+        if completed_process.returncode != self.COMMAND_SUCCESS:
+            raise RuntimeError(
+                f"Export CANN Profiling data failed, please verify that the ascend-toolkit is installed and set-env.sh "
+                f"is sourced. or you can execute the command to confirm the CANN Profiling export result: "
+                f"msprof --export=on --output={self._cann_path}")
+
+
+    def get_file_list_by_type(self, file_type: CANNDataEnum) -> set:
+        return self._file_dict.get(file_type, set())
+
+    def _file_dispatch(self):
+        all_file_list = PathManager.get_device_all_file_list_by_type(self._cann_path, self.SUMMARY)
+        all_file_list += PathManager.get_device_all_file_list_by_type(self._cann_path, self.TIMELINE)
+        all_file_list += PathManager.get_analyze_all_file(self._cann_path, self.ANALYZE)
+        for file_path in all_file_list:
+            if not os.path.isfile(file_path):
+                continue
+            for data_type, re_match_exp_list in self.CANN_DATA_MATCH.items():
+                for re_match_exp in re_match_exp_list:
+                    if re.match(re_match_exp, os.path.basename(file_path)):
+                        self._file_dict.setdefault(data_type, set()).add(file_path)
+
+    def _del_summary_and_timeline_data(self):
+        device_path = PathManager.get_device_path(self._cann_path)
+        if not device_path:
+            return
+        summary_path = os.path.join(device_path, "summary")
+        timeline_path = os.path.join(device_path, "timeline")
+        FileManager.remove_file_safety(summary_path)
+        FileManager.remove_file_safety(timeline_path)

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

@@ -0,0 +1,56 @@
+import csv
+import json
+import os.path
+import shutil
+from warnings import warn
+
+MAX_FILE_SIZE = 1024 * 1024 * 1024 * 10
+MAX_CSV_SIZE = 1024 * 1024 * 1024 * 5
+
+class FileManager:
+    @classmethod
+    def file_read_all(cls, file_path: str, mode: str = "r") -> any:
+        if not os.path.isfile(file_path):
+            return ''
+        file_size = os.path.getsize(file_path)
+        if file_size <= 0:
+            return ''
+        if file_size > MAX_FILE_SIZE:
+            warn(f"The file size exceeds the preset value {MAX_FILE_SIZE / 1024 / 1024}MB, "
+                 f"please check the file: {file_path}")
+            return ''
+        try:
+            with open(file_path, mode) as file:
+                return file.read()
+        except Exception:
+            raise RuntimeError(f"Can't read file: {file_path}")
+
+    @classmethod
+    def read_csv_file(cls, file_path: str, class_bean: any) -> list:
+        if not os.path.isfile(file_path):
+            return []
+        file_size = os.path.getsize(file_path)
+        if file_size <= 0:
+            return []
+        if file_size > MAX_CSV_SIZE:
+            warn(f"The file size exceeds the preset value {MAX_CSV_SIZE / 1024 / 1024}MB, "
+                 f"please check the file: {file_path}")
+            return []
+        result_data = []
+        try:
+            with open(file_path, newline="") as csv_file:
+                reader = csv.DictReader(csv_file)
+                for row in reader:
+                    result_data.append(class_bean(row))
+        except Exception:
+            raise RuntimeError(f"Failed to read the file: {file_path}")
+        return result_data
+
+    @classmethod
+    def remove_file_safety(cls, path: str):
+        if os.path.exists(path):
+            try:
+                shutil.rmtree(path)
+            except Exception:
+                print(f"[WARNING] [{os.getpid()}] profiler.py: Can't remove the directory: {path}")
+

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

@@ -0,0 +1,23 @@
+from .op_summary_headers import OpSummaryHeaders
+
+class OpSummaryBean:
+    headers = []
+
+    def __init__(self, data: list):
+        self._data = data
+
+    @property
+    def row(self) -> list:
+        row = []
+        read_headers = OpSummaryBean.headers if OpSummaryBean.headers else self._data.keys()
+        for field_name in read_headers:
+            row.append(self._data.get(field_name, ""))
+        return row
+
+    @property
+    def ts(self) -> float:
+        return float(self._data.get(OpSummaryHeaders.TASK_START_TIME, 0))
+
+    @property
+    def all_headers(self) -> list:
+        return list(self._data.keys())

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

@@ -0,0 +1,8 @@
+class OpSummaryHeaders(object):
+    # op_summary
+    TASK_START_TIME = "Task Start Time(us)"
+    AIC_TOTAL_CYCLES = "aic_total_cycles"
+    AIV_TOTAL_CYCLES = "aiv_total_cycles"
+    TASK_DURATION = "Task Duration(us)"
+    OP_SUMMARY_SHOW_HEADERS = ["Op Name", "OP Type", "Task Type", TASK_START_TIME, TASK_DURATION,
+                               "Task Wait Time(us)", "Block Dim" ,AIC_TOTAL_CYCLES, AIV_TOTAL_CYCLES]

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

@@ -0,0 +1,42 @@
+from .op_summary_headers import OpSummaryHeaders
+from .file_manager import FileManager
+from .op_summary_bean import OpSummaryBean
+from .cann_file_parser import CANNFileParser, CANNDataEnum
+
+
+class OpSummaryParser():
+
+    def __init__(self, profiler_path: str):
+        self._profiler_path = profiler_path
+
+    @classmethod
+    def _project_map_for_headers(cls, input_headers: list):
+        project_map_dict = {OpSummaryHeaders.OP_SUMMARY_SHOW_HEADERS[i]: OpSummaryHeaders.OP_SUMMARY_KERNEL_BASE_HEADERS[i] for i in
+                            range(len(OpSummaryHeaders.OP_SUMMARY_SHOW_HEADERS))}
+        output_headers = []
+        for header in input_headers:
+            if header in project_map_dict:
+                output_headers.append(project_map_dict.get(header))
+            else:
+                output_headers.append(header)
+        return output_headers
+
+    def generate_op_summary_data(self) -> dict:
+        op_summary_file_set = CANNFileParser(self._profiler_path).get_file_list_by_type(CANNDataEnum.OP_SUMMARY)
+        summary_data = []
+        for file_path in op_summary_file_set:
+            all_data = FileManager.read_csv_file(file_path, OpSummaryBean)
+            if all_data:
+                OpSummaryBean.headers = OpSummaryHeaders.OP_SUMMARY_SHOW_HEADERS
+                output_headers = OpSummaryBean.headers
+                summary_data.extend([data.row for data in all_data])
+            else:
+                raise RuntimeError("parse op summary csv failed.")
+
+        return self.create_dict(summary_data[0],output_headers)
+
+    def create_dict(self,summary_data,headers):
+        summary_dict={}
+        for i in range(len(summary_data)):
+            summary_dict[headers[i]] = summary_data[i]
+        return summary_dict

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

@@ -0,0 +1,65 @@
+import os
+import re
+
+class PathManager:
+
+    @classmethod
+    def get_cann_path(cls, profiler_path: str) -> str:
+        sub_dirs = os.listdir(os.path.realpath(profiler_path))
+        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):
+                return sub_path
+        return ""
+
+    @classmethod
+    def get_device_path(cls, cann_path: str) -> str:
+        sub_dirs = os.listdir(os.path.realpath(cann_path))
+        for sub_dir in sub_dirs:
+            sub_path = os.path.join(cann_path, sub_dir)
+            if os.path.isdir(sub_path) and re.match(r"^device_\d", sub_dir):
+                return sub_path
+        return ""
+
+    @classmethod
+    def get_start_info_path(cls, cann_path: str) -> str:
+        start_info_path = os.path.join(cann_path, "host", "start_info")
+        if os.path.exists(start_info_path):
+            return start_info_path
+        device_path = cls.get_device_path(cann_path)
+        if not device_path:
+            return ""
+        device_path_split = os.path.basename(device_path).split("_")
+        if len(device_path_split) != 2:
+            return ""
+        start_info_file = f"start_info.{device_path_split[1]}"
+        start_info_path = os.path.join(device_path, start_info_file)
+        if os.path.exists(start_info_path):
+            return start_info_path
+        return ""
+
+    @classmethod
+    def get_device_all_file_list_by_type(cls, profiler_path: str, summary_or_timeline: str) -> list:
+        file_list = []
+        _path = os.path.join(cls.get_device_path(profiler_path), summary_or_timeline)
+        if not os.path.isdir(_path):
+            return file_list
+        sub_files = os.listdir(os.path.realpath(_path))
+        if not sub_files:
+            return file_list
+        for sub_file in sub_files:
+            file_list.append(os.path.join(_path, sub_file))
+        return file_list
+
+    @classmethod
+    def get_analyze_all_file(cls, profiler_path: str, analyze: str) -> list:
+        file_list = []
+        _path = os.path.join(profiler_path, analyze)
+        if not os.path.isdir(_path):
+            return file_list
+        sub_files = os.listdir(os.path.realpath(_path))
+        if not sub_files:
+            return file_list
+        for sub_file in sub_files:
+            file_list.append(os.path.join(_path, sub_file))
+        return file_list

mindspore/_akg/akg/utils/kernel_exec.py

@@ -43,10 +43,12 @@ from akg.utils import custom_tiling as ct_util
 from akg.utils import validation_check as vc_util
 from akg.utils.dsl_create import TensorUtils
 from akg.utils.util import parse_kwargs
-from akg.backend.parsing_profiling_data import HWTSLogParser
+from akg.backend.parsing_profiling_data import HWTSLogParser, max_time_consume
 from akg.backend.parsing_profiling_data import validate_and_normalize_path
 from akg.backend import aic_model
-
+from .ascend_profilier.cann_file_parser import CANNFileParser
+from .ascend_profilier.op_summary_parser import OpSummaryParser
+from .ascend_profilier.op_summary_headers import OpSummaryHeaders
 sh = logging.StreamHandler(sys.stdout)
 logging.getLogger().addHandler(sh)
 logging.getLogger().setLevel(logging.INFO)
@@ -169,7 +171,7 @@ def gen_name_kernel(kernel, dtype, shapes):
     return res
 
 
-def profiling_mode_run(kernel_name, args, outputs, tuning, device_id):
+def profiling_mode_run(kernel_name, args, outputs, tuning, device_id, arch=None):
     """
     Function for collecting cycle data from device.
 
@@ -180,21 +182,45 @@ def profiling_mode_run(kernel_name, args, outputs, tuning, device_id):
         tuning: tuning model.
         device_id: device_id on device.
     """
-    akg.tvm.get_global_func("ascend_start_profiling")(device_id)
+    akg.tvm.get_global_func("ascend_start_profiling")(kernel_name)
     time_before_launch = time.time()
     output_data = ascend_run(kernel_name, args, outputs, device_id)
     akg.tvm.get_global_func("ascend_stop_profiling")()
-
-    cycle = profiling_analyse(device_id, time_before_launch)
-    logging.info('=====parsing cycles==============================')
+    cycle = 0
+    if arch is not None and "910B" in arch:
+        # for ascend910B profiling
+        cycle = profiling_analyse_910B(time_before_launch)
+    else:
+        cycle = profiling_analyse(device_id, time_before_launch)
+    logging.info('=====Task Duration(us)==============================')
     if cycle != PROF_ERROR_CODE:
         logging.info(cycle)
     else:
-        logging.error("OOPS, can't correctly parsing cycles!")
+        logging.error("OOPS, can't correctly Task Duration!")
     TestUtils.record_cycle(cycle)
-    logging.info('=====parsing cycles==============================')
+    logging.info('=====Task Duration(us)==============================')
     return output_data, {'run_time': cycle}
 
+def  profiling_analyse_910B(time_before_launch):
+    public_path = os.getenv('PROFILING_DIR')
+    if public_path is None:
+        raise RuntimeError("Environment PROFILING_DIR not set!")
+    public_path = validate_and_normalize_path(public_path)
+    CANNFileParser(public_path).export_cann_profiling()
+    cann_file_parser = OpSummaryParser(public_path)
+    profiler_file = cann_file_parser._profiler_path
+    logging.debug("prof file is: %s", os.path.basename(profiler_file))
+    file_create_time = os.path.getctime(profiler_file)
+    if file_create_time < time_before_launch:
+        raise RuntimeError("The PROF file is too old")
+    datas:dict = cann_file_parser.generate_op_summary_data()
+    task_duration = float(datas.get(OpSummaryHeaders.TASK_DURATION,max_time_consume))
+    # # aic_total_cycles means ai core cycle
+    # # aiv_total_cycles means ai vector cycle
+    # aiv_total_cycle = int(datas.get(OpSummaryHeaders.AIV_TOTAL_CYCLES,max_time_consume))
+    # aic_total_cycle = int(datas.get(OpSummaryHeaders.AIC_TOTAL_CYCLES,max_time_consume))
+    # return aiv_total_cycle+aic_total_cycle
+    return task_duration
 
 def profiling_analyse(device_id, time_before_launch):
     """analyse profiling."""
@@ -348,12 +374,12 @@ def get_kernel_name_from_mod(mod):
     return kernel_name
 
 
-def mod_launch_ascend_profiling(mod, args, outputs=(-1,), tuning=False, device_id=-1):
+def mod_launch_ascend_profiling(mod, args, outputs=(-1,), tuning=False, device_id=-1, arch=None):
     gc.collect()
     if device_id == -1:
         device_id = int(os.environ.get("DEVICE_ID", 0))
     kernel_name = get_kernel_name_from_mod(mod)
-    return profiling_mode_run(kernel_name, args, outputs, tuning, device_id)
+    return profiling_mode_run(kernel_name, args, outputs, tuning, device_id, arch=arch)
 
 
 def mod_launch_default(mod, args, outputs=(-1,), target=CUDA, tuning=False, device_id=-1, repeat_time=400):
@@ -387,7 +413,7 @@ def mod_launch_default(mod, args, outputs=(-1,), target=CUDA, tuning=False, devi
 
 
 @func_time_required
-def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None, repeat_time=400):
+def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None, repeat_time=400, arch=None):
     """
     unified run CCE kernel api.
 
@@ -398,7 +424,7 @@ def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None
         tuning (bool): tuning model.
         device_id: device_id on device.
         expect: when mode in ["compile_cloud", "compile_mini"], return it.
-
+        arch: Ascend arch type
     Returns:
         output numpy array, or tuple of numpy array if multi-output.
     """
@@ -414,7 +440,7 @@ def mod_launch(mod, args, outputs=(-1,), tuning=False, device_id=-1, expect=None
         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]
+            run_args = [kernel_name, args, outputs, tuning, device_id, arch]
             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.")
@@ -1031,7 +1057,7 @@ def npu_op_build(s, op_var, shape_var, kernel_name="", binds=None, attrs=None,
         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))
+        is_success = build_tbe_codegen(kernel_name, json_str, args_json, attrs)
         if not is_success:
             raise TypeError("npu_inference codegen failed.")
         return kernel_name

mindspore/_akg/akg/utils/tbe_codegen_utils.py

@@ -15,22 +15,40 @@
 # limitations under the License.
 import os
 import logging
+from akg.utils.util import write_code
+import json
 
 logging.getLogger().setLevel(logging.INFO)
 
+def create_directory(dir):
+    if not os.path.exists(dir):
+        os.mkdir(dir)
 
-def copy_to_akg_kernel_meta(kernel_name, postfixs):
-    akg_kernel_mate_str = "akg_kernel_meta"
+def set_workspace_for_json(json_path, workspace_dict):
+    if workspace_dict is None:
+        return
+    with open(json_path, 'r') as f:
+        js_info = json.loads(f.read())
+    js_info["workspace"] = workspace_dict
+    write_code(js_info, json_path)
+
+
+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
     source = source + "/" + "kernel_meta/" + kernel_name
-    if not os.path.exists(akg_kernel_mate_str):
-        os.mkdir(akg_kernel_mate_str)
+    if source == target:
+        return True
+    akg_kernel_mate_str = os.path.abspath(akg_kernel_mate_str)
+    create_directory(akg_kernel_mate_str)
     for postfix in postfixs:
         if os.path.exists(source + postfix):
             try:
                 shutil.move(source + postfix, target + postfix)
+                if postfix == ".json" and workspace_dict is not None:
+                    set_workspace_for_json(target + postfix, workspace_dict)
             except IOError as e:
                 logging.error("Unable to move file. {}".format(e))
             except Exception as e:
@@ -152,7 +170,7 @@ def build_npu_for_akg(kernel_name,
     generate_cce_code(mod, "cce", None)
 
 
-def build_tbe_codegen(kernel_name, stmt_json, arg_json, ascend_type=None, is_dynamic=False):
+def build_tbe_codegen(kernel_name, stmt_json, arg_json, attr, ascend_type=None):
     import sys
     copy_modules = sys.modules.copy()
     clean_env()
@@ -186,6 +204,9 @@ def build_tbe_codegen(kernel_name, stmt_json, arg_json, ascend_type=None, is_dyn
     arg_list = []
     for buff in arg_json:
         arg_list.append(tbe.tvm.ir.load_json(buff))
+
+    is_dynamic = attr.get("dynamic", False)
+    workspace_dict = attr.get("workspace", None)
     with AscendPassContext(config=cfg):
         build_npu_for_akg(kernel_name,
                           stmt,
@@ -193,6 +214,6 @@ def build_tbe_codegen(kernel_name, stmt_json, arg_json, ascend_type=None, is_dyn
                           is_dynamic=is_dynamic,
                           cfg=cfg)
     postfixs = [".o", ".cce", ".json"]
-    is_success = copy_to_akg_kernel_meta(kernel_name, postfixs)
+    is_success = copy_to_akg_kernel_meta(kernel_name, postfixs, workspace_dict)
     sys.modules = copy_modules
     return is_success

mindspore/_akg/akg/utils/util.py

@@ -19,6 +19,7 @@ import os
 import json
 import types
 import akg.tvm
+from akg.utils.format_transform import get_shape
 
 
 def parse_int_const(value):
@@ -29,6 +30,35 @@ def parse_int_const(value):
     return None
 
 
+def parse_workspace_map(attrs):
+    if not isinstance(attrs, akg.tvm.container.Map):
+        return None
+    if "workspace" not in attrs:
+        return None
+
+    workspace_map = attrs["workspace"]
+    if not isinstance(workspace_map, akg.tvm.container.Map):
+        return None
+    if "num" not in workspace_map or "size" not in workspace_map:
+        return None
+
+    worksapce_num =  parse_int_const(workspace_map["num"])
+    if not isinstance(workspace_map["size"], akg.tvm.container.Array):
+        return None
+
+    workspace_size = []
+    tt = get_shape(workspace_map["size"])
+    for item in tt:
+        workspace_size.append(parse_int_const(item))
+    if len(workspace_size) != worksapce_num:
+        return None
+
+    workspace_dict = {
+        "num": worksapce_num,
+        "size": workspace_size
+    }
+    return workspace_dict
+
 def parse_workspace(workspace):
     if not isinstance(workspace, akg.tvm.container.Map):
         return None
@@ -63,3 +93,11 @@ def write_code(js_dict, fname):
     with os.fdopen(os.open(fname, os.O_WRONLY | os.O_CREAT, 0o400), 'w') as f:
         json.dump(js_dict, f, sort_keys=True, indent=4, separators=(',', ':'))
 
+def get_ascend_type(desc):
+    if "target_info" not in desc.keys():
+        return None
+
+    target_info_type = desc["target_info"]
+    if target_info_type.get("arch"):
+        return target_info_type.get("arch")
+    return None

mindspore/_checkparam.py

@@ -720,9 +720,9 @@ def check_value_type(arg_name, arg_value, valid_types, prim_name=None):
         type_names = [t.__name__ if hasattr(t, '__name__') else str(t) for t in valid_types]
         num_types = len(valid_types)
         msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-        type_name_msg = f'{type_names if num_types > 1 else type_names[0]}'
-        msg = f'type of \'{arg_name}\' should be{"one of " if num_types > 1 else ""} \'{type_name_msg}\''
-        raise TypeError(f'{msg_prefix} {msg}, but got type \'{type(arg_value).__name__}\'.')
+        raise TypeError(f'{msg_prefix} type of \'{arg_name}\' should be {"one of " if num_types > 1 else ""}' \
+                        f'\'{type_names if num_types > 1 else type_names[0]}\', ' \
+                        f'but got type \'{type(arg_value).__name__}\'.')
 
     # Notice: bool is subclass of int, so `check_value_type('x', True, [int])` will check fail, and
     #         `check_value_type('x', True, [bool, int])` will check pass