mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/pytorch/api_accuracy_checker/generate_op_script/op_generator.py

@@ -1,480 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-# Copyright (c) 2024-2025, 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 argparse
-import json
-import os
-import re
-
-import math
-import numpy as np
-import torch
-
-from msprobe.pytorch.api_accuracy_checker.compare.compare_utils import binary_standard_api, absolute_standard_api, \
-ulp_standard_api, thousandth_standard_api
-from msprobe.core.common.file_utils import FileOpen, load_json, save_json
-from msprobe.core.common.utils import check_file_or_directory_path, check_op_str_pattern_valid, is_int
-from msprobe.core.common.const import Const, MonitorConst, MsgConst, FileCheckConst
-from msprobe.core.common.log import logger
-from msprobe.core.common.file_utils import make_dir, change_mode
-from msprobe.core.common.decorator import recursion_depth_decorator
-
-TENSOR_DATA_LIST = ["torch.Tensor", "torch.nn.parameter.Parameter"]
-TORCH_BOOL_TYPE = ["torch.bool"]
-TORCH_INT_TYPE = ["torch.uint8", "torch.int8", "torch.int16", "torch.short", "torch.int32", "torch.int",
-                  "torch.int64", "torch.long"]
-TORCH_FLOAT_TYPE = ["torch.float16", "torch.half", "torch.bfloat16", "torch.float32", "torch.float",
-                    "torch.float64", "torch.double"]
-TORCH_COMPLEX_TYPE = ["torch.complex32", "torch.chalf", "torch.complex64", "torch.cfloat", "torch.complex128",
-                      "torch.cdouble"]
-OPERATOR_TYPE = ("Functional", "Tensor", "Torch")
-
-API_INFO = 2
-FOUR_SEGMENT = 4
-FIVE_SEGMENT = 5
-DATA_NAME = "data_name"
-API_MAX_LENGTH = 300
-PROPAGATION_LIST = [Const.FORWARD, Const.BACKWARD]
-DATAMODE_LIST = ["random_data", "real_data"]
-ITER_MAX_TIMES = 1000
-
-
-class APIInfo:
-    def __init__(self, api_full_name, api_info_dict, backward_info=None):
-        self.api_full_name = api_full_name
-        self.api_info_dict = api_info_dict
-        self.backward_info = backward_info
-
-    @property
-    def api_type(self):
-        return self.api_full_name.split(Const.SEP, -1)[0]
-
-    @classmethod
-    def from_json(cls, json_content, propagation):
-        forward_name, forward_dict = list(json_content.items())[0]
-        forward_info = cls(api_full_name=forward_name, api_info_dict=forward_dict)
-
-        if propagation == Const.BACKWARD:
-            backward_name, backward_dict = list(json_content.items())[1]
-            backward_info = cls(api_full_name=backward_name, api_info_dict=backward_dict)
-            forward_info.backward_info = backward_info
-
-        if not forward_info.is_supported_type():
-            raise ValueError(f"type {forward_info.api_type} of API is not supported!")
-
-        return forward_info
-
-    def is_supported_type(self):
-        return self.api_type in OPERATOR_TYPE
-
-
-class CommonConfig:
-    def __init__(self, json_config):
-        self.dump_json_path = json_config.get('dump_json_path')
-        self.api_name = json_config.get('api_name')
-        self.extract_api_path = json_config.get('extract_api_path')
-        self.propagation = json_config.get('propagation')
-        self.data_mode = json_config.get('data_mode')
-        self.random_seed = json_config.get('random_seed')
-        self.iter_times = json_config.get('iter_times')
-        self._check_config()
-
-
-    def check_user_settings(self):
-        iter_t = self.iter_times
-        if iter_t <= 0:
-            raise ValueError("iter_times should be an integer bigger than zero!")
-        if iter_t > ITER_MAX_TIMES:
-            raise ValueError("iter_times should not be greater than 1000!")
-
-        json_file = self.extract_api_path
-        propagation = self.propagation
-
-        json_content = load_json(json_file)
-
-        # ensure the dict is not empty
-        if not json_content:
-            raise ValueError(f'json file is empty!')
-
-        # ensure json_content is of type dict
-        if not isinstance(json_content, dict):
-            raise ValueError(f'content of json file is not a dict!')
-
-        # ensure the length of json_content is within allowed limits
-        if len(json_content) > API_INFO:
-            raise ValueError(f'json file has more than one API, the API only contains forward and backward info')
-
-        # Retrieve the first API name and dictionary
-        forward_item = next(iter(json_content.items()), None)
-        if not forward_item or not isinstance(forward_item[1], dict) or not forward_item[1]:
-            raise ValueError(f'Invalid forward API data in json_content!')
-
-        # if propagation is backward, ensure json file contains forward and backward info
-        if propagation == Const.BACKWARD and len(json_content) < API_INFO:
-            raise ValueError(f'Backward propagation requires contains forward and backward info!')
-
-        # if propagation is backward, ensure it has valid data
-        if propagation == Const.BACKWARD:
-            backward_item = list(json_content.items())[1]
-            if not isinstance(backward_item[1], dict) or not backward_item[1]:
-                raise ValueError(f'Invalid backward API data in json_content!')
-
-        return json_content
-
-
-    def _check_config(self):
-        if self.dump_json_path:
-            check_file_or_directory_path(self.dump_json_path)
-        if self.api_name:
-            check_op_str_pattern_valid(self.api_name)
-            if len(self.api_name) > API_MAX_LENGTH:
-                raise ValueError(f'API name {self.api_name} is too long!')
-        make_dir(os.path.dirname(self.extract_api_path))
-        if self.propagation and self.propagation not in PROPAGATION_LIST:
-            raise ValueError(f'propagation is invalid, it should be one of {PROPAGATION_LIST}')
-        if self.data_mode and self.data_mode not in DATAMODE_LIST:
-            raise ValueError(f'data_mode is invalid, it should be one of {DATAMODE_LIST}')
-        if not is_int(self.random_seed):
-            raise ValueError(f'random_seed is invalid, it should be an int')
-        if not is_int(self.iter_times):
-            raise ValueError(f'iter_times is invalid, it should be an int')
-
-
-class APIExtractor:
-    def __init__(self, api_name, dump_json_path, output_file):
-        self.api_name = api_name
-        self.dump_json_path = dump_json_path
-        self.output_file = output_file
-        self.data = None
-
-    def extract_op(self):
-        self.data = load_json(self.dump_json_path)
-        new_data = {}
-        extract_key_pattern = re.compile(f"^{re.escape(self.api_name)}\..+")
-        real_data_path = self.data.get('dump_data_dir', '')
-        for key, value in self.data.get('data', {}).items():
-            if extract_key_pattern.match(key):
-                if real_data_path:
-                    value = self.load_real_data_path(value, real_data_path)
-                new_data[key] = value
-        if not new_data:
-            logger.warning(f"Warning: The api '{self.api_name}' does not exist in the file.")
-        else:
-            save_json(self.output_file, new_data, indent=4)
-            logger.info(
-                f"The api '{self.api_name}' has been successfully extracted and saved in: {self.output_file}")
-
-    def load_real_data_path(self, value, dump_data_dir):
-        parameters = [Const.INPUT_ARGS, Const.GRAD_INPUT, Const.INPUT, Const.OUTPUT, Const.GRAD_OUTPUT]
-        for parameter in parameters:
-            for v in value.get(parameter, []):
-                if v is not None:
-                    self.update_data_name(v, dump_data_dir)
-        return value
-
-    @recursion_depth_decorator("OpGenerator: APIExtractor.update_data_name")
-    def update_data_name(self, data, dump_data_dir):
-        if isinstance(data, list):
-            for item in data:
-                self.update_data_name(item, dump_data_dir)
-        elif DATA_NAME in data:
-            data[DATA_NAME] = os.path.join(dump_data_dir, data[DATA_NAME])
-
-
-class OperatorScriptGenerator:
-    def __init__(self, common_config, args_info_forward, kwargs_info_forward, args_info_backward):
-        self.common_config = common_config
-        self.args_info_forward = args_info_forward
-        self.kwargs_info_forward = kwargs_info_forward
-        self.args_info_backward = args_info_backward
-
-    @staticmethod
-    def get_compare_standard(api_name):
-        api_standard_map = {
-            "binary_standard_api": "CompareStandard.BINARY_EQUALITY_STANDARD",
-            "absolute_standard_api": "CompareStandard.ABSOLUTE_THRESHOLD_STANDARD",
-            "ulp_standard_api": "CompareStandard.ULP_ERROR_STANDARD",
-            "thousandth_standard_api": "CompareStandard.THOUSANDTH_STANDARD"
-        }
-        for standard_api, standard_value in api_standard_map.items():
-            if api_name in globals()[standard_api]:
-                return standard_value
-        return "CompareStandard.BENCHMARK_STANDARD"
-
-    @staticmethod
-    def extract_detailed_api_segments(full_api_name):
-        """
-        Function Description:
-            Extract the name of the API.
-        Parameter:
-            full_api_name_with_direction_status: Full name of the API. Example: torch.matmul.0.forward.output.0
-        Return:
-            api_name: Name of api. Example: matmul, mul, etc.
-            full_api_name: Full name of api. Example: torch.matmul.0
-            direction_status: Direction status of api. Example: forward, backward, etc.
-        """
-        api_parts = full_api_name.split(Const.SEP)
-        api_parts_length = len(api_parts)
-        api_type, api_name, api_order = None, None, None
-        if api_parts_length == FOUR_SEGMENT:
-            api_type, api_name, api_order, _ = api_parts
-        elif api_parts_length == FIVE_SEGMENT:
-            api_type, prefix, api_name, api_order, _ = api_parts
-            api_name = Const.SEP.join([prefix, api_name])
-        return api_type, api_name, api_order
-
-    def get_settings(self, api_full_name):
-        '''
-        internal_settings contain all information needed for the operator program.
-        keys:
-            api_full_name: api_type.api_name.ordinal_number
-            api_type: type of API, one of torch.nn.functional, torch.Tensor or Torch
-            api_name: name of API
-            ordinal_number: how many times the same api has been called
-            direction_status: forward
-            random_seed: if mode is random_data, random seed is random_seed
-            iter_times: if mode is random_data, generate iter_times group of data; if mode is real_data,
-            iter_times does not matter
-            args_element_assignment: code for args assignment
-            args_list_generator_device: code for generate args list on device
-            args_list_generator_bench: code for generate args list on bench
-            kwargs_value_assignment: code for kwargs assignment
-            kwargs_dict_generator_device: code for generate kwargs dict on device
-            kwargs_dict_generator_bench: code for generate kwargs dict on bench
-        '''
-        # Generate an internal setting dictionary based on user settings
-        # including API name, type, comparison standard, random seed, number of iterations and other information
-        internal_settings = {}
-        internal_settings["propagation"] = self.common_config.propagation
-        internal_settings["api_full_name"] = api_full_name
-        api_type, api_name, ordinal_number = self.extract_detailed_api_segments(api_full_name)
-        if api_type == "Functional":
-            internal_settings["api_type"] = "torch.nn.functional"
-        elif api_type == "Tensor":
-            internal_settings["api_type"] = "torch.Tensor"
-        else:
-            internal_settings["api_type"] = "torch"
-        internal_settings["api_name"] = api_name
-        internal_settings["compare_standard"] = self.get_compare_standard(api_name)
-        internal_settings["ordinal_number"] = ordinal_number
-        internal_settings["direction_status"] = self.common_config.propagation
-        internal_settings["random_seed"] = self.common_config.random_seed
-        if self.common_config.data_mode == "real_data":
-            internal_settings["iter_times"] = 1
-        else:
-            internal_settings["iter_times"] = self.common_config.iter_times
-        internal_settings["args_element_assignment"] = \
-                            self.generate_args_element_assignment_code(self.args_info_forward)
-        internal_settings["args_list_generator_device"] = \
-                            self.generate_args_list(self.args_info_forward, flag_device=True)
-        internal_settings["args_list_generator_bench"] = \
-                            self.generate_args_list(self.args_info_forward, flag_device=False)
-        internal_settings["kwargs_value_assignment"] = \
-                            self.generate_kwargs_value_assignment_code(self.kwargs_info_forward)
-        internal_settings["kwargs_dict_generator_device"] = \
-                            self.generate_kwargs_dict(self.kwargs_info_forward, flag_device=True)
-        internal_settings["kwargs_dict_generator_bench"] = \
-                            self.generate_kwargs_dict(self.kwargs_info_forward, flag_device=False)
-        if self.common_config.propagation == Const.BACKWARD:
-            internal_settings["args_element_assignment_backward"] = self.generate_args_element_assignment_code(
-                self.args_info_backward)
-            internal_settings["args_list_generator_device_backward"] = \
-                            self.generate_args_list(self.args_info_backward, flag_device=True)
-            internal_settings["args_list_generator_bench_backward"] = \
-                            self.generate_args_list(self.args_info_backward, flag_device=False)
-        else:
-            internal_settings["args_element_assignment_backward"] = ''
-            internal_settings["args_list_generator_device_backward"] = ''
-            internal_settings["args_list_generator_bench_backward"] = ''
-
-        return internal_settings
-
-    @recursion_depth_decorator("OpGenerator: OperatorScriptGenerator.recursive_args_element_assignment")
-    def recursive_args_element_assignment(self, args_info, name_number):
-        args_element_assignment = ""
-        for index, arg in enumerate(args_info):
-            if isinstance(arg, (list, tuple)):
-                new_args_element_assignment = \
-                    self.recursive_args_element_assignment(arg, name_number + "_" + str(index))
-                args_element_assignment += new_args_element_assignment
-            else:
-                arg["parameter_name"] = "arg" + name_number + "_" + str(index)
-                args_element_assignment += "    " + "arg_info" + name_number + "_" + str(index) + " = " + \
-                    "{}".format(str(arg)) + MsgConst.SPECIAL_CHAR[0]
-                args_element_assignment += "    " + "arg" + name_number + "_" + str(index) + " = " + \
-                    "generate_data(arg_info" + name_number + "_" + str(index) + ")" + MsgConst.SPECIAL_CHAR[0]
-        return args_element_assignment
-
-
-    def generate_args_element_assignment_code(self, args_info):
-        args_element_assignment = self.recursive_args_element_assignment(args_info, "")
-        return args_element_assignment
-
-    @recursion_depth_decorator("OpGenerator: OperatorScriptGenerator.recursive_args_list")
-    def recursive_args_list(self, args_info, flag_device=False, flag_bench=False):
-        args_list_generator = ""
-        for _, arg in enumerate(args_info):
-            if isinstance(arg, (list, tuple)):
-                (left_bracket, right_bracket) = ("[", "]") if isinstance(arg, list) else ("(", ")")
-                args_list_generator += left_bracket
-                new_args_list_generator = self.recursive_args_list(arg, flag_device=flag_device, flag_bench=flag_bench)
-                args_list_generator += new_args_list_generator
-                args_list_generator += right_bracket
-            else:
-                args_list_generator += arg.get("parameter_name")
-                if arg.get("type") in TENSOR_DATA_LIST:
-                    if flag_device:
-                        args_list_generator += ".to(device)"
-                    if flag_bench:
-                        args_list_generator += '.to(torch.device("cpu"))'
-                        args_list_generator += ".to(RAISE_PRECISION.get(str(" + arg.get("parameter_name") + \
-                            ".dtype), " + arg.get("parameter_name") + ".dtype))"
-            args_list_generator += Const.COMMA
-        return args_list_generator
-
-    def generate_args_list(self, args_info, flag_device):
-        if flag_device:
-            args_list_generator = self.recursive_args_list(args_info, flag_device=True)
-        else:
-            args_list_generator = self.recursive_args_list(args_info, flag_bench=True)
-        return args_list_generator
-
-    @recursion_depth_decorator("OpGenerator: OperatorScriptGenerator.recursive_kwargs_value_assignment")
-    def recursive_kwargs_value_assignment(self, info, key_name, name_number):
-        kwargs_value_assignment = ""
-        if isinstance(info, dict):
-            if info.get("type") == "torch.device" or info.get("type") == "torch.dtype":
-                kwargs_value_assignment += "    " + "kwarg_" + key_name + name_number + " = " + info.get("value")
-            else:
-                kwargs_value_assignment += "    " + "kwarg_info_" + key_name + name_number + " = " + \
-                    "{}".format(str(info)) + MsgConst.SPECIAL_CHAR[0]
-                kwargs_value_assignment += "    " + "kwarg_" + key_name + name_number + " = " + \
-                    "generate_data(kwarg_info_" + key_name + name_number + ")" + MsgConst.SPECIAL_CHAR[0]
-            info["parameter_name"] = "kwarg_" + key_name + name_number
-        else:
-            for index, arg in enumerate(info):
-                new_kwargs_value_assignment = self.recursive_kwargs_value_assignment(arg, key_name, name_number + \
-                    "_" + str(index))
-                kwargs_value_assignment += new_kwargs_value_assignment
-        return kwargs_value_assignment
-
-    def generate_kwargs_value_assignment_code(self, kwargs_info):
-        kwargs_value_assignment = ""
-        for key, value in kwargs_info.items():
-            kwargs_value_assignment += self.recursive_kwargs_value_assignment(value, key, "")
-        return kwargs_value_assignment
-
-    @recursion_depth_decorator("OpGenerator: OperatorScriptGenerator.recursive_kwargs_dict")
-    def recursive_kwargs_dict(self, info, flag_device=False, flag_bench=False):
-        kwargs_dict_generator = ""
-        if isinstance(info, dict):
-            kwargs_dict_generator += info.get("parameter_name")
-            if info.get("type") in TENSOR_DATA_LIST:
-                if flag_device:
-                    kwargs_dict_generator += ".to(device)"
-                if flag_bench:
-                    kwargs_dict_generator += '.to(torch.device("cpu"))'
-                    kwargs_dict_generator += ".to(RAISE_PRECISION.get(str(" + info.get("parameter_name") + \
-                        ".dtype), " + info.get("parameter_name") + ".dtype))"
-        else:
-            (left_bracket, right_bracket) = ("[", "]") if isinstance(info, list) else ("(", ")")
-            kwargs_dict_generator += left_bracket
-            for arg in info:
-                kwargs_dict_generator += self.recursive_kwargs_dict(arg, flag_device=flag_device, flag_bench=flag_bench)
-                kwargs_dict_generator += Const.COMMA
-            kwargs_dict_generator += right_bracket
-        return kwargs_dict_generator
-
-
-    def generate_kwargs_dict(self, kwargs_info, flag_device):
-        kwargs_dict_generator = ""
-        for key, value in kwargs_info.items():
-            kwargs_dict_generator += '"' + key + '"' + MonitorConst.NAME_SEP
-            if flag_device:
-                kwargs_dict_generator += self.recursive_kwargs_dict(value, flag_device=True) + Const.COMMA
-            else:
-                kwargs_dict_generator += self.recursive_kwargs_dict(value, flag_bench=True) + Const.COMMA
-        return kwargs_dict_generator
-
-
-def _op_generator_parser(parser):
-    parser.add_argument("-i", "--config_input", dest="config_input", type=str,
-                        help="<Required> Path of config json file", required=True)
-    parser.add_argument("-o", "--api_output_path", dest="api_output_path", type=str,
-                        help="<Required> Path of extract api_name.json.", required=True)
-
-
-def parse_json_config(json_file_path):
-    if not json_file_path:
-        raise Exception("config_input path can not be empty, please check.")
-    json_config = load_json(json_file_path)
-    common_config = CommonConfig(json_config)
-    return common_config
-
-
-def _run_operator_generate_commond(cmd_args):
-    common_config = parse_json_config(cmd_args.config_input)
-
-    if common_config.dump_json_path:
-        api_extract = APIExtractor(common_config.api_name, common_config.dump_json_path, common_config.extract_api_path)
-        api_extract.extract_op()
-    check_file_or_directory_path(common_config.extract_api_path)
-    check_file_or_directory_path(cmd_args.api_output_path, isdir=True)
-    json_content = common_config.check_user_settings()
-    api_info = APIInfo.from_json(json_content, common_config.propagation)
-
-    if common_config.propagation == Const.BACKWARD:
-        # read and check json
-        api_full_name_forward, api_info_dict_forward = api_info.api_full_name, api_info.api_info_dict
-        api_full_name_backward, api_info_dict_backward = (api_info.backward_info.api_full_name,
-                                                          api_info.backward_info.api_info_dict)
-        args_info_forward = api_info_dict_forward.get(Const.INPUT_ARGS)
-        kwargs_info_forward = api_info_dict_forward.get(Const.INPUT_KWARGS)
-        if Const.GRAD_INPUT in api_info_dict_backward:
-            args_info_backward = api_info_dict_backward.get(Const.GRAD_INPUT)
-        elif Const.INPUT in api_info_dict_backward:
-            args_info_backward = api_info_dict_backward.get(Const.INPUT)
-        op_generate = OperatorScriptGenerator(common_config, args_info_forward, kwargs_info_forward, args_info_backward)
-        internal_settings = op_generate.get_settings(api_full_name_backward)
-    else:
-        # read and check json
-        api_full_name_forward, api_info_dict_forward = api_info.api_full_name, api_info.api_info_dict
-        args_info_forward = api_info_dict_forward.get(Const.INPUT_ARGS)
-        kwargs_info_forward = api_info_dict_forward.get(Const.INPUT_KWARGS)
-        op_generate = OperatorScriptGenerator(common_config, args_info_forward, kwargs_info_forward, None)
-        internal_settings = op_generate.get_settings(api_full_name_forward)
-
-    template_path = os.path.join(os.path.dirname(__file__), "operator_replication.template")
-    operator_script_path = os.path.join(cmd_args.api_output_path, 
-                                        "{0}.py".format(internal_settings.get("api_full_name")))
-
-    try:
-        with FileOpen(template_path, 'r') as ftemp, FileOpen(operator_script_path, 'w') as fout:
-            code_template = ftemp.read()
-            fout.write(code_template.format(**internal_settings))
-    except OSError:
-        logger.error(f"Failed to open file. Please check file {template_path} or {operator_script_path}.")
-    change_mode(operator_script_path, FileCheckConst.DATA_FILE_AUTHORITY)
-
-    logger.info(f"Generate operator script successfully and the name is {operator_script_path}.")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    _op_generator_parser(parser)
-    cmd_args = parser.parse_args()
-    _run_operator_generate_commond(cmd_args)