mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.0__py3-none-any.whl

msprobe/core/compare/acc_compare.py

@@ -1,14 +1,30 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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 multiprocessing
 import os
-import json
 import pandas as pd
-from msprobe.core.common.file_utils import FileOpen
+from tqdm import tqdm
+from msprobe.core.common.file_utils import load_json
 from msprobe.core.common.const import CompareConst, Const
 from msprobe.core.common.exceptions import FileCheckException
 from msprobe.core.common.log import logger
-from msprobe.core.common.utils import add_time_with_xlsx, CompareException
+from msprobe.core.common.utils import add_time_with_xlsx, CompareException, check_op_str_pattern_valid
 from msprobe.core.common.file_utils import remove_path
-from msprobe.core.compare.check import check_graph_mode, check_struct_match, fuzzy_check_op
+from msprobe.core.compare.check import check_graph_mode, check_struct_match, fuzzy_check_op, check_dump_json_str, \
+                                        check_stack_json_str
 from msprobe.core.compare.highlight import find_compare_result_error_rows, highlight_rows_xlsx
 from msprobe.core.compare.utils import read_op, merge_tensor, get_un_match_accuracy, get_accuracy
 from msprobe.core.compare.multiprocessing_compute import _handle_multi_process, ComparisonResult, _save_cmp_result
@@ -21,10 +37,53 @@ class Comparator:
     
     def __init__(self):
         pass
+
+    @staticmethod
+    def get_result_md5_compare(ms_op_name, bench_op_name, npu_ops_all, bench_ops_all, *args):
+        result_item = [ms_op_name, bench_op_name, npu_ops_all.get(ms_op_name).get('struct')[0],
+                       bench_ops_all.get(bench_op_name).get('struct')[0],
+                       npu_ops_all.get(ms_op_name).get('struct')[1],
+                       bench_ops_all.get(bench_op_name).get('struct')[1],
+                       npu_ops_all.get(ms_op_name).get('struct')[2],
+                       bench_ops_all.get(bench_op_name).get('struct')[2],
+                       CompareConst.PASS if npu_ops_all.get(ms_op_name).get('struct')[2]
+                                            == bench_ops_all.get(bench_op_name).get('struct')[2]
+                       else CompareConst.DIFF]
+        if args[0]:
+            result_item.extend(args[1])
+        else:
+            result_item.append(CompareConst.NONE)
+        return result_item
+
+    @staticmethod
+    def calculate_summary_data(npu_summary_data, bench_summary_data, result_item):
+        err_msg = ""
+        start_idx = CompareConst.SUMMARY_COMPARE_RESULT_HEADER.index(CompareConst.MAX_DIFF)
+        warning_flag = False
+        for i, (npu_val, bench_val) in enumerate(zip(npu_summary_data, bench_summary_data)):
+            if isinstance(npu_val, (float, int)) and isinstance(bench_val, (float, int)):
+                diff = npu_val - bench_val
+                if bench_val != 0:
+                    relative = str(abs((diff / bench_val) * 100)) + '%'
+                else:
+                    relative = "N/A"
+                result_item[start_idx + i] = diff
+                result_item[start_idx + i + 4] = relative
+                magnitude_diff = abs(diff) / (max(abs(npu_val), abs(bench_val)) + 1e-10)
+                if magnitude_diff > 0.5:
+                    warning_flag = True
+            else:
+                result_item[start_idx + i] = CompareConst.NONE
+        accuracy_check = CompareConst.WARNING if warning_flag else ""
+        err_msg += "Need double check api accuracy." if warning_flag else ""
+        for i in range(start_idx, len(result_item)):
+            if str(result_item[i]) in ('inf', '-inf', 'nan'):
+                result_item[i] = f'{result_item[i]}\t'
+        result_item.append(accuracy_check)
+        result_item.append(err_msg)
     
     @classmethod
-    def make_result_table(cls,result, md5_compare, summary_compare, stack_mode):
-        header = []
+    def make_result_table(cls, result, md5_compare, summary_compare, stack_mode):
         if md5_compare:
             header = CompareConst.MD5_COMPARE_RESULT_HEADER[:]
         elif summary_compare:
@@ -47,17 +106,22 @@ class Comparator:
             else:
                 for row in result:
                     del row[-1]
-        result_df = pd.DataFrame(result, columns=header)
+        result_df = pd.DataFrame(result, columns=header, dtype='object')
         return result_df   
     
     @classmethod
-    def gen_merge_list(self, json_data, op_name,stack_json_data, summary_compare, md5_compare):
+    def gen_merge_list(cls, json_data, op_name, stack_json_data, summary_compare, md5_compare):
         op_data = json_data['data'][op_name]
+        check_dump_json_str(op_data, op_name)
         op_parsed_list = read_op(op_data, op_name)
-        if op_name in stack_json_data:
-            op_parsed_list.append({'full_op_name': op_name, 'full_info': stack_json_data[op_name]})
-        else:
-            op_parsed_list.append({'full_op_name': op_name, 'full_info': None})
+
+        stack_info = stack_json_data.get(op_name)
+        if stack_info is not None:
+            check_stack_json_str(stack_info, op_name)
+        op_parsed_list.append({
+            'full_op_name': op_name,
+            'full_info': stack_info
+        })
             
         merge_list = merge_tensor(op_parsed_list, summary_compare, md5_compare)
         return merge_list
@@ -67,7 +131,7 @@ class Comparator:
         b_op_name = bench_dict["op_name"]
         graph_mode = check_graph_mode(a_op_name[0], b_op_name[0])
         
-        frame_name = getattr(self,"frame_name")
+        frame_name = getattr(self, "frame_name")
         if frame_name == "PTComparator":
             from msprobe.pytorch.compare.match import graph_mapping
             if graph_mode:
@@ -94,11 +158,11 @@ class Comparator:
                 return n_index, len(bench_queue) - 1
         return -1, -1
     
-    def compare_process(self, file_handles, stack_mode, fuzzy_match, summary_compare=False, md5_compare=False):
-        npu_json_handle, bench_json_handle, stack_json_handle = file_handles
-        npu_json_data = json.load(npu_json_handle)
-        bench_json_data = json.load(bench_json_handle)
-        stack_json_data = json.load(stack_json_handle)
+    def compare_process(self, file_lists, stack_mode, fuzzy_match, summary_compare=False, md5_compare=False):
+        npu_json_path, bench_json_path, stack_json_path = file_lists
+        npu_json_data = load_json(npu_json_path)
+        bench_json_data = load_json(bench_json_path)
+        stack_json_data = load_json(stack_json_path)
 
         if fuzzy_match:
             logger.warning("This task uses fuzzy matching, which may affect the accuracy of the comparison.")
@@ -114,14 +178,19 @@ class Comparator:
         last_npu_ops_len = 0
         last_bench_ops_len = 0
 
+        npu_api_nums = len(npu_json_data['data'])
+        progress_bar = tqdm(total=npu_api_nums, desc="API/Module Read Progress", unit="item", ncols=100)
+
         while True:
             if not read_err_npu and not read_err_bench:
                 break
             try:
                 last_npu_ops_len = len(npu_ops_queue)
                 op_name_npu = next(ops_npu_iter)
+                check_op_str_pattern_valid(op_name_npu)
                 read_err_npu = True
-                npu_merge_list = self.gen_merge_list(npu_json_data,op_name_npu,stack_json_data,summary_compare,md5_compare)
+                npu_merge_list = self.gen_merge_list(npu_json_data, op_name_npu, stack_json_data,
+                                                     summary_compare, md5_compare)
                 if npu_merge_list:
                     npu_ops_queue.append(npu_merge_list)
             except StopIteration:
@@ -129,12 +198,16 @@ class Comparator:
             try:
                 last_bench_ops_len = len(bench_ops_queue)
                 op_name_bench = next(ops_bench_iter)
-                bench_merge_list = self.gen_merge_list(bench_json_data,op_name_bench,stack_json_data,summary_compare,md5_compare)
+                check_op_str_pattern_valid(op_name_bench)
+                bench_merge_list = self.gen_merge_list(bench_json_data, op_name_bench, stack_json_data,
+                                                       summary_compare, md5_compare)
                 if bench_merge_list:
                     bench_ops_queue.append(bench_merge_list)
             except StopIteration:
                 read_err_bench = False
 
+            progress_bar.update(1)
+
             # merge all boolean expressions
             both_empty = not npu_ops_queue and not bench_ops_queue
             no_change = (len(npu_ops_queue) == last_npu_ops_len) and (len(bench_ops_queue) == last_bench_ops_len)
@@ -163,7 +236,91 @@ class Comparator:
                 
         result_df = self.make_result_table(result, md5_compare, summary_compare, stack_mode)
         return result_df
-    
+
+    def merge_data(self, json_data, stack_json_data, summary_compare, md5_compare):
+        ops_all = {}
+        for op_name in json_data.get('data', {}):
+            merge_list = self.gen_merge_list(json_data, op_name, stack_json_data, summary_compare,
+                                             md5_compare)
+            if merge_list:
+                input_index, output_index = 0, 0
+                for index, input_or_output in enumerate(merge_list['op_name']):
+                    input_or_output_list = input_or_output.split(Const.SEP)
+                    data_name = merge_list.get('data_name')
+                    data_name = data_name[index] if data_name else None
+                    if Const.INPUT in input_or_output_list or Const.KWARGS in input_or_output_list:
+                        ops_all[input_or_output] = {'struct': merge_list.get('input_struct')[input_index],
+                                                    'summary': merge_list.get('summary')[index],
+                                                    'data_name': data_name,
+                                                    'stack_info': merge_list.get('stack_info')}
+                        input_index += 1
+
+                    elif Const.OUTPUT in input_or_output_list:
+                        ops_all[input_or_output] = {'struct': merge_list.get('output_struct')[output_index],
+                                                    'summary': merge_list.get('summary')[index],
+                                                    'data_name': data_name,
+                                                    'stack_info': merge_list.get('stack_info')}
+                        output_index += 1
+        return ops_all
+
+    def get_accuracy(self, npu_ops_all, bench_ops_all, summary_compare, md5_compare):
+        result = []
+        for ms_op_name, bench_op_name in self.data_mapping_dict.items():
+            if ms_op_name in npu_ops_all and bench_op_name in bench_ops_all:
+                npu_stack_info = npu_ops_all.get(ms_op_name).get("stack_info", None)
+                bench_stack_info = bench_ops_all.get(bench_op_name).get("stack_info", None)
+                has_stack = npu_stack_info and bench_stack_info
+                if md5_compare:
+                    result.append(self.get_result_md5_compare(ms_op_name, bench_op_name, npu_ops_all,
+                                                              bench_ops_all, has_stack, npu_stack_info))
+                    continue
+                if summary_compare:
+                    result_item = [ms_op_name, bench_op_name, npu_ops_all.get(ms_op_name).get('struct')[0],
+                                   bench_ops_all.get(bench_op_name).get('struct')[0],
+                                   npu_ops_all.get(ms_op_name).get('struct')[1],
+                                   bench_ops_all.get(bench_op_name).get('struct')[1],
+                                   " ", " ", " ", " ", " ", " ", " ", " "]
+                else:
+                    result_item = [ms_op_name, bench_op_name, npu_ops_all.get(ms_op_name).get('struct')[0],
+                                   bench_ops_all.get(bench_op_name).get('struct')[0],
+                                   npu_ops_all.get(ms_op_name).get('struct')[1],
+                                   bench_ops_all.get(bench_op_name).get('struct')[1],
+                                   " ", " ", " ", " ", " "]
+                npu_summary_data = npu_ops_all.get(ms_op_name).get("summary")
+                result_item.extend(npu_summary_data)
+                bench_summary_data = bench_ops_all.get(bench_op_name).get("summary")
+                result_item.extend(bench_summary_data)
+                if summary_compare:
+                    self.calculate_summary_data(npu_summary_data, bench_summary_data, result_item)
+                else:
+                    result_item.append(CompareConst.ACCURACY_CHECK_YES)
+                    result_item.append("")
+                if has_stack:
+                    result_item.extend(npu_stack_info)
+                else:
+                    result_item.append(CompareConst.NONE)
+                if not (summary_compare or md5_compare):
+                    result_item.append(npu_ops_all.get(ms_op_name).get("data_name", None))
+                result.append(result_item)
+            elif ms_op_name not in npu_ops_all:
+                logger.warning(f'Can not find npu op name : `{ms_op_name}` in npu dump json file.')
+            elif bench_op_name not in npu_ops_all:
+                logger.warning(f'Can not find bench op name : `{bench_op_name}` in bench dump json file.')
+        return result
+
+    def compare_process_custom(self, file_lists, stack_mode, summary_compare=False, md5_compare=False):
+        npu_json_path, bench_json_path, stack_json_path = file_lists
+        npu_json_data = load_json(npu_json_path)
+        bench_json_data = load_json(bench_json_path)
+        stack_json_data = load_json(stack_json_path)
+
+        npu_ops_all = self.merge_data(npu_json_data, stack_json_data, summary_compare, md5_compare)
+        bench_ops_all = self.merge_data(bench_json_data, stack_json_data, summary_compare, md5_compare)
+
+        result = self.get_accuracy(npu_ops_all, bench_ops_all, summary_compare, md5_compare)
+        result_df = self.make_result_table(result, md5_compare, summary_compare, stack_mode)
+        return result_df
+
     def compare_by_op(self, npu_op_name, bench_op_name, op_name_mapping_dict, input_param):
         npu_bench_name_list = op_name_mapping_dict[npu_op_name]
         data_name = npu_bench_name_list[1]
@@ -178,9 +335,11 @@ class Comparator:
                 if frame_name == "MSComparator":
                     n_value = read_npy_data(input_param.get("npu_dump_data_dir"), npu_op_name + Const.NUMPY_SUFFIX)
                     if self.cross_frame:
-                        b_value = read_npy_data(input_param.get("bench_dump_data_dir"), bench_op_name + Const.PT_SUFFIX, load_pt_file=True)
+                        b_value = read_npy_data(input_param.get("bench_dump_data_dir"),
+                                                bench_op_name + Const.PT_SUFFIX, load_pt_file=True)
                     else:
-                        b_value = read_npy_data(input_param.get("bench_dump_data_dir"), bench_op_name + Const.NUMPY_SUFFIX)
+                        b_value = read_npy_data(input_param.get("bench_dump_data_dir"),
+                                                bench_op_name + Const.NUMPY_SUFFIX)
                 else:
                     n_value = read_npy_data(input_param.get("npu_dump_data_dir"), npu_op_name + Const.PT_SUFFIX)
                     b_value = read_npy_data(input_param.get("bench_dump_data_dir"), bench_op_name + Const.PT_SUFFIX)
@@ -237,19 +396,31 @@ class Comparator:
         file_path = os.path.join(os.path.realpath(output_path), file_name)
         remove_path(file_path)
         highlight_dict = {'red_rows': [], 'yellow_rows': []}
-        
-        with FileOpen(input_parma.get("npu_json_path"), "r") as npu_json, \
-                FileOpen(input_parma.get("bench_json_path"), "r") as bench_json, \
-                FileOpen(input_parma.get("stack_json_path"), "r") as stack_json:
+
+        npu_json = input_parma.get("npu_json_path")
+        bench_json = input_parma.get("bench_json_path")
+        stack_json = input_parma.get("stack_json_path")
+        if self.data_mapping:
+            result_df = self.compare_process_custom([npu_json, bench_json, stack_json], stack_mode,
+                                                    summary_compare, md5_compare)
+        else:
             result_df = self.compare_process([npu_json, bench_json, stack_json], stack_mode, fuzzy_match,
-                                        summary_compare, md5_compare)
+                                             summary_compare, md5_compare)
+
+        if not result_df.values.tolist():
+            logger.warning("Can`t match any op.")
+            return
 
         if not md5_compare and not summary_compare:
             result_df = self._do_multi_process(input_parma, result_df)
+
+        logger.info("Highlight suspicious API/Module start.")
         find_compare_result_error_rows(result_df, highlight_dict, summary_compare, md5_compare)
         highlight_rows_xlsx(result_df, highlight_dict, file_path)
+        logger.info("Highlight suspicious API/Module finish.")
+
         if auto_analyze:
-            advisor = Advisor(result_df, output_path)
+            advisor = Advisor(result_df, output_path, suffix)
             advisor.analysis()
     
     def compare_ops(self, idx, dump_path_dict, result_df, lock, input_param):
@@ -265,13 +436,14 @@ class Comparator:
             bench_op_name = result_df.iloc[i, 1]
             if is_print_compare_log:
                 logger.info("start compare: {}".format(npu_op_name))
-            cos_sim, max_abs_err, max_relative_err, one_thousand_err_ratio, five_thousand_err_ratio, err_msg = self.compare_by_op(
-                npu_op_name, bench_op_name, dump_path_dict, input_param)
+            cos_sim, max_abs_err, max_relative_err, one_thousand_err_ratio, five_thousand_err_ratio, err_msg = \
+                self.compare_by_op(npu_op_name, bench_op_name, dump_path_dict, input_param)
             if is_print_compare_log:
                 logger.info(
-                    "[{}] Compare result: cosine {}, max_abs_err {}, max_relative_err {}, {}, one_thousand_err_ratio {}, "
-                    "five_thousand_err_ratio {}".format(npu_op_name, cos_sim, max_abs_err, max_relative_err, err_msg,
-                                                        one_thousand_err_ratio, five_thousand_err_ratio))
+                    "[{}] Compare result: cosine {}, max_abs_err {}, max_relative_err {}, {}, \
+                    one_thousand_err_ratio {}, "
+                    "five_thousand_err_ratio {}".format(npu_op_name, cos_sim, max_abs_err, max_relative_err,
+                                                        err_msg, one_thousand_err_ratio, five_thousand_err_ratio))
             cos_result.append(cos_sim)
             max_err_result.append(max_abs_err)
             max_relative_err_result.append(max_relative_err)
@@ -290,9 +462,10 @@ class Comparator:
 
         return _save_cmp_result(idx, cr, result_df, lock)   
     
-    def _do_multi_process(self,input_parma, result_df):
+    def _do_multi_process(self, input_parma, result_df):
         try:
-            result_df = _handle_multi_process(self.compare_ops, input_parma, result_df, multiprocessing.Manager().RLock())
+            result_df = _handle_multi_process(self.compare_ops, input_parma, result_df,
+                                              multiprocessing.Manager().RLock())
             return result_df
         except ValueError as e:
             logger.error('result dataframe is not found.')

msprobe/core/compare/check.py

@@ -1,5 +1,22 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from msprobe.core.common.log import logger
-from msprobe.core.compare.utils import rename_api 
+from msprobe.core.compare.utils import rename_api
+from msprobe.core.common.utils import check_op_str_pattern_valid, CompareException
+from msprobe.core.common.const import Const
 
 
 dtype_mapping = {
@@ -34,8 +51,15 @@ def check_struct_match(npu_dict, bench_dict, cross_frame=False):
     if not is_match:
         if len(npu_struct_in) == 0 or len(bench_struct_in) == 0 or len(npu_struct_in) != len(bench_struct_in):
             return False
-        struct_in_is_match = check_type_shape_match(npu_struct_in, bench_struct_in)
-        struct_out_is_match = check_type_shape_match(npu_struct_out, bench_struct_out)
+        try:
+            struct_in_is_match = check_type_shape_match(npu_struct_in, bench_struct_in)
+            struct_out_is_match = check_type_shape_match(npu_struct_out, bench_struct_out)
+        except CompareException as error:
+            err_msg = f'index out of bounds error occurs in npu or bench api, please check!\n' \
+                      f'npu_dict: {npu_dict}' \
+                      f'bench_dict: {bench_dict}'
+            logger.error(err_msg)
+            raise CompareException(CompareException.INDEX_OUT_OF_BOUNDS_ERROR) from error
         is_match = struct_in_is_match and struct_out_is_match
     return is_match
 
@@ -43,17 +67,27 @@ def check_struct_match(npu_dict, bench_dict, cross_frame=False):
 def check_type_shape_match(npu_struct, bench_struct):
     shape_type_match = False
     for npu_type_shape, bench_type_shape in zip(npu_struct, bench_struct):
-        npu_type = npu_type_shape[0]
-        npu_shape = npu_type_shape[1]
-        bench_type = bench_type_shape[0]
-        bench_shape = bench_type_shape[1]
+        try:
+            npu_type = npu_type_shape[0]
+            npu_shape = npu_type_shape[1]
+            bench_type = bench_type_shape[0]
+            bench_shape = bench_type_shape[1]
+        except IndexError as error:
+            logger.error(f'length of npu_type_shape: {npu_type_shape} and bench_type_shape: {bench_type_shape} '
+                         f'should both be 2, please check!')
+            raise CompareException(CompareException.INDEX_OUT_OF_BOUNDS_ERROR) from error
         shape_match = npu_shape == bench_shape
         type_match = npu_type == bench_type
         if not type_match:
-            ms_type=[["Float16", "Float32"], ["Float32", "Float16"],["Float16", "BFloat16"],["BFloat16", "Float16"]] 
-            torch_type=[["torch.float16", "torch.float32"], ["torch.float32", "torch.float16"],
-                                ["torch.float16", "torch.bfloat16"], ["torch.bfloat16", "torch.float16"]]
-            if ([npu_type, bench_type] in ms_type)or  ([npu_type, bench_type] in torch_type):                    
+            ms_type = [
+                [Const.FLOAT16, Const.FLOAT32], [Const.FLOAT32, Const.FLOAT16],
+                [Const.FLOAT16, Const.BFLOAT16], [Const.BFLOAT16, Const.FLOAT16]
+            ]
+            torch_type = [
+                [Const.TORCH_FLOAT16, Const.TORCH_FLOAT32], [Const.TORCH_FLOAT32, Const.TORCH_FLOAT16],
+                [Const.TORCH_FLOAT16, Const.TORCH_BFLOAT16], [Const.TORCH_BFLOAT16, Const.TORCH_FLOAT16]
+            ]
+            if ([npu_type, bench_type] in ms_type) or ([npu_type, bench_type] in torch_type):
                 type_match = True
             else:
                 type_match = False
@@ -64,9 +98,9 @@ def check_type_shape_match(npu_struct, bench_struct):
 
 
 def check_graph_mode(a_op_name, b_op_name):
-    if "Aten" in a_op_name and "Aten" not in b_op_name:
+    if Const.ATEN in a_op_name and Const.ATEN not in b_op_name:
         return True
-    if "Aten" not in a_op_name and "Aten" in b_op_name:
+    if Const.ATEN not in a_op_name and Const.ATEN in b_op_name:
         return True
     return False
 
@@ -83,13 +117,64 @@ def fuzzy_check_op(npu_name_list, bench_name_list):
 
 
 def fuzzy_check_name(npu_name, bench_name):
-    if "forward" in npu_name and "forward" in bench_name:
-        is_match = rename_api(npu_name, "forward") == rename_api(bench_name, "forward")
-    elif "backward" in npu_name and "backward" in bench_name:
-        is_match = rename_api(npu_name, "backward") == rename_api(bench_name, "backward")
+    if Const.FORWARD in npu_name and Const.FORWARD in bench_name:
+        is_match = rename_api(npu_name, Const.FORWARD) == rename_api(bench_name, Const.FORWARD)
+    elif Const.BACKWARD in npu_name and Const.BACKWARD in bench_name:
+        is_match = rename_api(npu_name, Const.BACKWARD) == rename_api(bench_name, Const.BACKWARD)
     else:
         is_match = npu_name == bench_name
     return is_match
 
 
+def check_dump_json_str(op_data, op_name):
+    input_list = op_data.get(Const.INPUT_ARGS, None) if op_data.get(Const.INPUT_ARGS, None) else op_data.get(
+        Const.INPUT, None)
+    input_kwargs = op_data.get(Const.INPUT_KWARGS, None)
+    output_list = op_data.get(Const.OUTPUT, None)
+
+    args = [input_list, input_kwargs, output_list]
+    for arg in args:
+        if not arg:
+            continue
+        if isinstance(arg, dict):
+            check_json_key_value(arg, op_name)
+        else:
+            for ele in arg:
+                if not ele:
+                    continue
+                check_json_key_value(ele, op_name)
+
+
+def check_json_key_value(input_output, op_name, depth=0):
+    if depth > Const.MAX_DEPTH:
+        logger.error(f"string check of data info of {op_name} exceeds the recursion limit.")
+        return
+    if isinstance(input_output, list):
+        for item in input_output:
+            check_json_key_value(item, op_name, depth+1)
+    elif isinstance(input_output, dict):
+        for key, value in input_output.items():
+            if isinstance(value, dict):
+                check_json_key_value(value, op_name, depth+1)
+            else:
+                valid_key_value(key, value, op_name)
+
 
+def valid_key_value(key, value, op_name):
+    if key == "shape" and not isinstance(value, (list, tuple)):
+        logger.error(f"shape of input or output of {op_name} is not list or tuple, please check!")
+        raise CompareException(CompareException.INVALID_OBJECT_TYPE_ERROR)
+    elif key == "requires_grad" and not isinstance(value, bool):
+        logger.error(f"requires_grad of input or output of {op_name} is not bool, please check!")
+        raise CompareException(CompareException.INVALID_OBJECT_TYPE_ERROR)
+    else:
+        check_op_str_pattern_valid(value, op_name)
+
+
+def check_stack_json_str(stack_info, op_name):
+    if isinstance(stack_info, list):
+        for item in stack_info:
+            check_op_str_pattern_valid(item, op_name, stack=True)
+    else:
+        logger.error(f"Expected stack_info to be a list, but got {type(stack_info).__name__} for '{op_name}'")
+        raise CompareException(CompareException.INVALID_OBJECT_TYPE_ERROR)

msprobe/core/compare/compare_cli.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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 json
 from msprobe.core.common.file_utils import FileOpen, check_file_type
 from msprobe.core.common.const import FileCheckConst, Const
@@ -23,8 +38,11 @@ def compare_cli(args):
         input_param["bench_json_path"] = input_param.pop("bench_path")
         input_param["stack_json_path"] = input_param.pop("stack_path")
         if frame_name == Const.PT_FRAMEWORK:
+            kwargs = {
+                "data_mapping": args.data_mapping
+            }
             compare(input_param, args.output_path, stack_mode=args.stack_mode, auto_analyze=auto_analyze,
-                    fuzzy_match=args.fuzzy_match)
+                    fuzzy_match=args.fuzzy_match, **kwargs)
         else:
             kwargs = {
                 "stack_mode": args.stack_mode,
@@ -32,6 +50,8 @@ def compare_cli(args):
                 "fuzzy_match": args.fuzzy_match,
                 "cell_mapping": args.cell_mapping,
                 "api_mapping": args.api_mapping,
+                "data_mapping": args.data_mapping,
+                "layer_mapping": args.layer_mapping
             }
 
             ms_compare(input_param, args.output_path, **kwargs)

msprobe/core/compare/highlight.py

@@ -1,5 +1,21 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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 math
 import abc
+import re
 from collections import namedtuple
 import numpy as np
 import openpyxl
@@ -7,7 +23,7 @@ from openpyxl.styles import PatternFill
 from msprobe.core.common.utils import get_header_index
 from msprobe.core.common.file_utils import save_workbook
 from msprobe.core.common.log import logger
-from msprobe.core.common.const import CompareConst
+from msprobe.core.common.const import CompareConst, FileCheckConst
 
 
 class HighlightCheck(abc.ABC):
@@ -34,9 +50,11 @@ class CheckOneThousandErrorRatio(HighlightCheck):
     def apply(self, info, color_columns, summary_compare=True):
         api_in, api_out, num = info
         one_thousand_index = get_header_index('One Thousandth Err Ratio', summary_compare)
-        if not isinstance(api_in[one_thousand_index], (float, int)) or not isinstance(api_out[one_thousand_index], (float, int)):
+        if (not isinstance(api_in[one_thousand_index], (float, int)) or
+                not isinstance(api_out[one_thousand_index], (float, int))):
             return
-        if api_in[one_thousand_index] > CompareConst.ONE_THOUSAND_ERROR_IN_RED and api_out[one_thousand_index] < CompareConst.ONE_THOUSAND_ERROR_OUT_RED:
+        if (api_in[one_thousand_index] > CompareConst.ONE_THOUSAND_ERROR_IN_RED and
+                api_out[one_thousand_index] < CompareConst.ONE_THOUSAND_ERROR_OUT_RED):
             color_columns.red.append(num)
         elif api_in[one_thousand_index] - api_out[one_thousand_index] > CompareConst.ONE_THOUSAND_ERROR_DIFF_YELLOW:
             color_columns.yellow.append(num)
@@ -66,7 +84,8 @@ class CheckMaxRelativeDiff(HighlightCheck):
             return
         if output_max_relative_diff > CompareConst.MAX_RELATIVE_OUT_RED:
             color_columns.red.append(num)
-        elif output_max_relative_diff > CompareConst.MAX_RELATIVE_OUT_YELLOW and input_max_relative_diff < CompareConst.MAX_RELATIVE_IN_YELLOW:
+        elif (output_max_relative_diff > CompareConst.MAX_RELATIVE_OUT_YELLOW and
+              input_max_relative_diff < CompareConst.MAX_RELATIVE_IN_YELLOW):
             color_columns.yellow.append(num)
 
 
@@ -193,7 +212,8 @@ def find_compare_result_error_rows(result_df, highlight_dict, summary_compare, m
             input_num = num
         else:
             output_num = num
-        find_error_rows(result[start:start + input_num + output_num], start, input_num, highlight_dict, summary_compare, md5_compare)
+        find_error_rows(result[start:start + input_num + output_num], start, input_num, highlight_dict,
+                        summary_compare, md5_compare)
 
 
 def highlight_rows_xlsx(result_df, highlight_dict, file_path):
@@ -205,12 +225,16 @@ def highlight_rows_xlsx(result_df, highlight_dict, file_path):
 
     # write header
     for j, col_name in enumerate(result_df.columns, start=1):
+        if not csv_value_is_valid(col_name):
+            raise RuntimeError(f"Malicious value [{col_name}] is not allowed to be written into the xlsx: {file_path}.")
         ws.cell(row=1, column=j, value=col_name)
 
     for i, row in enumerate(result_df.iterrows(), start=2):
         for j, value in enumerate(row[1], start=1):
             if not isinstance(value, (float, int)):
                 value = f'{str(value)}\t' if str(value) in ('inf', '-inf', 'nan') else str(value)
+            if not csv_value_is_valid(value):
+                raise RuntimeError(f"Malicious value [{value}] is not allowed to be written into the xlsx: {file_path}.")
             ws.cell(row=i, column=j, value=f'{str(value)}\t' if str(value) in ('inf', '-inf', 'nan') else value)
 
             if (i - 2) in highlight_dict['red_rows']:
@@ -221,3 +245,15 @@ def highlight_rows_xlsx(result_df, highlight_dict, file_path):
                                                             end_color=CompareConst.YELLOW, fill_type="solid")
 
     save_workbook(wb, file_path)
+
+
+def csv_value_is_valid(value: str) -> bool:
+    if not isinstance(value, str):
+        return True
+    try:
+        # -1.00 or +1.00 should be consdiered as digit numbers
+        float(value)
+    except ValueError:
+        # otherwise, they will be considered as formular injections
+        return not bool(re.compile(FileCheckConst.CSV_BLACK_LIST).search(value))
+    return True