| # Copyright 2024 Arm Limited and/or its affiliates. |
| # |
| # This source code is licensed under the BSD-style license found in the |
| # LICENSE file in the root directory of this source tree. |
| |
| import json |
| import logging |
| import os |
| import re |
| import shutil |
| import subprocess |
| import tempfile |
| |
| from pathlib import Path |
| from typing import Dict, List, Optional, Tuple |
| |
| import numpy as np |
| import torch |
| |
| from executorch.backends.arm.test.common import arm_test_options, is_option_enabled |
| |
| from torch.export import ExportedProgram |
| from torch.fx.node import Node |
| |
| logger = logging.getLogger(__name__) |
| logger.setLevel(logging.WARNING) |
| |
| |
| class QuantizationParams: |
| __slots__ = ["node_name", "zp", "scale", "qmin", "qmax", "dtype"] |
| |
| # todo: zps and scales can be per tensors or per channel => a list?? |
| def __init__( |
| self, |
| node_name: str, |
| zp: int, |
| scale: float, |
| qmin: int, |
| qmax: int, |
| dtype: torch.dtype, |
| ): |
| self.node_name = node_name # not need I think, but good for error check |
| self.zp = zp |
| self.scale = scale |
| self.qmin = qmin |
| self.qmax = qmax |
| self.dtype = dtype |
| |
| |
| def _get_input_names(program: ExportedProgram) -> list[str]: |
| """ |
| Get a list[str] with the names of the inputs to this model. |
| |
| Args: |
| program (ExportedProgram): The program to get input names from. |
| Returns: |
| A list of strings with the names of the model input. |
| """ |
| input_names = [] |
| |
| # E.g. bias and weights are 'placeholders' as well. This is used to |
| # get only the use inputs. |
| usr_inputs = program.graph_signature.user_inputs |
| for node in program.graph.nodes: |
| if node.op == "placeholder" and node.name in usr_inputs: |
| input_names.append(node.name) |
| |
| return input_names |
| |
| |
| def _get_input_quantization_params( |
| program: ExportedProgram, |
| ) -> list[QuantizationParams]: |
| """ |
| Get input QuantizationParams in a program, maximum one per input to the program. |
| Args: |
| program (ExportedProgram): The program to get input quantization parameters from. |
| Returns: |
| list[QuantizationParams]: The found quantization parameters. |
| Raises: |
| RuntimeError if no quantization parameters are found. |
| """ |
| |
| quant_params = [] |
| input_names = _get_input_names(program) |
| num_inputs = len(input_names) |
| for node in program.graph.nodes: |
| if ( |
| node.target == torch.ops.quantized_decomposed.quantize_per_tensor.default |
| and node.args[0].name in input_names |
| ): |
| qp = QuantizationParams( |
| node_name=node.args[0].name, |
| scale=node.args[1], |
| zp=node.args[2], |
| qmin=node.args[3], |
| qmax=node.args[4], |
| dtype=node.args[5], |
| ) |
| quant_params.append(qp) |
| if ( |
| len(quant_params) == num_inputs |
| ): # break early if we have all the inputs quantized parameters |
| break |
| if len(quant_params) == 0: |
| raise RuntimeError("No Quantization parameters found in exported model.") |
| return quant_params |
| |
| |
| def _get_output_node(program: ExportedProgram) -> Node: |
| """ |
| Get output node to this model. |
| |
| Args: |
| program (ExportedProgram): The program to get output node from. |
| Returns: |
| The node that is the output of 'program'. |
| """ |
| |
| for node in program.graph.nodes: |
| if node.op == "output": |
| return node |
| raise RuntimeError("No output node found.") |
| |
| |
| def _get_output_quantization_params( |
| program: ExportedProgram, output_node: Node |
| ) -> QuantizationParams: |
| """ |
| Get output QuantizationParams from a program. |
| Args: |
| program (ExportedProgram): The program to get output quantization parameters from. |
| Returns: |
| QuantizationParams: The found quantization parameters. |
| Raises: |
| RuntimeError if no output quantization parameters are found. |
| """ |
| |
| quant_params = None |
| for node in program.graph.nodes: |
| if ( |
| node.target == torch.ops.quantized_decomposed.dequantize_per_tensor.default |
| and node == output_node.args[0][0] |
| ): |
| quant_params = QuantizationParams( |
| node_name=node.args[0].name, |
| scale=node.args[1], |
| zp=node.args[2], |
| qmin=node.args[3], |
| qmax=node.args[4], |
| dtype=node.args[5], |
| ) |
| break # break early, there's only one output node |
| if quant_params is None: |
| raise RuntimeError("No Quantization parameters not found in exported model.") |
| return quant_params |
| |
| |
| """ |
| A class to store parameters needed for running programs, either in tosa or .pte format. |
| """ |
| |
| |
| class RunnerUtil: |
| def __init__( |
| self, |
| intermediate_path: str, |
| tosa_ref_model_path: Optional[str] = None, |
| ): |
| self.intermediate_path = intermediate_path |
| self.tosa_ref_model_path = tosa_ref_model_path or "tosa_reference_model" |
| assert os.path.exists( |
| self.intermediate_path |
| ), f"TOSA artifact path don't exist! Path: {self.intermediate_path}" |
| |
| self.is_quantized: bool = False |
| self.input_names: list[str] = None |
| self.output_name: str = None |
| self.qp_input: list[QuantizationParams] = None |
| self.qp_output: QuantizationParams = None |
| self.timeout = 120 |
| self.target_board: str = None |
| |
| self._has_init_run = False |
| |
| def init_run( |
| self, |
| exported_program: ExportedProgram, |
| edge_program: ExportedProgram, |
| is_quantized: bool, |
| target_board: str, |
| ): |
| |
| if target_board not in ["corstone-300", "corstone-320"]: |
| raise RuntimeError(f"Unknown target board: {target_board}") |
| |
| self.input_names = _get_input_names(edge_program) |
| self.output_node = _get_output_node(exported_program) |
| self.output_name = self.output_node.name |
| self.is_quantized = is_quantized |
| self.target_board = target_board |
| |
| if is_quantized: |
| self.qp_input = _get_input_quantization_params(exported_program) |
| self.qp_output = _get_output_quantization_params( |
| exported_program, self.output_node |
| ) |
| else: |
| self.qp_input = [None] * len(self.input_names) |
| self.qp_output = None |
| |
| self._has_init_run = True |
| |
| def set_timeout(self, timeout: int): |
| self.timeout = timeout |
| |
| def run_corstone( |
| self, |
| inputs: Tuple[torch.Tensor], |
| ) -> list[torch.Tensor]: |
| |
| assert ( |
| self._has_init_run |
| ), "RunnerUtil needs to be initialized using init_run() before running Corstone300." |
| |
| pte_path = os.path.join(self.intermediate_path, "program.pte") |
| assert os.path.exists(pte_path), f"Pte path '{pte_path}' not found." |
| |
| for input_name, quant_param, data in zip( |
| self.input_names, self.qp_input, inputs |
| ): |
| save_bytes(self.intermediate_path, data, False, input_name, quant_param) |
| |
| out_path = os.path.join(self.intermediate_path, "out") |
| out_path_with_suffix = out_path + "-0.bin" |
| input_paths = [] |
| for name in self.input_names: |
| input_paths.append( |
| os.path.join(self.intermediate_path, f"{name}.bin"), |
| ) |
| elf_path = os.path.join( |
| "cmake-out", |
| f"arm_semihosting_executor_runner_{self.target_board}", |
| "arm_executor_runner", |
| ) |
| assert os.path.exists( |
| elf_path |
| ), f"Did not find build arm_executor_runner in path {elf_path}, run setup_testing.sh?" |
| |
| cmd_line = f"executor_runner -m {pte_path} -o {out_path}" |
| for input_path in input_paths: |
| cmd_line += f" -i {input_path}" |
| |
| ethos_u_extra_args = "" |
| if is_option_enabled(arm_test_options.fast_fvp): |
| ethos_u_extra_args = ethos_u_extra_args + "--fast" |
| |
| command_args = { |
| "corstone-300": [ |
| "FVP_Corstone_SSE-300_Ethos-U55", |
| "-C", |
| "ethosu.num_macs=128", |
| "-C", |
| "mps3_board.visualisation.disable-visualisation=1", |
| "-C", |
| "mps3_board.telnetterminal0.start_telnet=0", |
| "-C", |
| "mps3_board.uart0.out_file='-'", |
| "-C", |
| "cpu0.CFGITCMSZ=11", |
| "-C", |
| "cpu0.semihosting-enable=1", |
| "-C", |
| "cpu0.semihosting-stack_base=0", |
| "-C", |
| f"ethosu.extra_args='{ethos_u_extra_args}'", |
| "-C", |
| "cpu0.semihosting-heap_limit=0", |
| "-C", |
| f"cpu0.semihosting-cmd_line='{cmd_line}'", |
| "-a", |
| elf_path, |
| "--timelimit", |
| f"{self.timeout}", |
| ], |
| "corstone-320": [ |
| "FVP_Corstone_SSE-320", |
| "-C", |
| "mps4_board.subsystem.ethosu.num_macs=128", |
| "-C", |
| "mps4_board.visualisation.disable-visualisation=1", |
| "-C", |
| "vis_hdlcd.disable_visualisation=1", |
| "-C", |
| "mps4_board.telnetterminal0.start_telnet=0", |
| "-C", |
| "mps4_board.uart0.out_file='-'", |
| "-C", |
| "mps4_board.uart0.unbuffered_output=1", |
| "-C", |
| "mps4_board.uart0.shutdown_on_eot=1", |
| "-C", |
| "mps4_board.subsystem.cpu0.semihosting-enable=1", |
| "-C", |
| "mps4_board.subsystem.cpu0.semihosting-stack_base=0", |
| "-C", |
| "mps4_board.subsystem.cpu0.semihosting-heap_limit=0", |
| "-C", |
| f"mps4_board.subsystem.ethosu.extra_args='{ethos_u_extra_args}'", |
| "-C", |
| f"mps4_board.subsystem.cpu0.semihosting-cmd_line='{cmd_line}'", |
| "-a", |
| elf_path, |
| "--timelimit", |
| f"{self.timeout}", |
| ], |
| } |
| |
| result = _run_cmd(command_args[self.target_board], check=False) |
| if result.returncode != 0: |
| raise RuntimeError( |
| f"Failed to run {command_args[self.target_board]}\nError: {result.stderr.decode()}" |
| ) |
| result_stdout = result.stdout.decode() |
| |
| error_regex = r"(^[EF][: ].*$)|(^.*Hard fault.*$)|(^.*Assertion.*$)" |
| |
| # Check for errors in the output |
| # regex to check for error or fault messages in stdout from FVP |
| if re.compile(error_regex, re.MULTILINE).search(result_stdout): |
| raise RuntimeError( |
| f"Corstone simulation failed:\ncmd: {command_args[self.target_board]}\n, log: \n {result_stdout}\n{result.stderr.decode()}" |
| ) |
| |
| tosa_ref_output = np.fromfile(out_path_with_suffix, dtype=np.float32) |
| output_shape = self.output_node.args[0][0].meta["val"].shape |
| tosa_ref_output = torch.from_numpy(tosa_ref_output).reshape(output_shape) |
| return [tosa_ref_output] |
| |
| def run_tosa_ref_model( |
| self, |
| inputs: Tuple[torch.Tensor], |
| ) -> list[torch.Tensor]: |
| """ |
| Run TOSA reference model using the tosa_reference_model program. |
| |
| In order to do that we need: |
| 1. desc.json, which points to files needed by tosa_reference_model. |
| 2. output.tosa, which is the TOSA buffer that describes the model we're |
| trying to run. |
| |
| These two files are created by arm_backend.py as part of partition stage |
| |
| All these files are saved on disk in self.intermediate_path. |
| |
| Args: |
| inputs (Tuple[torch.Tensor]): The input data to run the TOSA |
| |
| Returns: |
| torch.Tensor: The output of the TOSA reference model, as a torch |
| tensor. |
| |
| Here's a sample desc.json file: |
| { |
| "tosa_file": "output.tosa", |
| "ifm_name": [ |
| "arg0_1" |
| ], |
| "ifm_file": [ |
| "arg0_1.npy" |
| ], |
| "ofm_name": [ |
| "quantized_decomposed_dequantize_per_tensor_default_1" |
| ], |
| "ofm_file": [ |
| "ref-quantized_decomposed_dequantize_per_tensor_default_1.npy" |
| ], |
| "expected_return_code": 0, |
| "expected_failure": false |
| } |
| |
| Todo: |
| * It would be nice to not rely on files on disk. Should be possible |
| as a next step. See: |
| https://review.mlplatform.org/plugins/gitiles/tosa/reference_model/#executable-usage |
| """ |
| |
| assert ( |
| self._has_init_run |
| ), "RunnerUtil needs to be initialized using init_run() before running tosa reference." |
| |
| all_desc_file_paths = [ |
| str(path) for path in Path(self.intermediate_path).glob("desc*.json") |
| ] |
| assert ( |
| all_desc_file_paths |
| ), f"No TOSA description file found in '{self.intermediate_path}'." |
| if len(all_desc_file_paths) != 1: |
| raise NotImplementedError( |
| "Graphs with more than one partition are currently not supported." |
| ) |
| |
| desc_file_path = all_desc_file_paths[0] |
| assert os.path.exists( |
| desc_file_path |
| ), f"desc_file_path: {desc_file_path} does not exist" |
| |
| # Save the input data to disk as a .npy file, since that's what the TOSA |
| # reference model expects. Name of the file must match the name in |
| # desc.json, which is the tensor name from the graph + .npy |
| for input_name, quant_param, data in zip( |
| self.input_names, self.qp_input, inputs, strict=True |
| ): |
| save_npy( |
| self.intermediate_path, data, self.is_quantized, input_name, quant_param |
| ) |
| |
| # Run the TOSA reference model via command line, this will produce a |
| # .npy file with the result (aka OFM). |
| assert ( |
| shutil.which(self.tosa_ref_model_path) is not None |
| ), f"tosa_reference_model tool not found, did you run examples/arm/setup.sh? Path: {self.tosa_ref_model_path}" |
| loglevel_map = { |
| logging.INFO: "INFO", |
| logging.CRITICAL: "LOW", |
| logging.ERROR: "LOW", |
| logging.WARNING: "MED", |
| logging.DEBUG: "HIGH", |
| logging.NOTSET: "MED", |
| } |
| clamped_logging_level = max(min(logger.level // 10 * 10, 50), 0) |
| cmd_ref_model = [ |
| self.tosa_ref_model_path, |
| "--test_desc", |
| desc_file_path, |
| "-l", |
| loglevel_map[clamped_logging_level], |
| ] |
| _run_cmd(cmd_ref_model) |
| |
| # Load desc.json, just to get the name of the output file above |
| with open(desc_file_path) as f: |
| desc_json = json.load(f) |
| |
| tosa_ref_outputs = [] |
| for ofm_file in desc_json["ofm_file"]: |
| ofm_file_npy = os.path.join(self.intermediate_path, ofm_file) |
| |
| # Load the output file (OFM) and return it as a numpy array |
| tosa_ref_output = np.load(ofm_file_npy) |
| |
| if self.is_quantized: |
| # Need to dequant back to FP32 for comparison with torch output |
| # Convert to int32 prior to dequantize the output |
| if tosa_ref_output.dtype == np.int8: |
| tosa_ref_output = tosa_ref_output.astype(np.int32) |
| quant_param = self.qp_output |
| assert ( |
| quant_param is not None |
| ), "There are no quantization parameters, check output parameters" |
| tosa_ref_output = (tosa_ref_output - quant_param.zp) * quant_param.scale |
| |
| if tosa_ref_output.dtype == np.double: |
| tosa_ref_output = tosa_ref_output.astype("float32") |
| |
| # tosa_output is a numpy array, convert to torch tensor for comparison |
| tosa_ref_outputs.append(torch.from_numpy(tosa_ref_output)) |
| |
| return tosa_ref_outputs |
| |
| |
| def prep_data_for_save( |
| data, is_quantized: bool, input_name: str, quant_param: QuantizationParams |
| ): |
| data_np = np.array(data.detach(), order="C").astype( |
| f"{data.dtype}".replace("torch.", "") |
| ) |
| |
| if is_quantized: |
| assert quant_param.node_name in input_name, ( |
| f"The quantization params name '{quant_param.node_name}' does not " |
| f"match the input tensor name '{input_name}'." |
| ) |
| data_np = ( |
| ((data_np / np.float32(quant_param.scale)) + quant_param.zp) |
| .round() |
| .clip(quant_param.qmin, quant_param.qmax) |
| .astype( |
| f"{quant_param.dtype}".replace("torch.", "") |
| ) # Use string format of dtype to convert to numpy dtype |
| ) |
| return data_np |
| |
| |
| def save_npy( |
| path: str, |
| data, |
| is_quantized: bool, |
| input_name: str, |
| quant_param: QuantizationParams, |
| ) -> str: |
| """Serializes and saves 'data' as a .npy file, possibly quantizing it before. |
| |
| Parameters: |
| path: the directory where to save the data. |
| data: the data to save. |
| is_quantized: whether to quantize the data before saving it. |
| input_name: the name of the file, without file-ending. |
| quant_param: the parameters to use for quantization. |
| Returns: |
| the full file path of the output. |
| """ |
| data_np = prep_data_for_save(data, is_quantized, input_name, quant_param) |
| file_path = os.path.join(path, input_name + ".npy") |
| np.save(file_path, data_np, allow_pickle=False) |
| |
| return file_path |
| |
| |
| def save_bytes( |
| path: str, |
| data, |
| is_quantized: bool, |
| input_name: str, |
| quant_param: QuantizationParams, |
| ) -> str: |
| """Serializes and saves 'data' in byte format, possibly quantizing it before. |
| |
| Parameters: |
| path: the directory where to save the data. |
| data: the data to save. |
| is_quantized: whether to quantize the data before saving it. |
| input_name: the name of the file, without file-ending. |
| quant_param: the parameters to use for quantization. |
| Returns: |
| the full file path of the output. |
| """ |
| data_np = prep_data_for_save(data, is_quantized, input_name, quant_param) |
| file_path = os.path.join(path, input_name + ".bin") |
| with open(file_path, "w+b") as f: |
| data_np_bytes = data_np.tobytes() |
| f.write(data_np_bytes) |
| |
| return file_path |
| |
| |
| def _run_cmd(cmd: List[str], check=True) -> subprocess.CompletedProcess[bytes]: |
| """ |
| Run a command and check for errors. |
| |
| Args: |
| cmd (List[str]): The command to run as a list. |
| """ |
| try: |
| result = subprocess.run(cmd, check=check, capture_output=True) |
| return result |
| except subprocess.CalledProcessError as e: |
| arg_string = " ".join(cmd) |
| raise RuntimeError( |
| f"Failed running command {arg_string}\nStderr: {e.stderr.decode()}\nStdout: {e.stdout.decode()}" |
| ) |
| |
| |
| def dbg_tosa_fb_to_json(tosa_fb: bytes) -> Dict: |
| """ |
| This function is used to dump the TOSA flatbuffer to a human readable |
| format, using flatc. It is used for debugging purposes. |
| """ |
| |
| tmp = tempfile.mkdtemp() |
| tosa_input_file = os.path.join(tmp, "output.tosa") |
| with open(tosa_input_file, "wb") as f: |
| f.write(tosa_fb) |
| |
| arm_backend_path = os.path.realpath(os.path.dirname(__file__) + "/..") |
| tosa_schema_file = os.path.join( |
| arm_backend_path, "third-party/serialization_lib/schema/tosa.fbs" |
| ) |
| assert os.path.exists( |
| tosa_schema_file |
| ), f"tosa_schema_file: {tosa_schema_file} does not exist" |
| |
| assert shutil.which("flatc") is not None |
| cmd_flatc = [ |
| "flatc", |
| "--json", |
| "--strict-json", |
| "-o", |
| tmp, |
| "--raw-binary", |
| "-t", |
| tosa_schema_file, |
| "--", |
| tosa_input_file, |
| ] |
| _run_cmd(cmd_flatc) |
| with open(os.path.join(tmp, "output.json"), "r") as f: |
| json_out = json.load(f) |
| |
| # Cast float tensors to proper dtype. |
| try: |
| for region in json_out["regions"]: |
| for block in region["blocks"]: |
| for tensor in block["tensors"]: |
| if "data" in tensor: |
| if tensor["type"] == "FP32": |
| data = np.array(tensor["data"]) |
| data = data.astype(np.int8) |
| data = np.frombuffer(data, dtype=np.float32) |
| data = data.reshape(tensor["shape"]) |
| tensor["data"] = data |
| except Exception: |
| # This is just nice-to-have if it works, don't care if it fails. |
| pass |
| |
| return json_out |
