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

msprobe/mindspore/code_mapping/bind.py

@@ -0,0 +1,264 @@
+# 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 os
+import time
+import glob
+from typing import Dict, List
+from pathlib import Path
+
+import pandas as pd
+
+from msprobe.core.common.const import Const
+from msprobe.core.common.file_utils import (
+    check_file_or_directory_path,
+    FileOpen,
+    create_directory,
+    write_csv,
+    check_path_before_create,
+    read_csv,
+    write_df_to_csv
+)
+from msprobe.mindspore.code_mapping.graph import GraphNode
+from msprobe.mindspore.common.log import logger
+
+
+# 定义Trie节点
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end_of_key = False
+        self.value = None
+
+
+# 定义Trie树
+class Trie:
+    def __init__(self):
+        self.root = TrieNode()
+
+    # 向Trie中插入一个键
+    def insert(self, key, value):
+        node = self.root
+        for key_char in key:
+            if key_char not in node.children:
+                node.children[key_char] = TrieNode()
+            node = node.children[key_char]
+        # 标记结束位置
+        node.is_end_of_key = True
+        node.value = value
+
+    # 在name字符串中查找所有匹配的键
+    def search_in_string(self, string):
+        matched_values = []
+        for i in range(len(string)):
+            node = self.root
+            j = i
+            # 从字符串的每个字符开始，逐字符查找匹配
+            while j < len(string) and string[j] in node.children:
+                node = node.children[string[j]]
+                if node.is_end_of_key:
+                    matched_values.append(node.value)
+                j += 1
+        return matched_values
+
+
+# 定义匹配函数
+def match_codes(trie, name):
+    matched_nodes = trie.search_in_string(name)
+    matched_codes = [Const.NEW_LINE.join(node.code_info) for node in matched_nodes]
+    return Const.NEW_LINE.join(matched_codes)
+
+
+def match_names(trie, name):
+    matched_nodes = trie.search_in_string(name)
+    matched_names = [node.scope for node in matched_nodes]
+    return Const.NEW_LINE.join(matched_names)
+
+
+def map_op_names_to_codes_and_scopes(df, match_dict):
+    # 构建Trie树并插入所有键
+    trie = Trie()
+    for key, value in match_dict.items():
+        trie.insert(key, value)
+
+    df[Const.CODE_STACK] = df[Const.OP_NAME].apply(lambda name: match_codes(trie, name))
+    df[Const.SCOPE_NAME] = df[Const.OP_NAME].apply(lambda name: match_names(trie, name))
+    return df
+
+
+def find_npy_files(npy_path):
+    """
+    查找指定路径下所有的.npy文件。
+
+    Parameters:
+        npy_path (str): 搜索的路径，可以是文件或目录。
+
+    Returns:
+        List[Path]: 找到的.npy文件路径列表。
+    """
+    npy_files = []
+    npy_path_obj = Path(npy_path)
+
+    # 检查当前路径是否是一个以 .npy 结尾的文件
+    if npy_path_obj.suffix == Const.NUMPY_SUFFIX and npy_path_obj.is_file():
+        check_file_or_directory_path(npy_path_obj)
+        npy_files.append(npy_path_obj.resolve())
+        return npy_files
+
+    # 如果是目录，使用Path.rglob查找所有.npy文件
+    if npy_path_obj.is_dir():
+        for file in npy_path_obj.rglob(Const.NUMPY_PATTERN):
+            check_file_or_directory_path(file)
+            npy_files.append(file.resolve())
+    else:
+        logger.info(f"The specified path is neither an .npy file nor a directory: {npy_path}")
+
+    return npy_files
+
+
+def write_to_csv(param: Dict, output_dir: str):
+    """
+    将参数写入CSV文件。
+
+    Parameters:
+        param (Dict): 要写入的数据，格式为{文件名: (代码堆栈, 作用域名称)}。
+        output_dir (str): 输出目录路径。
+    """
+    create_directory(output_dir)
+
+    # 使用时间戳生成文件名
+    timestamp = time.strftime("%Y%m%d%H%M%S", time.localtime())
+    file_path = Path(output_dir) / f"code_mapping_{timestamp}.csv"
+    check_path_before_create(file_path)
+    data = [(name, res1, res2) for name, (res1, res2) in param.items()]
+    df = pd.DataFrame(data, columns=[Const.FILE_PATH, Const.CODE_STACK, Const.SCOPE_NAME])
+    write_df_to_csv(df, file_path)
+
+
+def find_statistic_files(path):
+    if not os.path.isdir(path):
+        if os.path.basename(path) == 'statistic.csv':
+            return [path]
+        else:
+            return []
+    pattern = os.path.join(path, '**', "statistic.csv")
+
+    statistic_files = list(glob.glob(pattern, recursive=True))
+    return statistic_files
+
+
+def check_and_fix_header(file_path: str):
+    """
+    检查 CSV 文件的表头是否以逗号结尾，如果没有则添加一个逗号。
+
+    Parameters:
+        file_path (str): CSV 文件的路径。
+
+    Returns:
+        bool: 如果表头被修改，返回 True；否则，返回 False。
+    """
+
+    with FileOpen(file_path, "r") as f:
+        lines = f.readlines()
+
+    if not lines:
+        logger.warning(f"The file {file_path} is empty.")
+        return False
+
+    # 获取表头并去除末尾的换行符
+    header = lines[0].rstrip(Const.NEW_LINE).rstrip('\r')
+
+    if not header.endswith(','):
+        logger.info(f"The header does not end with a comma. Adding a comma to the file: {file_path}.")
+        # 添加逗号并恢复换行符
+        lines[0] = header + Const.CSV_NEWLINE_SEPARATOR
+
+        # 写回修复后的内容到文件
+        with FileOpen(file_path, "w") as f:
+            f.writelines(lines)
+        logger.info(f"Added a trailing comma to the file: {file_path}.")
+        return True
+    else:
+        logger.info(f"The header already ends with a comma. No modification needed for the file: {file_path}.")
+        return False
+
+
+def bind_for_statistic(statistic_files: List[str], match_dict: Dict):
+    """
+    处理统计文件，绑定代码信息。
+
+    Parameters:
+        statistic_files (List[str]): 统计文件路径列表。
+        match_dict (Dict): 匹配字典，用于复杂映射。
+    """
+    for statistic_file in statistic_files:
+        # 使用FileOpen安全打开文件
+        header_modified = check_and_fix_header(statistic_file)
+        if header_modified:
+            logger.info(f"The header of the file {statistic_file} has been fixed.")
+
+        df = read_csv(statistic_file, as_pd=True)
+
+        # 进行复杂映射
+        df = map_op_names_to_codes_and_scopes(df, match_dict)
+
+        # 使用write_csv安全写入文件
+        write_df_to_csv(df, statistic_file)
+
+
+def bind_code_info_for_data(input_dir: str, nodes: Dict[str, GraphNode]) -> Dict[str, str]:
+    # 待重构后优化性能
+    match_dict = {}
+    for node in nodes.values():
+        # 屏蔽子图节点
+        if node.is_subgraph:
+            continue
+        # 获取规范化后的scope name
+        scope_name = node.scope.replace(Const.SCOPE_SEPARATOR, Const.REPLACEMENT_CHARACTER)
+        match_dict[scope_name] = node
+    npy_files = find_npy_files(input_dir)
+
+    bind_result = {}
+    if not npy_files:
+        statistic_files = find_statistic_files(input_dir)
+        if statistic_files:
+            bind_for_statistic(statistic_files, match_dict)
+        return bind_result
+
+    for npy_file in npy_files:
+        directory, file_name = os.path.split(npy_file)  # 拆分路径
+        name_without_ext = os.path.splitext(file_name)[0]  # 提取文件名（去掉扩展名）
+        if name_without_ext.isdigit():
+            # 3. 读取find.csv文件
+            csv_file_path = os.path.join(directory, 'mapping.csv')
+            check_file_or_directory_path(csv_file_path)
+            df = read_csv(csv_file_path, header=None)
+
+            # 4. 查找是否有与xxx.npy匹配的条目
+            matching_row = df[df[0] == file_name]  # 假设A列存储文件名
+            if not matching_row.empty:
+                corresponding_name = matching_row[1].values[0]
+            else:
+                corresponding_name = None
+            name_without_ext = os.path.splitext(corresponding_name)[0]
+        npy_path = os.path.realpath(npy_file)
+        node_scope = name_without_ext.split(".")[1]
+        trie = Trie()
+        for key, value in match_dict.items():
+            trie.insert(key, value)
+        bind_code = match_codes(trie, node_scope)
+        bind_name = match_names(trie, node_scope)
+        bind_result[npy_path] = (bind_code, bind_name)
+    return bind_result

msprobe/mindspore/code_mapping/cmd_parser.py

@@ -0,0 +1,40 @@
+# 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 os
+
+from msprobe.core.common.file_utils import check_file_or_directory_path, create_directory, FileChecker
+
+
+def add_ir_parser_arguments(parser):
+    parser.add_argument('--ir', type=str, required=True, help="Path to the graph file")
+    parser.add_argument('--dump_data', type=str, required=True, default=None, help="Path to data dir")
+    parser.add_argument('--output', type=str, required=False, default="./", help="Path to output dir")
+
+
+def check_args(args):
+    args.ir = os.path.abspath(args.ir)
+
+    check_file_or_directory_path(args.ir)
+
+    args.dump_data = os.path.abspath(args.dump_data)
+    if os.path.isdir(args.dump_data):
+        check_file_or_directory_path(args.dump_data, isdir=True)
+    else:
+        check_file_or_directory_path(args.dump_data, isdir=False)
+
+    args.output = os.path.abspath(args.output)
+    create_directory(args.output)
+

msprobe/mindspore/code_mapping/graph.py

@@ -0,0 +1,49 @@
+# 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 typing import List, Dict, Union
+
+
+class GraphNode:
+    def __init__(self, name: str, pos: int = -1, unique_name: str = "", operator_name: str = "",
+                 return_variable: str = "", return_value: str = "",
+                 var_inputs: List[str] = None, has_constant_input: bool = False,
+                 unique_id: str = "", scope: str = "", code_info: List[str] = None,
+                 is_subgraph: bool = False, attrs: Union[Dict[str, str], List[str]] = None):
+        self.name = name
+        self.unique_name = unique_name
+        self.pos = pos
+        self.operator_name = operator_name
+        self.return_variable = return_variable
+        self.return_value = return_value
+        self.var_inputs = var_inputs if var_inputs else []
+        self.has_constant_input = has_constant_input
+        self.unique_id = unique_id
+        self.scope = scope
+        self.code_info = code_info if code_info else []
+        self.attrs = attrs if attrs else ({} if not is_subgraph else [])
+        self.nodes = {}  # Internal nodes if this is a subgraph
+        self.predecessors = []  # Predecessor nodes
+        self.successors = []    # Successor nodes
+        self.is_subgraph = is_subgraph
+
+    def trace_back_ancestors(self, ancestors: List[str], visited: Dict[str, bool], parser) -> None:
+        if visited[self.unique_name]:
+            return
+        visited[self.unique_name] = True
+        ancestors.append(self.unique_name)
+        for predecessor in self.predecessors:
+            predecessor.trace_back_ancestors(ancestors, visited, parser)
+

msprobe/mindspore/code_mapping/graph_parser.py

@@ -0,0 +1,226 @@
+# 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 re
+import logging
+from typing import Tuple, List, Dict
+from msprobe.mindspore.code_mapping.graph import GraphNode
+
+
+class Parser:
+    def __init__(self):
+        self.nodes = {}
+        self.local_dict = {}
+        self.number_dict = {}
+
+    @staticmethod
+    def parse_subgraph_attributes(text: str, subgraph_node: GraphNode, start_pos: int, end_pos: int) -> None:
+        subgraph_attr_pattern = re.compile(r'subgraph attr:\s*(.*)', re.DOTALL)
+        match = subgraph_attr_pattern.search(text, start_pos, end_pos)
+        if match:
+            attrs = match.group(1).strip().split('\n')
+            if isinstance(subgraph_node.attrs, list):
+                subgraph_node.attrs.extend(attrs)
+
+    @staticmethod
+    def parse_graph_attributes(text: str, graph_node: GraphNode) -> None:
+        attr_pattern = re.compile(r'# Attrs:\s*(.*)', re.DOTALL)
+        match = attr_pattern.search(text, graph_node.pos)
+        if match:
+            attrs = match.group(1).strip().split('\n')
+            for attr in attrs:
+                if not attr:
+                    break
+                key, value = attr.split(':')
+                if isinstance(graph_node.attrs, dict):
+                    graph_node.attrs[key.strip()] = value.strip()
+
+    @staticmethod
+    def parse_code_info(text: str, start_pos: int, end_pos: int) -> List[str]:
+        code_info = []
+        code_info_pattern = re.compile(r'# .*', re.MULTILINE)
+        final_pos = end_pos if end_pos else len(text) - 1
+        lines = text[start_pos + 1:final_pos].split('\n')
+        for line in lines:
+            match = code_info_pattern.search(line)
+            if not match:
+                break
+            code_info.append(match.group(0).strip('# ').strip('/'))
+        return code_info
+
+    @staticmethod
+    def extract_bracket_content(text: str, start_pos: int) -> Tuple[str, int]:
+        stack = []
+        content = []
+        for i in range(start_pos, len(text)):
+            char = text[i]
+            if char == '(':
+                stack.append('(')
+            elif char == ')':
+                stack.pop()
+                if not stack:
+                    content.append(char)
+                    return ''.join(content), i
+            content.append(char)
+        raise ValueError("Mismatched parentheses")
+
+    @staticmethod
+    def find_matching_brace(text: str, start_pos: int) -> int:
+        stack = []
+        for i in range(start_pos, len(text)):
+            if text[i] == '{':
+                stack.append('{')
+            elif text[i] == '}':
+                stack.pop()
+                if not stack:
+                    return i
+        raise ValueError("Matching closing brace not found")
+
+    @staticmethod
+    def extract_constants(inputs_str: str) -> List[str]:
+        constant_pattern = re.compile(r'\b(\w+\(.*?\))')
+        constants = constant_pattern.findall(inputs_str)
+        return constants
+
+    def parse_func_graph(self, text: str) -> None:
+        func_graph_pattern = re.compile(r'# IR entry: @(\S+)')
+        matches = func_graph_pattern.finditer(text)
+        for match in matches:
+            func_name = match.group(1)
+            func_graph_info = GraphNode(name=func_name, pos=match.start(), is_subgraph=False)
+            self.nodes[func_name] = func_graph_info
+
+    def parse_nodes(self, text: str, subgraph_info: GraphNode) -> None:
+        node_pattern = re.compile(r'(%\d+)\((\S+)\)\s*=\s*(\S+)\(')
+        matches = list(node_pattern.finditer(text))
+        for i, match in enumerate(matches):
+            series_number = match.group(1)
+            variable_name = match.group(2)
+            operator_name = match.group(3)
+            unique_name = "&".join([series_number, variable_name])
+            self.local_dict[series_number] = unique_name
+
+            args_str, end_pos = self.__class__.extract_bracket_content(text, match.end() - 1)
+            inputs = re.findall(r'%\w+', args_str)
+            subgraph_inputs = re.findall(r'@\w+', args_str)
+            inputs += subgraph_inputs
+
+            constants = self.__class__.extract_constants(args_str)
+
+            scope_pattern = re.compile(r'# .*scope.*:\s*\((.*?)\)', re.IGNORECASE | re.MULTILINE)
+
+            scope_match = scope_pattern.search(text, end_pos)
+            scope = scope_match.group(1) if scope_match else ""
+
+            id_pattern = re.compile(r'.*cnode_primal_attrs:'
+                                    r'\s*\{.*\b(?:forward_unique_id|unique_id):\s*\"(\d+)\".*', re.IGNORECASE)
+            unique_id_match = id_pattern.search(text, end_pos, scope_match.start())
+            unique_id = unique_id_match.group(1) if unique_id_match else None
+
+            if scope:
+                next_match = matches[i + 1].start() - 1 if i < len(matches) - 1 else None
+                code_info = self.__class__.parse_code_info(text, scope_match.end(), next_match)
+            else:
+                code_info = None
+
+            node_info = GraphNode(name=variable_name, unique_name=unique_name, operator_name=operator_name,
+                                  var_inputs=inputs + constants, unique_id=unique_id, scope=scope, code_info=code_info)
+
+            if unique_id and scope and not scope.startswith("Gradients"):
+                self.number_dict[unique_id] = node_info
+            
+            if subgraph_info:
+                subgraph_info.nodes[variable_name] = node_info
+
+            if not self.nodes.get(unique_name, None):
+                self.nodes[unique_name] = node_info
+            else:
+                pass
+
+            for const in constants:
+                if const not in self.nodes:
+                    const_node = GraphNode(name=const, operator_name="Constant", var_inputs=[], has_constant_input=True)
+                    if not self.nodes.get(const_node, None):
+                        self.nodes[const] = const_node
+                    if subgraph_info:
+                        subgraph_info.nodes[const] = const_node
+                    self.local_dict[const] = const
+
+            for input_var in node_info.var_inputs:
+                if input_var in self.local_dict or input_var in self.nodes:
+                    input_name = self.local_dict.get(input_var, input_var)
+                    input_node = self.nodes.get(input_name, None)
+                    if input_node:
+                        node_info.predecessors.append(input_node)
+                        input_node.successors.append(node_info)
+                else:
+                    param_node = GraphNode(name=input_var, operator_name="Param", var_inputs=[],
+                                           has_constant_input=False)
+                    if not self.nodes.get(input_var, None):
+                        self.nodes[input_var] = param_node
+                    node_info.predecessors.append(param_node)
+                    param_node.successors.append(node_info)
+
+    def extract_callees(self, text: str) -> None:
+        for node_info in self.nodes.values():
+            func_start_pos = node_info.pos
+            func_end_pos = text.find('}', func_start_pos)
+            func_text = text[func_start_pos:func_end_pos]
+            callee_pattern = re.compile(r'Partial\(@(\S+)\(')
+            callee_matches = callee_pattern.finditer(func_text)
+            for callee_match in callee_matches:
+                callee_name = callee_match.group(1)
+                if callee_name not in node_info.var_inputs:
+                    node_info.var_inputs.append(callee_name)
+
+    def parse_subgraphs(self, text: str) -> None:
+        subgraph_pattern = re.compile(r'subgraph\s+@(\S+)(\([^\)]*\))?\s+.*\{')
+        matches = list(subgraph_pattern.finditer(text))
+        end_pos = 0
+        for match in matches:
+            last_pos = end_pos + 2
+            subgraph_name = match.group(1).split('(')[0]
+            start_pos = match.start()
+            end_pos = self.__class__.find_matching_brace(text, start_pos)
+            subgraph_text = text[start_pos:end_pos + 1]
+            attr_text = text[last_pos:start_pos]
+            subgraph_info = GraphNode(name=subgraph_name, pos=start_pos, is_subgraph=True)
+            self.nodes[subgraph_name] = subgraph_info
+            self.__class__.parse_subgraph_attributes(text, subgraph_info, last_pos, start_pos)
+            self.parse_nodes(subgraph_text, subgraph_info)
+            subgraph_info.end = end_pos
+            logging.info('Parsed subgraph: %s', subgraph_name)
+
+    def count_nodes(self) -> Tuple[int, int]:
+        total_nodes = len(self.nodes)
+        total_cnodes = sum(1 for node in self.nodes.values() if node.name.startswith('CNode'))
+        return total_nodes, total_cnodes
+
+    def create_backward_map(self):
+        for node in self.nodes.values():
+            if node.scope and node.scope.startswith("Gradients"):
+                related_forward_node = self.number_dict.get(node.unique_id, None)
+                if related_forward_node:
+                    node.code_info = related_forward_node.code_info
+
+    def parse(self, text: str) -> None:
+        self.parse_func_graph(text)
+        self.parse_subgraphs(text)
+        self.parse_nodes(text, None)
+        self.extract_callees(text)
+        self.create_backward_map()
+
+    def get_nodes(self) -> Dict[str, GraphNode]:
+        return self.nodes

msprobe/mindspore/code_mapping/main.py

@@ -0,0 +1,24 @@
+# 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.mindspore.code_mapping.processor import process
+from msprobe.mindspore.code_mapping.cmd_parser import check_args
+
+
+def code_mapping_main(args):
+    check_args(args)
+    process(args)
+
+

msprobe/mindspore/code_mapping/processor.py

@@ -0,0 +1,34 @@
+# 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.mindspore.code_mapping.graph_parser import Parser
+from msprobe.mindspore.code_mapping.bind import bind_code_info_for_data, write_to_csv
+from msprobe.core.common.file_utils import FileOpen
+
+
+def process(args):
+    ir_file_path = args.ir
+    with FileOpen(ir_file_path, 'r') as f:
+        input_text = f.read()
+
+    parser = Parser()
+    parser.parse(input_text)
+
+    nodes = parser.get_nodes()
+
+    bind_result = bind_code_info_for_data(args.dump_data, nodes)
+    if bind_result:
+        write_to_csv(bind_result, args.output)
+

msprobe/mindspore/common/const.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");
@@ -48,6 +48,8 @@ class Const:
     MINT_DATA_PREFIX = "Mint."
     MINT_NN_FUNC_DATA_PREFIX = "MintFunctional."
     DISTRIBUTED_DATA_PREFIX = "Distributed."
+    TORCH_DATA_PREFIX = "Torch."
+    TORCH_NPU_DATA_PREFIX = "NPU."
 
     SUPPORTED_API_LIST_FILE = "support_wrap_ops.yaml"
     SUPPORTED_TENSOR_LIST_KEY = "tensor"