blob: 2e49c3076cbae21e19ca2db0b1ff18a423b9984c [file] [log] [blame]
# 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