build_tools/benchmarks/common/benchmark_driver.py - 3p/openxla/iree - Git at Google

 # 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 json
 import pathlib
 import time
 from typing import Dict, Optional, Sequence, Tuple
 from common.benchmark_suite import BenchmarkCase, BenchmarkSuite
 from common.benchmark_config import BENCHMARK_RESULTS_REL_PATH, CAPTURES_REL_PATH, BenchmarkConfig
 from common.benchmark_definition import BenchmarkInfo, BenchmarkResults, BenchmarkRun, DeviceInfo


 class BenchmarkDriver(object):
   """Abstract driver runs the whole benchmark flow."""

   def __init__(self,
                device_info: DeviceInfo,
                benchmark_config: BenchmarkConfig,
                benchmark_suite: BenchmarkSuite,
                benchmark_grace_time: float = 0.0,
                verbose: bool = False):
     self.device_info = device_info
     self.config = benchmark_config
     self.benchmark_suite = benchmark_suite
     self.benchmark_grace_time = benchmark_grace_time
     self.verbose = verbose
     self.finished_benchmarks: Dict[str, Tuple[BenchmarkInfo, pathlib.Path]] = {}
     self.finished_captures: Dict[str, Tuple[BenchmarkInfo, pathlib.Path]] = {}
     self.benchmark_errors = []

   def run_benchmark_case(self, benchmark_case: BenchmarkCase,
                          benchmark_results_filename: Optional[pathlib.Path],
                          capture_filename: Optional[pathlib.Path]) -> None:
     """Runs the benchmark case and returns the results.

     Args:
       benchmark_case: the benchmark_case.
       benchmark_results_filename: the path to store benchmark results.
         Benchmarking is required if set.
       capture_filename: the path to store captured trace. Trace capturing is
         required if set.

     Raises:
       Exception during benchmarking.
     """
     raise NotImplementedError("Should be overwritten by a subclass.")

   def add_previous_benchmarks_and_captures(
       self, previous_directory: pathlib.Path) -> None:
     """Collect names of previous benchmarks and captures that should be skipped
     and merged into the results.
     """

     def get_key_value_pair(path: pathlib.Path):
       name = path.stem
       info = BenchmarkInfo.from_device_info_and_name(self.device_info, name)
       return (str(info), (info, path))

     previous_benchmark_filenames = set()
     previous_capture_filenames = set()
     previous_benchmarks_dir = previous_directory / BENCHMARK_RESULTS_REL_PATH
     if previous_benchmarks_dir.is_dir():
       previous_benchmark_filenames = set(
           previous_benchmarks_dir / p
           for p in previous_benchmarks_dir.iterdir()
           if p.suffix == ".json")

     previous_captures_dir = previous_directory / CAPTURES_REL_PATH
     if previous_captures_dir.is_dir():
       previous_capture_filenames = set(previous_captures_dir / p
                                        for p in previous_captures_dir.iterdir()
                                        if p.suffix == ".tracy")

     self.finished_benchmarks.update(
         get_key_value_pair(p) for p in previous_benchmark_filenames)
     self.finished_captures.update(
         get_key_value_pair(p) for p in previous_capture_filenames)

   def run(self) -> None:
     """Execute the benchmark flow.

     It performs the following steps:
       1. Enumerate all categories in the benchmark suites.
       2. For each category, enumerate and filter benchmark cases.
       3. Call 'run_benchmark_case' for each benchmark case.
       4. Collect the benchmark results and captures.
     """

     do_capture = self.config.trace_capture_config is not None

     self.config.benchmark_results_dir.mkdir(parents=True, exist_ok=True)
     if do_capture:
       self.config.trace_capture_config.capture_tmp_dir.mkdir(parents=True,
                                                              exist_ok=True)

     cpu_target_arch = self.device_info.get_iree_cpu_arch_name()
     gpu_target_arch = self.device_info.get_iree_gpu_arch_name()
     drivers, loaders = self.__get_available_drivers_and_loaders()

     for category, _ in self.benchmark_suite.list_categories():
       benchmark_cases = self.benchmark_suite.filter_benchmarks_for_category(
           category=category,
           available_drivers=drivers,
           available_loaders=loaders,
           cpu_target_arch_filter=f"^{cpu_target_arch}$",
           gpu_target_arch_filter=f"^{gpu_target_arch}$",
           driver_filter=self.config.driver_filter,
           mode_filter=self.config.mode_filter,
           model_name_filter=self.config.model_name_filter)

       for benchmark_case in benchmark_cases:
         (benchmark_info, benchmark_results_filename,
          capture_filename) = self.__get_benchmark_info_and_output_paths(
              category, benchmark_case)

         # Skip if no need to benchmark and capture.
         if not benchmark_results_filename and not capture_filename:
           continue

         benchmark_key = str(benchmark_info)
         print(f"--> Benchmark started: {benchmark_key} <--")

         try:
           self.run_benchmark_case(benchmark_case, benchmark_results_filename,
                                   capture_filename)
         except Exception as e:
           if not self.config.keep_going:
             raise e

           print(f"Processing of benchmark failed with: {e}")
           self.benchmark_errors.append(e)
           continue
         finally:
           # Some grace time.
           time.sleep(self.benchmark_grace_time)

         print("Benchmark completed")

         if benchmark_results_filename:
           self.finished_benchmarks[benchmark_key] = (benchmark_info,
                                                      benchmark_results_filename)
         if capture_filename:
           self.finished_captures[benchmark_key] = (benchmark_info,
                                                    capture_filename)

   def get_benchmark_results(self) -> BenchmarkResults:
     """Returns the finished benchmark results."""

     results = BenchmarkResults()
     results.set_commit(self.config.git_commit_hash)

     finished_benchmarks = list(self.finished_benchmarks.items())
     finished_benchmarks.sort(key=lambda b: b[0])

     for _, value in finished_benchmarks:
       benchmark_info, path = value
       with open(path) as f:
         result_json_object = json.loads(f.read())
       benchmark_run = BenchmarkRun(benchmark_info,
                                    result_json_object["context"],
                                    result_json_object["benchmarks"])
       results.benchmarks.append(benchmark_run)

     return results

   def get_benchmark_result_filenames(self) -> Sequence[pathlib.Path]:
     """Returns the json file paths of finished benchmarks."""
     return list(path for _, path in self.finished_benchmarks.values())

   def get_capture_filenames(self) -> Sequence[pathlib.Path]:
     """Returns the tracy file paths of finished captures."""
     return list(path for _, path in self.finished_captures.values())

   def get_benchmark_errors(self):
     """Returns the exceptions captured during benchmarking."""
     return self.benchmark_errors

   def __get_benchmark_info_and_output_paths(self, category: str,
                                             benchmark_case: BenchmarkCase):
     """Get benchmark info and paths for the results and capture. The path of
     results/capture is None if the benchmark/capture doesn't need to be run.
     """
     benchmark_info = self.__get_benchmark_info_from_case(
         category=category, benchmark_case=benchmark_case)
     benchmark_name = str(benchmark_info)

     benchmark_results_filename = None
     if (benchmark_name not in self.finished_benchmarks and
         self.config.normal_benchmark_tool_dir):
       benchmark_results_filename = self.config.benchmark_results_dir / f"{benchmark_name}.json"

     capture_filename = None
     if (benchmark_name not in self.finished_captures and
         self.config.trace_capture_config):
       capture_filename = self.config.trace_capture_config.capture_tmp_dir / f"{benchmark_name}.tracy"

     return (benchmark_info, benchmark_results_filename, capture_filename)

   def __get_benchmark_info_from_case(
       self, category: str, benchmark_case: BenchmarkCase) -> BenchmarkInfo:
     if benchmark_case.run_config is None:
       # TODO(#11076): Remove legacy path.
       return BenchmarkInfo(model_name=benchmark_case.model_name,
                            model_tags=benchmark_case.model_tags,
                            model_source=category,
                            bench_mode=benchmark_case.bench_mode,
                            driver_info=benchmark_case.driver_info,
                            device_info=self.device_info)

     run_tags = benchmark_case.run_config.module_execution_config.tags
     compile_tags = benchmark_case.run_config.module_generation_config.compile_config.tags
     return BenchmarkInfo(model_name=benchmark_case.model_name,
                          model_tags=benchmark_case.model_tags,
                          model_source=category,
                          bench_mode=run_tags,
                          compile_tags=compile_tags,
                          driver_info=benchmark_case.driver_info,
                          device_info=self.device_info,
                          run_config_id=benchmark_case.run_config.composite_id)

   def __get_available_drivers_and_loaders(
       self) -> Tuple[Sequence[str], Sequence[str]]:
     any_tool_dir = (self.config.normal_benchmark_tool_dir
                     if self.config.normal_benchmark_tool_dir else
                     self.config.trace_capture_config.traced_benchmark_tool_dir)
     config_txt_file_path = any_tool_dir / "build_config.txt"
     config_txt_file_lines = config_txt_file_path.read_text().splitlines()

     available_drivers = []
     available_loaders = []
     for line in config_txt_file_lines:
       name, value = line.strip().split("=")
       if value != "ON":
         continue
       if name == "IREE_HAL_DRIVER_CUDA":
         available_drivers.append("cuda")
       elif name == "IREE_HAL_DRIVER_LOCAL_SYNC":
         available_drivers.append("local-sync")
       elif name == "IREE_HAL_DRIVER_LOCAL_TASK":
         available_drivers.append("local-task")
       elif name == "IREE_HAL_DRIVER_VULKAN":
         available_drivers.append("vulkan")
       elif name == "IREE_HAL_EXECUTABLE_LOADER_EMBEDDED_ELF":
         available_loaders.append("embedded-elf")
       elif name == "IREE_HAL_EXECUTABLE_LOADER_SYSTEM_LIBRARY":
         available_loaders.append("system-library")
       elif name == "IREE_HAL_EXECUTABLE_LOADER_VMVX_MODULE":
         available_loaders.append("vmvx-module")
       else:
         continue

     if self.verbose:
       available_drivers_str = ', '.join(available_drivers)
       print(f"Available drivers: {available_drivers_str}")
       available_loaders_str = ', '.join(available_loaders)
       print(f"Available loaders: {available_loaders_str}")

     return available_drivers, available_loaders
	# 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 json
	import pathlib
	import time
	from typing import Dict, Optional, Sequence, Tuple
	from common.benchmark_suite import BenchmarkCase, BenchmarkSuite
	from common.benchmark_config import BENCHMARK_RESULTS_REL_PATH, CAPTURES_REL_PATH, BenchmarkConfig
	from common.benchmark_definition import BenchmarkInfo, BenchmarkResults, BenchmarkRun, DeviceInfo


	class BenchmarkDriver(object):
	"""Abstract driver runs the whole benchmark flow."""

	def __init__(self,
	device_info: DeviceInfo,
	benchmark_config: BenchmarkConfig,
	benchmark_suite: BenchmarkSuite,
	benchmark_grace_time: float = 0.0,
	verbose: bool = False):
	self.device_info = device_info
	self.config = benchmark_config
	self.benchmark_suite = benchmark_suite
	self.benchmark_grace_time = benchmark_grace_time
	self.verbose = verbose
	self.finished_benchmarks: Dict[str, Tuple[BenchmarkInfo, pathlib.Path]] = {}
	self.finished_captures: Dict[str, Tuple[BenchmarkInfo, pathlib.Path]] = {}
	self.benchmark_errors = []

	def run_benchmark_case(self, benchmark_case: BenchmarkCase,
	benchmark_results_filename: Optional[pathlib.Path],
	capture_filename: Optional[pathlib.Path]) -> None:
	"""Runs the benchmark case and returns the results.

	Args:
	benchmark_case: the benchmark_case.
	benchmark_results_filename: the path to store benchmark results.
	Benchmarking is required if set.
	capture_filename: the path to store captured trace. Trace capturing is
	required if set.

	Raises:
	Exception during benchmarking.
	"""
	raise NotImplementedError("Should be overwritten by a subclass.")

	def add_previous_benchmarks_and_captures(
	self, previous_directory: pathlib.Path) -> None:
	"""Collect names of previous benchmarks and captures that should be skipped
	and merged into the results.
	"""

	def get_key_value_pair(path: pathlib.Path):
	name = path.stem
	info = BenchmarkInfo.from_device_info_and_name(self.device_info, name)
	return (str(info), (info, path))

	previous_benchmark_filenames = set()
	previous_capture_filenames = set()
	previous_benchmarks_dir = previous_directory / BENCHMARK_RESULTS_REL_PATH
	if previous_benchmarks_dir.is_dir():
	previous_benchmark_filenames = set(
	previous_benchmarks_dir / p
	for p in previous_benchmarks_dir.iterdir()
	if p.suffix == ".json")

	previous_captures_dir = previous_directory / CAPTURES_REL_PATH
	if previous_captures_dir.is_dir():
	previous_capture_filenames = set(previous_captures_dir / p
	for p in previous_captures_dir.iterdir()
	if p.suffix == ".tracy")

	self.finished_benchmarks.update(
	get_key_value_pair(p) for p in previous_benchmark_filenames)
	self.finished_captures.update(
	get_key_value_pair(p) for p in previous_capture_filenames)

	def run(self) -> None:
	"""Execute the benchmark flow.

	It performs the following steps:
	1. Enumerate all categories in the benchmark suites.
	2. For each category, enumerate and filter benchmark cases.
	3. Call 'run_benchmark_case' for each benchmark case.
	4. Collect the benchmark results and captures.
	"""

	do_capture = self.config.trace_capture_config is not None

	self.config.benchmark_results_dir.mkdir(parents=True, exist_ok=True)
	if do_capture:
	self.config.trace_capture_config.capture_tmp_dir.mkdir(parents=True,
	exist_ok=True)

	cpu_target_arch = self.device_info.get_iree_cpu_arch_name()
	gpu_target_arch = self.device_info.get_iree_gpu_arch_name()
	drivers, loaders = self.__get_available_drivers_and_loaders()

	for category, _ in self.benchmark_suite.list_categories():
	benchmark_cases = self.benchmark_suite.filter_benchmarks_for_category(
	category=category,
	available_drivers=drivers,
	available_loaders=loaders,
	cpu_target_arch_filter=f"^{cpu_target_arch}$",
	gpu_target_arch_filter=f"^{gpu_target_arch}$",
	driver_filter=self.config.driver_filter,
	mode_filter=self.config.mode_filter,
	model_name_filter=self.config.model_name_filter)

	for benchmark_case in benchmark_cases:
	(benchmark_info, benchmark_results_filename,
	capture_filename) = self.__get_benchmark_info_and_output_paths(
	category, benchmark_case)

	# Skip if no need to benchmark and capture.
	if not benchmark_results_filename and not capture_filename:
	continue

	benchmark_key = str(benchmark_info)
	print(f"--> Benchmark started: {benchmark_key} <--")

	try:
	self.run_benchmark_case(benchmark_case, benchmark_results_filename,
	capture_filename)
	except Exception as e:
	if not self.config.keep_going:
	raise e

	print(f"Processing of benchmark failed with: {e}")
	self.benchmark_errors.append(e)
	continue
	finally:
	# Some grace time.
	time.sleep(self.benchmark_grace_time)

	print("Benchmark completed")

	if benchmark_results_filename:
	self.finished_benchmarks[benchmark_key] = (benchmark_info,
	benchmark_results_filename)
	if capture_filename:
	self.finished_captures[benchmark_key] = (benchmark_info,
	capture_filename)

	def get_benchmark_results(self) -> BenchmarkResults:
	"""Returns the finished benchmark results."""

	results = BenchmarkResults()
	results.set_commit(self.config.git_commit_hash)

	finished_benchmarks = list(self.finished_benchmarks.items())
	finished_benchmarks.sort(key=lambda b: b[0])

	for _, value in finished_benchmarks:
	benchmark_info, path = value
	with open(path) as f:
	result_json_object = json.loads(f.read())
	benchmark_run = BenchmarkRun(benchmark_info,
	result_json_object["context"],
	result_json_object["benchmarks"])
	results.benchmarks.append(benchmark_run)

	return results

	def get_benchmark_result_filenames(self) -> Sequence[pathlib.Path]:
	"""Returns the json file paths of finished benchmarks."""
	return list(path for _, path in self.finished_benchmarks.values())

	def get_capture_filenames(self) -> Sequence[pathlib.Path]:
	"""Returns the tracy file paths of finished captures."""
	return list(path for _, path in self.finished_captures.values())

	def get_benchmark_errors(self):
	"""Returns the exceptions captured during benchmarking."""
	return self.benchmark_errors

	def __get_benchmark_info_and_output_paths(self, category: str,
	benchmark_case: BenchmarkCase):
	"""Get benchmark info and paths for the results and capture. The path of
	results/capture is None if the benchmark/capture doesn't need to be run.
	"""
	benchmark_info = self.__get_benchmark_info_from_case(
	category=category, benchmark_case=benchmark_case)
	benchmark_name = str(benchmark_info)

	benchmark_results_filename = None
	if (benchmark_name not in self.finished_benchmarks and
	self.config.normal_benchmark_tool_dir):
	benchmark_results_filename = self.config.benchmark_results_dir / f"{benchmark_name}.json"

	capture_filename = None
	if (benchmark_name not in self.finished_captures and
	self.config.trace_capture_config):
	capture_filename = self.config.trace_capture_config.capture_tmp_dir / f"{benchmark_name}.tracy"

	return (benchmark_info, benchmark_results_filename, capture_filename)

	def __get_benchmark_info_from_case(
	self, category: str, benchmark_case: BenchmarkCase) -> BenchmarkInfo:
	if benchmark_case.run_config is None:
	# TODO(#11076): Remove legacy path.
	return BenchmarkInfo(model_name=benchmark_case.model_name,
	model_tags=benchmark_case.model_tags,
	model_source=category,
	bench_mode=benchmark_case.bench_mode,
	driver_info=benchmark_case.driver_info,
	device_info=self.device_info)

	run_tags = benchmark_case.run_config.module_execution_config.tags
	compile_tags = benchmark_case.run_config.module_generation_config.compile_config.tags
	return BenchmarkInfo(model_name=benchmark_case.model_name,
	model_tags=benchmark_case.model_tags,
	model_source=category,
	bench_mode=run_tags,
	compile_tags=compile_tags,
	driver_info=benchmark_case.driver_info,
	device_info=self.device_info,
	run_config_id=benchmark_case.run_config.composite_id)

	def __get_available_drivers_and_loaders(
	self) -> Tuple[Sequence[str], Sequence[str]]:
	any_tool_dir = (self.config.normal_benchmark_tool_dir
	if self.config.normal_benchmark_tool_dir else
	self.config.trace_capture_config.traced_benchmark_tool_dir)
	config_txt_file_path = any_tool_dir / "build_config.txt"
	config_txt_file_lines = config_txt_file_path.read_text().splitlines()

	available_drivers = []
	available_loaders = []
	for line in config_txt_file_lines:
	name, value = line.strip().split("=")
	if value != "ON":
	continue
	if name == "IREE_HAL_DRIVER_CUDA":
	available_drivers.append("cuda")
	elif name == "IREE_HAL_DRIVER_LOCAL_SYNC":
	available_drivers.append("local-sync")
	elif name == "IREE_HAL_DRIVER_LOCAL_TASK":
	available_drivers.append("local-task")
	elif name == "IREE_HAL_DRIVER_VULKAN":
	available_drivers.append("vulkan")
	elif name == "IREE_HAL_EXECUTABLE_LOADER_EMBEDDED_ELF":
	available_loaders.append("embedded-elf")
	elif name == "IREE_HAL_EXECUTABLE_LOADER_SYSTEM_LIBRARY":
	available_loaders.append("system-library")
	elif name == "IREE_HAL_EXECUTABLE_LOADER_VMVX_MODULE":
	available_loaders.append("vmvx-module")
	else:
	continue

	if self.verbose:
	available_drivers_str = ', '.join(available_drivers)
	print(f"Available drivers: {available_drivers_str}")
	available_loaders_str = ', '.join(available_loaders)
	print(f"Available loaders: {available_loaders_str}")

	return available_drivers, available_loaders