mindstudio-probe 1.1.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

msprobe/core/compare/highlight.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -13,19 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
 import abc
+import math
+import multiprocessing
 import re
 from collections import namedtuple
+
 import numpy as np
 import openpyxl
 from openpyxl.styles import PatternFill
+from openpyxl.utils.dataframe import dataframe_to_rows
 from tqdm import tqdm
-from msprobe.core.common.utils import get_header_index
+
+from msprobe.core.common.const import CompareConst, Const
 from msprobe.core.common.file_utils import save_workbook
 from msprobe.core.common.log import logger
-from msprobe.core.common.const import CompareConst, FileCheckConst, Const
-from msprobe.core.common.utils import safe_get_value
+from msprobe.core.common.utils import get_header_index, safe_get_value
+from msprobe.core.compare.utils import table_value_is_valid, get_name_and_state, CompareException
 
 
 class HighlightCheck(abc.ABC):
@@ -44,6 +48,7 @@ def add_highlight_row_info(color_list, num, highlight_err_msg):
 
 class CheckOrderMagnitude(HighlightCheck):
     """检查Max diff的数量级差异"""
+
     def apply(self, info, color_columns, dump_mode):
         api_in, api_out, num = info
         max_diff_index = get_header_index(CompareConst.MAX_DIFF if dump_mode == Const.SUMMARY
@@ -54,12 +59,13 @@ class CheckOrderMagnitude(HighlightCheck):
         out_order = 0 if abs(api_out[max_diff_index]) < 1 else math.log10(abs(api_out[max_diff_index]))
         if out_order - in_order >= CompareConst.ORDER_MAGNITUDE_DIFF_YELLOW:
             add_highlight_row_info(color_columns.yellow, num,
-                                   "maximum absolute error of both input and output exceed 1, "
+                                   "maximum absolute error of both input/parameters and output exceed 1, "
                                    "with the output larger by an order of magnitude")
 
 
 class CheckOneThousandErrorRatio(HighlightCheck):
     """检查千分误差比率"""
+
     def apply(self, info, color_columns, dump_mode):
         api_in, api_out, num = info
         one_thousand_index = get_header_index(CompareConst.ONE_THOUSANDTH_ERR_RATIO, dump_mode)
@@ -69,15 +75,17 @@ class CheckOneThousandErrorRatio(HighlightCheck):
         if (api_in[one_thousand_index] > CompareConst.ONE_THOUSAND_ERROR_IN_RED and
                 api_out[one_thousand_index] < CompareConst.ONE_THOUSAND_ERROR_OUT_RED):
             add_highlight_row_info(color_columns.red, num,
-                                   "The input's one thousandth err ratio exceeds 0.9, while the output's is below 0.6")
+                                   "The input/parameters's one thousandth err ratio exceeds 0.9, "
+                                   "while the output's is below 0.6")
         elif api_in[one_thousand_index] - api_out[one_thousand_index] > CompareConst.ONE_THOUSAND_ERROR_DIFF_YELLOW:
             add_highlight_row_info(color_columns.yellow, num,
                                    "The output's one thousandth err ratio decreases by more than 0.1 "
-                                   "compared to the input's")
+                                   "compared to the input/parameters's")
 
 
 class CheckCosineSimilarity(HighlightCheck):
     """检查余弦相似度"""
+
     def apply(self, info, color_columns, dump_mode):
         api_in, api_out, num = info
         cosine_index = get_header_index(CompareConst.COSINE, dump_mode)
@@ -85,17 +93,21 @@ class CheckCosineSimilarity(HighlightCheck):
             return
         if api_in[cosine_index] - api_out[cosine_index] > CompareConst.COSINE_DIFF_YELLOW:
             add_highlight_row_info(color_columns.yellow, num,
-                                   "The output's cosine decreases by more than 0.1 compared to the input's")
+                                   "The output's cosine decreases by more than 0.1 "
+                                   "compared to the input/parameters's")
 
 
 class CheckMaxRelativeDiff(HighlightCheck):
     """检查最大相对差异"""
+
     def apply(self, info, color_columns, dump_mode):
         api_in, api_out, num = info
         max_diff_index = get_header_index(CompareConst.MAX_DIFF, dump_mode)
         bench_max_index = get_header_index(CompareConst.BENCH_MAX, dump_mode)
-        input_max_relative_diff = np.abs(np.divide(api_in[max_diff_index], max(0.01, api_in[bench_max_index])))
-        output_max_relative_diff = np.abs(np.divide(api_out[max_diff_index], max(0.01, api_out[bench_max_index])))
+        input_max_relative_diff = np.abs(
+            np.divide(api_in[max_diff_index], max(Const.FLOAT_EPSILON, api_in[bench_max_index])))
+        output_max_relative_diff = np.abs(
+            np.divide(api_out[max_diff_index], max(Const.FLOAT_EPSILON, api_out[bench_max_index])))
         if not isinstance(input_max_relative_diff, (float, int)) or not isinstance(output_max_relative_diff,
                                                                                    (float, int)):
             return
@@ -104,11 +116,13 @@ class CheckMaxRelativeDiff(HighlightCheck):
         elif (output_max_relative_diff > CompareConst.MAX_RELATIVE_OUT_YELLOW and
               input_max_relative_diff < CompareConst.MAX_RELATIVE_IN_YELLOW):
             add_highlight_row_info(color_columns.yellow, num,
-                                   "The output's maximum relative error exceeds 0.1, while the input's is below 0.01")
+                                   "The output's maximum relative error exceeds 0.1, "
+                                   "while the input/parameters's is below 0.01")
 
 
 class CheckOverflow(HighlightCheck):
     """检查是否存在溢出"""
+
     def apply(self, info, color_columns, dump_mode):
         line, num = info
         npu_max_index = get_header_index(CompareConst.NPU_MAX, dump_mode)
@@ -141,9 +155,24 @@ class HighlightRules:
         "check_order_magnitude": CheckOrderMagnitude(),
         "check_max_relative_diff": CheckMaxRelativeDiff(),
     }
-    
 
-def find_error_rows(result, last_len, n_num_input, highlight_dict, dump_mode):
+
+def check_indices_numeric(api_items, indices: list):
+    """检查指定索引处的值是否都为数字类型（int 或 float）"""
+    return all(isinstance(api_items[i], (float, int)) for i in indices)
+
+
+def apply_comparison_rules(api_info, dump_mode, color_columns):
+    """output与input/params的比较"""
+    if dump_mode == Const.SUMMARY:
+        for rule in HighlightRules.summary_compare_rules.values():
+            rule.apply(api_info, color_columns, dump_mode)
+    else:
+        for rule in HighlightRules.compare_rules.values():
+            rule.apply(api_info, color_columns, dump_mode)
+
+
+def find_error_rows(result, api_batch, highlight_dict, dump_mode):
     """找到单个API中需要高亮的行"""
     if dump_mode == Const.MD5:
         return
@@ -158,35 +187,34 @@ def find_error_rows(result, last_len, n_num_input, highlight_dict, dump_mode):
     ColorColumns = namedtuple('ColorColumns', ['red', 'yellow'])
     color_columns = ColorColumns(red=red_lines, yellow=yellow_lines)
 
+    api_batch_start = api_batch.start  # result_df的input起始全局索引
+    api_batch_params_end_index = api_batch.params_end_index  # result_df的params结束全局索引 + 1
+    api_batch_output_end_index = api_batch.output_end_index  # result_df的output结束全局索引 + 1
+    api_batch_params_slice_index_local = api_batch_params_end_index - api_batch_start  # result的params结束局部切片索引
+    api_batch_output_slice_index_local = api_batch_output_end_index - api_batch_start  # result的output结束局部切片索引
+
     # 对单行API的输入或输出进行误差判断
     for i, line in enumerate(result):
-        num = last_len + i
-        line_info = LineInfo(line_data=line, num_pointer=num)
+        index = api_batch_start + i
+        line_info = LineInfo(line_data=line, num_pointer=index)
         for rule in HighlightRules.basic_rules.values():
             rule.apply(line_info, color_columns, dump_mode)
 
     # 对API的输出与输入比较，进行误差判断
-    for n, api_out in enumerate(result[n_num_input:len(result)]):
-        num = last_len + n_num_input + n
-        if num in red_lines:
+    for n, api_out in enumerate(result[api_batch_params_slice_index_local: api_batch_output_slice_index_local]):
+        index = api_batch_start + api_batch_params_slice_index_local + n
+        # 单行检查只有溢出检查（红色），如果已经溢出，不进一步检查
+        if index in red_lines:
             continue
-        if not isinstance(api_out[npu_max_index], (float, int)) \
-                or not isinstance(api_out[bench_max_index], (float, int)) \
-                or not isinstance(api_out[max_diff_index], (float, int)):
+        if not check_indices_numeric(api_out, [npu_max_index, bench_max_index, max_diff_index]):
             continue
-        for _, api_in in enumerate(result[0:n_num_input]):
-            if not isinstance(api_in[npu_max_index], (float, int)) \
-                    or not isinstance(api_in[bench_max_index], (float, int)) \
-                    or not isinstance(api_in[max_diff_index], (float, int)):
-                continue
 
-            api_info = ApiInfo(api_input=api_in, api_output=api_out, num_pointer=num)
-            if dump_mode == Const.SUMMARY:
-                for rule in HighlightRules.summary_compare_rules.values():
-                    rule.apply(api_info, color_columns, dump_mode)
-            else:
-                for rule in HighlightRules.compare_rules.values():
-                    rule.apply(api_info, color_columns, dump_mode)
+        # input/parameters的比较检查, 这里api_in包括input、parameters
+        for _, api_in in enumerate(result[0: api_batch_params_slice_index_local]):
+            if not check_indices_numeric(api_in, [npu_max_index, bench_max_index, max_diff_index]):
+                continue
+            api_info = ApiInfo(api_input=api_in, api_output=api_out, num_pointer=index)
+            apply_comparison_rules(api_info, dump_mode, color_columns)
 
     red_lines_num_set = {x[0] for x in red_lines}
     yellow_lines_num_set = {x[0] for x in yellow_lines}
@@ -196,78 +224,148 @@ def find_error_rows(result, last_len, n_num_input, highlight_dict, dump_mode):
     highlight_dict.get('yellow_lines', []).extend(yellow_lines)
 
 
-def get_name_and_state(name):
-    """Get api/module name and state"""
-    if Const.INPUT in name:
-        api_name = name.split(Const.INPUT)[0]
-        state = Const.INPUT
+class ApiBatch:
+    def __init__(self, api_name: str, start: int):
+        self.api_name = api_name
+        self.start = start
+        self.input_len = 1  # input的数量
+        self.params_end_index = start + 1  # params的结束index
+        self.output_end_index = start + 1  # output的结束index
+        self.params_grad_end_index = start + 1  # params_grad的结束index
+        # 内部state的标志("input", "output", "parameters", "parameters_grad"),
+        # 用于控制计算input_len, output_end_index, params_end_index, self.params_grad_end_index
+        self._state = Const.INPUT  # api_batch初始化为input
+
+    def set_state(self, state: str):
+        """设置当前状态"""
+        if state in {Const.INPUT, Const.OUTPUT, Const.KWARGS, Const.PARAMS, Const.PARAMS_GRAD}:
+            self._state = state
+        else:
+            raise ValueError(f"Invalid state: {state}")
+
+    def increment(self, state: str):
+        self.set_state(state)
+        if self._state == Const.INPUT or self._state == Const.KWARGS:
+            self.input_len += 1
+            self.params_end_index += 1
+            self.output_end_index += 1
+        if self._state == Const.PARAMS:
+            self.params_end_index += 1
+            self.output_end_index += 1
+        if self._state == Const.OUTPUT:
+            self.output_end_index += 1
+        self.params_grad_end_index += 1
+
+
+def api_batches_update(api_batches, api_name, state, index):
+    """
+    当一个api的所有item更新完后，input, output的索引范围：
+    input: [start: start+input_len]
+    output: [start+input_len: output_end_index]
+    params: [output_end_index: params_end_index]
+    """
+    if not api_batches:
+        api_batches.append(ApiBatch(api_name, index))
     else:
-        api_name = name.split(Const.OUTPUT)[0]
-        state = Const.OUTPUT
-    return api_name, state
+        api_batch = api_batches[-1]
+        if api_batch.api_name == api_name or (
+                not re.search(Const.REGEX_FORWARD_BACKWARD, api_name) and api_name in api_batch.api_name):
+            try:
+                api_batch.increment(state)
+            except ValueError as e:
+                logger.error(f"api_batch: {api_batch} with invalid state, please check! {e}")
+                raise CompareException(CompareException.INVALID_STATE_ERROR) from e
+        else:
+            api_batches.append(ApiBatch(api_name, index))
 
 
 def find_compare_result_error_rows(result_df, highlight_dict, dump_mode):
     """将dataframe根据API分组，并找到有误差的算子用于高亮"""
     result = result_df.values
-    start, input_num, output_num, end = 0, 0, 0, len(result_df)
-    last_api_name, last_state = None, None
-    num, last_len = 0, 0
-    progress_bar = tqdm(total=len(result), desc="API/Module Analyse Progress", unit="item", ncols=100)
-    for res_i in result:
+    api_batches = []
+    for i, res_i in enumerate(result):
         api_full_name = safe_get_value(res_i, 0, "res_i")
         api_name, state = get_name_and_state(api_full_name)
-        if last_api_name:
-            if api_name == last_api_name:
-                if state == last_state:
-                    num += 1
-                else:
-                    input_num = num
-                    num, last_state = 1, state
-            else:
-                output_num = num
-                find_error_rows(result[start:start + input_num + output_num], start, input_num, highlight_dict,
-                                dump_mode)
-                num, last_api_name, last_state = 1, api_name, state
-                start += input_num + output_num
-                input_num, output_num = 1, 0
+        api_batches_update(api_batches, api_name, state, i)
+    with tqdm(total=len(api_batches), desc="API/Module Analyse Progress", unit="item", ncols=100) as progress_bar:
+        for api_batch in api_batches:
+            find_error_rows(result[api_batch.start: api_batch.params_grad_end_index], api_batch, highlight_dict, 
+                            dump_mode)
+            progress_bar.update(1)
+
+
+def value_check(value, api_name=None, i=None, result_df_columns=None):
+    if not table_value_is_valid(value):
+        if result_df_columns:
+            logger.error(f"Malicious value [{value}] at api_name [{api_name}], column [{result_df_columns[i]}], "
+                         f"is not allowed to be written into the compare result xlsx.")
         else:
-            num, last_api_name, last_state = 1, api_name, state
-        progress_bar.update(1)
-    progress_bar.close()
-    if state:
-        if state == Const.INPUT:
-            input_num = num
-        else:
-            output_num = num
-        find_error_rows(result[start:start + input_num + output_num], start, input_num, highlight_dict,
-                        dump_mode)
+            logger.error(f"Malicious value [{value}] is not allowed to be written into the compare result xlsx.")
+
+
+def df_malicious_value_check(df_chunk, result_df_columns):
+    for row in df_chunk.itertuples(index=False):
+        api_name = row[0]
+        for i, value in enumerate(row):
+            value_check(value, api_name, i, result_df_columns)
+
+
+def handle_multi_process_malicious_value_check(func, result_df):
+    result_total_nums = len(result_df)
+    process_num = int((multiprocessing.cpu_count() + 1) / 2)
+
+    if result_total_nums <= process_num:
+        process_num = 1
+        chunks = [result_df]
+    else:
+        chunk_size = result_total_nums // process_num
+        chunks = [result_df.iloc[i: i + chunk_size] for i in range(0, result_total_nums, chunk_size)]
+
+    pool = multiprocessing.Pool(process_num)
+
+    def err_call(args):
+        logger.error("Multiprocessing malicious value check failed! Reason: {}".format(args))
+        try:
+            pool.terminate()
+        except OSError:
+            logger.error("Pool terminate failed")
+
+    result_df_columns = result_df.columns.tolist()
+    for column in result_df_columns:
+        value_check(column)
+    for df_chunk in chunks:
+        pool.apply_async(func, args=(df_chunk, result_df_columns,), error_callback=err_call)
+
+    pool.close()
+    pool.join()
+
+
+def compare_result_df_convert(value):
+    if not isinstance(value, (float, int)) or isinstance(value, bool):  # bool类型或者非数字类型转str
+        value = f"{str(value)}\t" if str(value) in ("inf", "-inf", "nan") else str(value)
+    if isinstance(value, float):
+        value = f"{str(value)}\t" if str(value) in ("inf", "-inf", "nan") else value
+    return value
 
 
 def highlight_rows_xlsx(result_df, highlight_dict, file_path):
     """Write and highlight results in Excel"""
 
-    update_highlight_err_msg(result_df, highlight_dict)     # add highlight err_msg
+    update_highlight_err_msg(result_df, highlight_dict)  # add highlight err_msg
 
     wb = openpyxl.Workbook()
     ws = wb.active
 
     # write header
     logger.info('Initializing Excel file.')
-    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)) or isinstance(value, bool):
-                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: "
-                                   f"{file_path}.")
-            ws.cell(row=i, column=j, value=f'{str(value)}\t' if str(value) in ('inf', '-inf', 'nan') else value)
-    
+
+    handle_multi_process_malicious_value_check(df_malicious_value_check, result_df)
+
+    result_df_convert = result_df.applymap(compare_result_df_convert)
+
+    for row in dataframe_to_rows(result_df_convert, index=False, header=True):
+        ws.append(row)
+
     # 对可疑数据标色
     logger.info('Coloring Excel in progress.')
     col_len = len(result_df.columns)
@@ -279,10 +377,11 @@ def highlight_rows_xlsx(result_df, highlight_dict, file_path):
     )
     for i in highlight_dict.get("red_rows", []):
         for j in range(1, col_len + 1):
-            ws.cell(row=i + 2, column=j).fill = red_fill
+            ws.cell(row=i + 2, column=j).fill = red_fill  # 2因为ws.cell中的row或column需要>=1,数据从第2行开始
     for i in highlight_dict.get("yellow_rows", []):
         for j in range(1, col_len + 1):
             ws.cell(row=i + 2, column=j).fill = yellow_fill
+
     logger.info('Saving Excel file to disk: %s' % file_path)
     save_workbook(wb, file_path)
 
@@ -314,15 +413,3 @@ def update_highlight_err_msg(result_df, highlight_dict):
                 red_lines_num_set.add(line_index)
 
     result_df[CompareConst.ERROR_MESSAGE] = err_msg
-
-
-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

msprobe/core/compare/layer_mapping/data_scope_parser.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -42,10 +42,10 @@ class DumpDataItem:
     construct_scope: str = ""
     scope_direction: Optional[str] = None
     scope_id: Optional[int] = None
+    state: str = ""
 
     # 类变量使用 ClassVar
-    framework2layername: ClassVar[Dict[str, str]] = {
-        Const.MS_FRAMEWORK: Const.CELL, Const.PT_FRAMEWORK: Const.MODULE}
+    layernames: ClassVar[set] = {Const.CELL, Const.MODULE}
     framework2stack_sign: ClassVar[Dict[str, Tuple[str, str]]] = {
         Const.MS_FRAMEWORK: ("Template", "construct"),
         Const.PT_FRAMEWORK: ("Template", r"in (for|back)ward,")
@@ -79,19 +79,30 @@ class DumpDataItem:
             )
             raise CompareException(CompareException.INVALID_DATA_ERROR)
 
+        if data_name_list[Const.LAST_INDEX] == Const.PARAMS_GRAD:
+            self.api_type = Const.PARAMS_GRAD
+            self.api_name = data_name_list[Const.PARAMS_GRAD_NAME_INDEX]
+            self.type_name = data_name_list[Const.PARAMS_GRAD_TYPE_NAME_INDEX]
+            self.state = Const.PARAMS_GRAD
+            return
+
         self.api_type = data_name_list[Const.API_TYPE_INDEX]
         self.type_name = data_name_list[Const.TYPE_NAME_INDEX]
-        if self.api_type == self.framework2layername.get(self.framework):
+        if self.api_type in self.layernames:
             self.api_name = data_name_list[Const.LAYER_NAME_INDEX]
+            self.state = data_name_list[Const.SCOPE_DIRECTION_INDEX]
         else:
             self.api_name = self.type_name
+            self.state = data_name_list[Const.LAST_INDEX]
 
     def set_layer_scope(self, construct_info: str) -> None:
         self.construct_scope = construct_info
-        if self.api_type == self.framework2layername.get(self.framework):
+        if self.api_type in self.layernames:
             # remove api name
             data_list = self.data_name.split(Const.SEP)
             data_list = data_list[:Const.LAYER_NAME_INDEX] + data_list[Const.TYPE_NAME_INDEX:]
+        elif self.api_type == Const.PARAMS_GRAD:
+            data_list = self.data_name.split(Const.SEP)
         elif construct_info:
             data_list = construct_info.split(Const.SEP)
         else:
@@ -100,7 +111,7 @@ class DumpDataItem:
         if data_list:
             self.layer_scope = Const.SEP.join(data_list[:Const.TYPE_NAME_INDEX])
         else:
-            self.layer_scope = self.framework2layername.get(self.framework)
+            self.layer_scope = Const.TOP_LAYER
         if construct_info:
             construct_list = construct_info.split(Const.SEP)
             if len(construct_list) < abs(Const.LAYER_NAME_INDEX):
@@ -115,7 +126,7 @@ class DumpDataItem:
 
     def set_stack_scope(self, stack_info: str) -> None:
         # Cell/Module has no stack info
-        if self.api_type == self.framework2layername.get(self.framework):
+        if self.api_type in self.layernames:
             return
 
         if self.api_type in Const.DATA_TYPE_SKIP_LIST or not stack_info:

msprobe/core/compare/layer_mapping/layer_mapping.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import os
+from collections import defaultdict
 
 from msprobe.core.common.const import CompareConst, Const
 from msprobe.core.common.file_utils import load_json, load_yaml, save_yaml
@@ -21,18 +22,20 @@ from msprobe.core.common.utils import (add_time_with_yaml,
                                        detect_framework_by_dump_json,
                                        get_stack_construct_by_dump_json_path)
 from msprobe.core.compare.layer_mapping.data_scope_parser import get_dump_data_items
-from msprobe.core.compare.utils import read_op
+from msprobe.core.compare.utils import read_op, reorder_op_name_list
+
 
 
 class LayerTrie:
     def __init__(self, type_name, framework=None):
         self.type_name = type_name
-        self.data_items = []
+        self.data_items = defaultdict(list)
         self.children = {}
         self.framework = framework
 
     def __repr__(self):
-        return f"Layer(type_name={self.type_name}, data_number={len(self.data_items)})"
+        data_nums = [{k: len(v)} for k, v in self.data_items.items()]
+        return f"Layer(type_name={self.type_name}, data_number={data_nums})"
 
     def get(self, name):
         return self.children.get(name)
@@ -46,10 +49,10 @@ class LayerTrie:
             if name not in node.children:
                 node.children[name] = LayerTrie(name, data_item.framework)
             node = node.children[name]
-        node.data_items.append(data_item)
+        node.data_items[data_item.state].append(data_item)
         node.type_name = data_item.type_name
 
-    def query_data(self, scope, index, default_value=None):
+    def query_data(self, scope, state, index, default_value=None):
         parts = scope.split(Const.SEP)
         node = self
         scope_name_list = parts[1:]
@@ -58,9 +61,9 @@ class LayerTrie:
             if name not in node.children:
                 return default_value
             node = node.children[name]
-        if index >= len(node.data_items):
+        if index >= len(node.data_items[state]):
             return default_value
-        return node.data_items[index]
+        return node.data_items[state][index]
 
     def save_to_yaml(self, output_path):
         result = {f"{self.type_name} @ {self}": self.convert_to_dict(self)}
@@ -70,7 +73,7 @@ class LayerTrie:
 
     def convert_to_dict(self, node):
         result = {}
-        result["data_item"] = [node.data_name for node in node.data_items]
+        result["data_item"] = {st: [dt.data_name for dt in dts] for st, dts in node.data_items.items()}
         for child_key, child_node in node.children.items():
             key = f"{child_key} @ {child_node}"
             result[key] = self.convert_to_dict(child_node)
@@ -102,10 +105,11 @@ def convert_scope(layer_trie, data_item, mapping=None):
             cur_node = child_node
             idx += 1
     index = -1
-    for idx, child in enumerate(cur_node.data_items):
+    state = data_item.state
+    for idx, child in enumerate(cur_node.data_items[state]):
         if data_item.data_name == child.data_name:
             index = idx
-    return new_scope, index
+    return new_scope, state, index
 
 
 def get_data_items_and_tree(dump_json_path, output_path):
@@ -122,8 +126,8 @@ def get_data_items_and_tree(dump_json_path, output_path):
 
 
 def convert_data_item(npu_tree, bench_tree, npu_data_item, mapping):
-    new_scope, index = convert_scope(npu_tree, npu_data_item, mapping)
-    bench_data_item = bench_tree.query_data(new_scope, index)
+    new_scope, state, index = convert_scope(npu_tree, npu_data_item, mapping)
+    bench_data_item = bench_tree.query_data(new_scope, state, index)
     return bench_data_item
 
 
@@ -223,7 +227,10 @@ def generate_data_mapping(npu_json_path, bench_json_path, api_mapping, output_pa
             continue
         npu_full_op_names = read_full_op_names(npu_data, npu_op_name)
         bench_full_op_names = read_full_op_names(bench_data, bench_op_name)
-        mapping = generate_op_data_mapping(npu_op_name, npu_full_op_names, bench_op_name, bench_full_op_names)
+        npu_full_op_names_reorder = reorder_op_name_list(npu_full_op_names)
+        bench_full_op_names_reorder = reorder_op_name_list(bench_full_op_names)
+        mapping = generate_op_data_mapping(npu_op_name, npu_full_op_names_reorder,
+                                           bench_op_name, bench_full_op_names_reorder)
         data_mapping.update(mapping)
     if output_path:
         file_name = add_time_with_yaml("data_mapping")

msprobe/core/compare/layer_mapping/postprocess_pass.py

@@ -29,9 +29,10 @@ def backward_pass(data_items, name2item):
         data_name_list = data_item.data_name.split(Const.SEP)
         if not data_name_list:
             continue
-        if Const.BACKWARD in data_name_list[Const.SCOPE_DIRECTION_INDEX :]:
-            data_name_list[Const.SCOPE_DIRECTION_INDEX :] = [
-                s.replace(Const.BACKWARD, Const.FORWARD) for s in data_name_list[Const.SCOPE_DIRECTION_INDEX :]
+        if Const.BACKWARD in data_name_list[Const.SCOPE_DIRECTION_INDEX:]:
+            data_name_list[Const.SCOPE_DIRECTION_INDEX:] = [
+                s.replace(Const.BACKWARD, Const.FORWARD)
+                for s in data_name_list[Const.SCOPE_DIRECTION_INDEX:]
             ]
             forward_name = Const.SEP.join(data_name_list)
             forward_item = name2item.get(forward_name, None)