Add BenchmarkSuite to load benchmarks. (#8753)
diff --git a/build_tools/benchmarks/common/CMakeLists.txt b/build_tools/benchmarks/common/CMakeLists.txt
index 315e08e..f9e2e94 100644
--- a/build_tools/benchmarks/common/CMakeLists.txt
+++ b/build_tools/benchmarks/common/CMakeLists.txt
@@ -31,12 +31,21 @@
"benchmark_config_test.py"
)
+iree_py_test(
+ NAME
+ benchmark_suite_test
+ SRCS
+ "benchmark_suite_test.py"
+)
+
# TODO(#8708): Temporary solution to fix python path for tests.
set_property(TEST "build_tools/benchmarks/common/linux_device_utils_test"
APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
set_property(TEST "build_tools/benchmarks/common/common_arguments_test"
- APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
+ APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
set_property(TEST "build_tools/benchmarks/common/benchmark_config_test"
- APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
+ APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
+set_property(TEST "build_tools/benchmarks/common/benchmark_suite_test"
+ APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
endif()
diff --git a/build_tools/benchmarks/common/benchmark_suite.py b/build_tools/benchmarks/common/benchmark_suite.py
index bd72617..564a657 100644
--- a/build_tools/benchmarks/common/benchmark_suite.py
+++ b/build_tools/benchmarks/common/benchmark_suite.py
@@ -29,16 +29,150 @@
└── <compiled-iree-model>-<sha1>.vmfb
"""
-import re
+import collections
import os
+import re
-from typing import List, Optional, Sequence
-
-from .benchmark_definition import DeviceInfo, BenchmarkInfo
+from common.benchmark_definition import DeviceInfo, BenchmarkInfo
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Sequence, Tuple
# All benchmarks' relative path against root build directory.
BENCHMARK_SUITE_REL_PATH = "benchmark_suites"
+MODEL_TOOLFILE_NAME = "tool"
+
+
+@dataclass
+class BenchmarkCase:
+ """Represents a benchmark case.
+
+ model_name_with_tags: the source model with tags, e.g., 'MobileSSD-fp32'.
+ bench_mode: the benchmark mode, e.g., '1-thread,big-core'.
+ target_arch: the target CPU/GPU architature, e.g., 'GPU-Adreno'.
+ driver: the IREE driver to run with, e.g., 'dylib'.
+ benchmark_case_dir: the path to benchmark case directory.
+ benchmark_tool_name: the benchmark tool, e.g., 'iree-benchmark-module'.
+ """
+
+ model_name_with_tags: str
+ bench_mode: str
+ target_arch: str
+ driver: str
+ benchmark_case_dir: str
+ benchmark_tool_name: str
+
+
+class BenchmarkSuite(object):
+ """Represents the benchmarks in benchmark suite directory."""
+
+ def __init__(self, suite_map: Dict[str, List[BenchmarkCase]]):
+ """Construct a benchmark suite.
+
+ Args:
+ suites: the map of benchmark cases keyed by category directories.
+ """
+ self.suite_map = suite_map
+ self.category_map = dict((os.path.basename(category_dir), category_dir)
+ for category_dir in self.suite_map.keys())
+
+ def list_categories(self) -> List[Tuple[str, str]]:
+ """Returns all categories and their directories.
+
+ Returns:
+ A tuple of (category name, category dir).
+ """
+ category_list = [(name, path) for name, path in self.category_map.items()]
+ # Fix the order of category list.
+ category_list.sort(key=lambda category: category[0])
+ return category_list
+
+ def filter_benchmarks_for_category(
+ self,
+ category: str,
+ available_drivers: Sequence[str],
+ cpu_target_arch_filter: str,
+ gpu_target_arch_filter: str,
+ driver_filter: Optional[str] = None,
+ mode_filter: Optional[str] = None,
+ model_name_filter: Optional[str] = None) -> Sequence[BenchmarkCase]:
+ """Filters benchmarks in a specific category for the given device.
+ Args:
+ category: the specific benchmark category.
+ available_drivers: list of drivers supported by the tools.
+ cpu_target_arch_filter: CPU target architecture filter regex.
+ gpu_target_arch_filter: GPU target architecture filter regex.
+ driver_filter: driver filter regex.
+ mode_filter: benchmark mode regex.
+ model_name_filter: model name regex.
+ Returns:
+ A list of matched benchmark cases.
+ """
+
+ category_dir = self.category_map.get(category)
+ if category_dir is None:
+ return []
+
+ chosen_cases = []
+ for benchmark_case in self.suite_map[category_dir]:
+ matched_driver = (benchmark_case.driver in available_drivers) and (
+ driver_filter is None or
+ re.match(driver_filter, benchmark_case.driver) is not None)
+ matched_arch = (re.match(cpu_target_arch_filter,
+ benchmark_case.target_arch) is not None or
+ re.match(gpu_target_arch_filter,
+ benchmark_case.target_arch) is not None)
+ matched_mode = (mode_filter is None or re.match(
+ mode_filter, benchmark_case.bench_mode) is not None)
+
+ # For backward compatibility, model_name_filter matches against the string:
+ # <model name with tags>/<benchmark case name>
+ model_and_case_name = f"{benchmark_case.model_name_with_tags}/{os.path.basename(benchmark_case.benchmark_case_dir)}"
+ matched_model_name = (model_name_filter is None or re.match(
+ model_name_filter, model_and_case_name) is not None)
+
+ if (matched_driver and matched_arch and matched_model_name and
+ matched_mode):
+ chosen_cases.append(benchmark_case)
+
+ return chosen_cases
+
+ @staticmethod
+ def load_from_benchmark_suite_dir(benchmark_suite_dir: str):
+ """Scans and loads the benchmarks under the directory."""
+
+ suite_map: Dict[str, List[BenchmarkCase]] = collections.defaultdict(list)
+ for benchmark_case_dir, _, _ in os.walk(benchmark_suite_dir):
+ model_dir, benchmark_name = os.path.split(benchmark_case_dir)
+ # Take the benchmark directory name and see if it matches the benchmark
+ # naming convention:
+ # <iree-driver>__<target-architecture>__<benchmark_mode>
+ segments = benchmark_name.split("__")
+ if len(segments) != 3 or not segments[0].startswith("iree-"):
+ continue
+
+ iree_driver, target_arch, bench_mode = segments
+ iree_driver = iree_driver[len("iree-"):].lower()
+ target_arch = target_arch.lower()
+
+ # The path of model_dir is expected to be:
+ # <benchmark_suite_dir>/<category>/<model_name>-<model_tags>
+ category_dir, model_name_with_tags = os.path.split(model_dir)
+
+ with open(os.path.join(benchmark_case_dir, MODEL_TOOLFILE_NAME),
+ "r") as f:
+ tool_name = f.read().strip()
+
+ suite_map[category_dir].append(
+ BenchmarkCase(model_name_with_tags=model_name_with_tags,
+ bench_mode=bench_mode,
+ target_arch=target_arch,
+ driver=iree_driver,
+ benchmark_case_dir=benchmark_case_dir,
+ benchmark_tool_name=tool_name))
+
+ return BenchmarkSuite(suite_map=suite_map)
+
def compose_info_object(device_info: DeviceInfo, benchmark_category_dir: str,
benchmark_case_dir: str) -> BenchmarkInfo:
@@ -137,7 +271,6 @@
chosen = False
if matched_driver and matched_arch and matched_model_name and matched_mode:
matched_benchmarks.append(root)
- chosen = True
if verbose:
print(f"dir: {root}")
diff --git a/build_tools/benchmarks/common/benchmark_suite_test.py b/build_tools/benchmarks/common/benchmark_suite_test.py
new file mode 100644
index 0000000..68085bc
--- /dev/null
+++ b/build_tools/benchmarks/common/benchmark_suite_test.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+# 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
+import tempfile
+import unittest
+
+from common.benchmark_suite import BenchmarkCase, BenchmarkSuite
+
+
+class BenchmarkSuiteTest(unittest.TestCase):
+
+ def test_list_categories(self):
+ suite = BenchmarkSuite({
+ "suite/TFLite": [],
+ "suite/PyTorch": [],
+ })
+
+ self.assertEqual(suite.list_categories(), [("PyTorch", "suite/PyTorch"),
+ ("TFLite", "suite/TFLite")])
+
+ def test_filter_benchmarks_for_category(self):
+ case1 = BenchmarkCase(model_name_with_tags="deepnet",
+ bench_mode="1-thread,full-inference",
+ target_arch="ARMv8",
+ driver="dylib",
+ benchmark_case_dir="case1",
+ benchmark_tool_name="tool")
+ case2 = BenchmarkCase(model_name_with_tags="deepnetv2-f32",
+ bench_mode="full-inference",
+ target_arch="Mali",
+ driver="vulkan",
+ benchmark_case_dir="case2",
+ benchmark_tool_name="tool")
+ suite = BenchmarkSuite({
+ "suite/TFLite": [case1, case2],
+ })
+
+ both_benchmarks = suite.filter_benchmarks_for_category(
+ category="TFLite",
+ available_drivers=["dylib", "vulkan"],
+ cpu_target_arch_filter="ARMv8",
+ gpu_target_arch_filter="Mali",
+ driver_filter=None,
+ mode_filter=".*full-inference.*",
+ model_name_filter="deepnet.*")
+ gpu_benchmarks = suite.filter_benchmarks_for_category(
+ category="TFLite",
+ available_drivers=["dylib", "vulkan"],
+ cpu_target_arch_filter="Unknown",
+ gpu_target_arch_filter="Mali",
+ driver_filter="vulkan",
+ mode_filter=".*full-inference.*",
+ model_name_filter="deepnet.*/case2")
+
+ self.assertEqual(both_benchmarks, [case1, case2])
+ self.assertEqual(gpu_benchmarks, [case2])
+
+ def test_filter_benchmarks_for_nonexistent_category(self):
+ suite = BenchmarkSuite({
+ "suite/TFLite": [],
+ })
+
+ benchmarks = suite.filter_benchmarks_for_category(
+ category="PyTorch",
+ available_drivers=[],
+ cpu_target_arch_filter="ARMv8",
+ gpu_target_arch_filter="Mali-G78")
+
+ self.assertEqual(benchmarks, [])
+
+ def test_load_from_benchmark_suite_dir(self):
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ tflite_dir = os.path.join(tmp_dir, "TFLite")
+ pytorch_dir = os.path.join(tmp_dir, "PyTorch")
+ case1 = BenchmarkSuiteTest.__create_bench(tflite_dir,
+ model="deepnet",
+ bench_mode="4-thread,full",
+ target_arch="cpu-armv8",
+ driver="dylib",
+ tool="run-cpu-bench")
+ case2 = BenchmarkSuiteTest.__create_bench(pytorch_dir,
+ model="deepnetv2",
+ bench_mode="full-inference",
+ target_arch="gpu-mali",
+ driver="vulkan",
+ tool="run-gpu-bench")
+
+ suite = BenchmarkSuite.load_from_benchmark_suite_dir(tmp_dir)
+
+ self.assertEqual(suite.list_categories(), [("PyTorch", pytorch_dir),
+ ("TFLite", tflite_dir)])
+ self.assertEqual(
+ suite.filter_benchmarks_for_category(
+ category="PyTorch",
+ available_drivers=["vulkan"],
+ cpu_target_arch_filter="cpu-armv8",
+ gpu_target_arch_filter="gpu-mali"), [case2])
+
+ @staticmethod
+ def __create_bench(dir_path: str, model: str, bench_mode: str,
+ target_arch: str, driver: str, tool: str):
+ case_name = f"iree-{driver}__{target_arch}__{bench_mode}"
+ bench_path = os.path.join(dir_path, model, case_name)
+ os.makedirs(bench_path)
+ with open(os.path.join(bench_path, "tool"), "w") as f:
+ f.write(tool)
+
+ return BenchmarkCase(model_name_with_tags=model,
+ bench_mode=bench_mode,
+ target_arch=target_arch,
+ driver=driver,
+ benchmark_case_dir=bench_path,
+ benchmark_tool_name=tool)
+
+
+if __name__ == "__main__":
+ unittest.main()
