| # Copyright 2022 The IREE Authors |
| # |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. |
| # See https://llvm.org/LICENSE.txt for license information. |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| import os |
| from typing import Optional |
| |
| from common.benchmark_command import * |
| from common.benchmark_command_factory import BenchmarkCommandFactory |
| |
| _DEFAULT_NUM_BENCHMARK_RUNS = 50 |
| _DEFAULT_NUM_THREADS = 1 |
| |
| |
| class TfliteWrapper(TFLiteBenchmarkCommand): |
| """Specializes the benchmark command to use TFLite.""" |
| |
| def __init__( |
| self, |
| benchmark_binary: str, |
| model_name: str, |
| model_path: str, |
| input_layer: Optional[str] = None, |
| input_shape: Optional[str] = None, |
| driver: str = "cpu", |
| num_threads: int = _DEFAULT_NUM_THREADS, |
| num_runs: int = _DEFAULT_NUM_BENCHMARK_RUNS, |
| taskset: Optional[str] = None, |
| ): |
| super().__init__( |
| benchmark_binary, |
| model_name, |
| model_path, |
| num_threads, |
| num_runs, |
| taskset=taskset, |
| ) |
| self.driver = driver |
| if input_layer and input_shape: |
| self.args.append("--input_layer=%s" % input_layer) |
| self.args.append("--input_layer_shape=%s" % input_shape) |
| |
| |
| class IreeWrapper(IreeBenchmarkCommand): |
| """Specializes the benchmark command to use IREE.""" |
| |
| def __init__( |
| self, |
| benchmark_binary: str, |
| model_name: str, |
| model_path: str, |
| function_input: str, |
| driver: str = "local-task", |
| num_threads: int = _DEFAULT_NUM_THREADS, |
| num_runs: int = _DEFAULT_NUM_BENCHMARK_RUNS, |
| taskset: Optional[str] = None, |
| ): |
| super().__init__( |
| benchmark_binary, |
| model_name, |
| model_path, |
| num_threads, |
| num_runs, |
| taskset=taskset, |
| ) |
| self.driver = driver |
| self.args.append("--function=main") |
| self.args.append('--input="%s"' % function_input) |
| |
| |
| class SimpleCommandFactory(BenchmarkCommandFactory): |
| """ |
| Generates `BenchmarkCommand` objects specific to running series of simple models. |
| |
| A model is considered simple if its inputs can be generically generated based |
| on expected signature only without affecting behavior. |
| """ |
| |
| def __init__( |
| self, |
| base_dir: str, |
| model_name: str, |
| function_input: str, |
| input_name: Optional[str] = None, |
| input_layer: Optional[str] = None, |
| ): |
| self._model_name = model_name |
| self._function_input = function_input |
| self._input_name = input_name |
| self._input_layer = input_layer |
| self._base_dir = base_dir |
| self._iree_benchmark_binary_path = os.path.join( |
| base_dir, "iree-benchmark-module" |
| ) |
| self._tflite_benchmark_binary_path = os.path.join(base_dir, "benchmark_model") |
| # Required to be set, but no test data used yet. |
| self._tflite_test_data_dir = os.path.join(self._base_dir, "test_data") |
| |
| def generate_benchmark_commands( |
| self, device: str, driver: str |
| ) -> list[BenchmarkCommand]: |
| if device == "desktop" and driver == "cpu": |
| return self._generate_cpu(device) |
| elif device == "desktop" and driver == "gpu": |
| return self._generate_gpu("cuda") |
| elif device == "mobile" and driver == "cpu": |
| return self._generate_cpu(device) |
| elif device == "mobile" and driver == "gpu": |
| return self._generate_gpu("vulkan") |
| else: |
| print("Warning! Not a valid configuration.") |
| return [] |
| |
| def _generate_cpu(self, device: str): |
| commands = [] |
| # Generate TFLite benchmarks. |
| tflite_model_path = os.path.join( |
| self._base_dir, "models", "tflite", self._model_name + ".tflite" |
| ) |
| tflite = TfliteWrapper( |
| self._tflite_benchmark_binary_path, |
| self._model_name, |
| tflite_model_path, |
| self._input_name, |
| driver="cpu", |
| ) |
| commands.append(tflite) |
| |
| tflite_noxnn = TfliteWrapper( |
| self._tflite_benchmark_binary_path, |
| self._model_name + "_noxnn", |
| tflite_model_path, |
| self._input_name, |
| driver="cpu", |
| ) |
| tflite_noxnn.args.append("--use_xnnpack=false") |
| commands.append(tflite_noxnn) |
| |
| # Generate IREE benchmarks. |
| driver = "local-task" |
| backend = "llvm-cpu" |
| |
| iree_model_path = os.path.join( |
| self._base_dir, "models", "iree", backend, self._model_name + ".vmfb" |
| ) |
| iree = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| self._model_name, |
| iree_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree) |
| |
| model_padfuse_name = self._model_name + "_padfuse" |
| iree_padfuse_model_path = os.path.join( |
| self._base_dir, "models", "iree", backend, model_padfuse_name + ".vmfb" |
| ) |
| iree_padfuse = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| model_padfuse_name, |
| iree_padfuse_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree_padfuse) |
| |
| # Test mmt4d only on mobile. |
| if device == "mobile": |
| model_mmt4d_name = self._model_name + "_mmt4d" |
| iree_mmt4d_model_path = os.path.join( |
| self._base_dir, "models", "iree", backend, model_mmt4d_name + ".vmfb" |
| ) |
| iree_mmt4d = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| model_mmt4d_name, |
| iree_mmt4d_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree_mmt4d) |
| |
| model_im2col_mmt4d_name = self._model_name + "_im2col_mmt4d" |
| iree_im2col_mmt4d_model_path = os.path.join( |
| self._base_dir, |
| "models", |
| "iree", |
| backend, |
| model_im2col_mmt4d_name + ".vmfb", |
| ) |
| iree_im2col_mmt4d = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| model_im2col_mmt4d_name, |
| iree_im2col_mmt4d_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree_im2col_mmt4d) |
| |
| return commands |
| |
| def _generate_gpu(self, driver: str): |
| commands = [] |
| tflite_model_path = os.path.join( |
| self._base_dir, "models", "tflite", self._model_name + ".tflite" |
| ) |
| tflite = TfliteWrapper( |
| self._tflite_benchmark_binary_path, |
| self._model_name, |
| tflite_model_path, |
| self._input_name, |
| self._input_layer, |
| driver="gpu", |
| ) |
| tflite.args.append("--gpu_precision_loss_allowed=false") |
| commands.append(tflite) |
| |
| tflite_noxnn = TfliteWrapper( |
| self._tflite_benchmark_binary_path, |
| self._model_name + "_noxnn", |
| tflite_model_path, |
| self._input_name, |
| self._input_layer, |
| driver="gpu", |
| ) |
| tflite.args.append("--use_xnnpack=false") |
| commands.append(tflite_noxnn) |
| |
| tflite_fp16 = TfliteWrapper( |
| self._tflite_benchmark_binary_path, |
| self._model_name + "_fp16", |
| tflite_model_path, |
| self._input_name, |
| self._input_layer, |
| driver="gpu", |
| ) |
| tflite.args.append("--gpu_precision_loss_allowed=true") |
| commands.append(tflite_fp16) |
| |
| iree_model_path = os.path.join( |
| self._base_dir, "models", "iree", driver, self._model_name + ".vmfb" |
| ) |
| iree = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| self._model_name, |
| iree_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree) |
| |
| iree_model_path = os.path.join( |
| self._base_dir, "models", "iree", driver, self._model_name + "_fp16.vmfb" |
| ) |
| iree = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| self._model_name + "_fp16", |
| iree_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree) |
| |
| iree_model_path = os.path.join( |
| self._base_dir, "models", "iree", driver, self._model_name + "_padfuse.vmfb" |
| ) |
| iree = IreeWrapper( |
| self._iree_benchmark_binary_path, |
| self._model_name + "_padfuse", |
| iree_model_path, |
| self._function_input, |
| driver=driver, |
| ) |
| commands.append(iree) |
| return commands |
