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

 #!/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
 """Collect compilation statistics from benchmark suites.

 The benchmark suites need to be built with ninja and enable the CMake option
 IREE_ENABLE_COMPILATION_BENCHMARKS.
 """

 import argparse
 import json
 import os
 import re
 import zipfile

 from dataclasses import asdict
 from typing import BinaryIO, Dict, Optional, TextIO
 from pathlib import PurePath

 from common.benchmark_definition import CompilationInfo, CompilationResults, CompilationStatistics, ModuleComponentSizes, get_git_commit_hash
 from common.benchmark_suite import BENCHMARK_SUITE_REL_PATH, BenchmarkSuite

 BENCHMARK_FLAGFILE = "flagfile"
 MODULE_DIR = "vmfb"
 MODULE_FILE_EXTENSION = ".vmfb"
 NINJA_LOG_HEADER = "ninja log v5"
 NINJA_BUILD_LOG = ".ninja_log"
 COMPILATION_STATS_MODULE_SUFFIX = "compile-stats"

 VM_COMPONENT_NAME = "module.fb"
 CONST_COMPONENT_NAME = "_const.bin"
 DISPATCH_COMPONENT_PATTERNS = [
     r".+_embedded_elf_.+\.so",
     r".+_vulkan_spirv_fb\.fb",
     r".+_cuda_nvptx_fb\.fb",
     r".+_vmvx_bytecode_fb\.bin",
 ]


 def match_module_cmake_target(module_path: str) -> Optional[str]:
   # Get the last 4 parts of module path. They are expected to be:
   # benchmark_suites/<category>/vmfb/<module filename>.vmfb
   path_parts = PurePath(module_path).parts[-4:]
   if len(path_parts) < 4:
     return None
   if path_parts[0] != BENCHMARK_SUITE_REL_PATH:
     return None
   if path_parts[2] != MODULE_DIR:
     return None
   if os.path.splitext(path_parts[3])[1] != MODULE_FILE_EXTENSION:
     return None
   return os.path.join(*path_parts)


 def parse_compilation_time_from_ninja_log(log: TextIO) -> Dict[str, int]:
   """Retrieve the compilation time (ms) from the Ninja build log.

   Returns:
     Map of target name and compilation time in ms.
   """

   target_build_time_map = {}
   header = log.readline()
   if NINJA_LOG_HEADER not in header:
     raise NotImplementedError(f"Unsupported ninja log version: {header}")

   for line in log:
     start_time, end_time, _, target, _ = line.strip().split("\t")
     cmake_target = match_module_cmake_target(target)
     if cmake_target is None:
       continue

     start_time = int(start_time)
     end_time = int(end_time)
     target_build_time_map[cmake_target] = end_time - start_time

   return target_build_time_map


 def get_module_component_info(module: BinaryIO,
                               module_file_bytes: int) -> ModuleComponentSizes:
   with zipfile.ZipFile(module) as module_zipfile:
     size_map = dict(
         (info.filename, info.file_size) for info in module_zipfile.infolist())

   vm_component_bytes = size_map[VM_COMPONENT_NAME]
   const_component_bytes = size_map[CONST_COMPONENT_NAME]
   identified_names = {VM_COMPONENT_NAME, CONST_COMPONENT_NAME}
   total_dispatch_component_bytes = 0
   for filename, size in size_map.items():
     for pattern in DISPATCH_COMPONENT_PATTERNS:
       if re.match(pattern, filename):
         total_dispatch_component_bytes += size
         identified_names.add(filename)
         break

   if identified_names != set(size_map.keys()):
     raise RuntimeError(
         f"Unrecognized components in the module: {size_map.keys()}.")

   return ModuleComponentSizes(
       file_bytes=module_file_bytes,
       vm_component_bytes=vm_component_bytes,
       const_component_bytes=const_component_bytes,
       total_dispatch_component_bytes=total_dispatch_component_bytes)


 def get_module_path(flag_file: TextIO) -> Optional[str]:
   """Retrieve the module path for compilation statistics from the flag file."""

   module_path = None
   for line in flag_file:
     match = re.match("--module_file=(.+)", line.strip())
     if match:
       module_name, module_ext = os.path.splitext(match.group(1))
       module_path = f"{module_name}-{COMPILATION_STATS_MODULE_SUFFIX}{module_ext}"
       break

   return module_path


 def parse_arguments():
   """Returns an argument parser with common options."""

   def check_dir_path(path):
     if os.path.isdir(path):
       return path
     else:
       raise argparse.ArgumentTypeError(path)

   parser = argparse.ArgumentParser()
   parser.add_argument("--output",
                       required=True,
                       help="Path to output JSON file.")
   parser.add_argument(
       "build_dir",
       metavar="<build-dir>",
       type=check_dir_path,
       help="Path to the build directory containing benchmark suites.")
   parser.add_argument("--verbose",
                       action="store_true",
                       help="Print internal information during execution.")

   return parser.parse_args()


 def main(args: argparse.Namespace):
   benchmark_suite_dir = os.path.join(args.build_dir, BENCHMARK_SUITE_REL_PATH)
   benchmark_suite = BenchmarkSuite.load_from_benchmark_suite_dir(
       benchmark_suite_dir)

   with open(os.path.join(args.build_dir, NINJA_BUILD_LOG), "r") as log_file:
     target_build_time_map = parse_compilation_time_from_ninja_log(log_file)

   compilation_statistics_list = []
   for category, _ in benchmark_suite.list_categories():
     benchmark_cases = benchmark_suite.filter_benchmarks_for_category(
         category=category)
     for benchmark_case in benchmark_cases:
       flag_file_path = os.path.join(benchmark_case.benchmark_case_dir,
                                     BENCHMARK_FLAGFILE)
       with open(flag_file_path, "r") as flag_file:
         module_path = get_module_path(flag_file)

       if module_path is None:
         raise RuntimeError(
             f"Can't find the module file in the flagfile: {flag_file_path}")
       module_path = os.path.abspath(
           os.path.join(benchmark_case.benchmark_case_dir, module_path))

       with open(module_path, "rb") as module_file:
         module_component_sizes = get_module_component_info(
             module_file,
             os.stat(module_path).st_size)

       cmake_target = match_module_cmake_target(module_path)
       if cmake_target is None:
         raise RuntimeError(
             f"Module path isn't a module cmake target: {module_path}")
       compilation_time_ms = target_build_time_map[cmake_target]

       compilation_info = CompilationInfo(model_name=benchmark_case.model_name,
                                          model_tags=benchmark_case.model_tags,
                                          model_source=category,
                                          target_arch=benchmark_case.target_arch,
                                          bench_mode=benchmark_case.bench_mode)
       compilation_statistics = CompilationStatistics(
           compilation_info=compilation_info,
           module_component_sizes=module_component_sizes,
           compilation_time_ms=compilation_time_ms)
       compilation_statistics_list.append(compilation_statistics)

   commit = get_git_commit_hash("HEAD")
   compilation_results = CompilationResults(
       commit=commit, compilation_statistics=compilation_statistics_list)

   json_object = asdict(compilation_results)
   with open(args.output, "w") as f:
     json.dump(json_object, f)

   if args.verbose:
     print(json.dumps(json_object, indent=4))


 if __name__ == "__main__":
   main(parse_arguments())
	#!/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
	"""Collect compilation statistics from benchmark suites.

	The benchmark suites need to be built with ninja and enable the CMake option
	IREE_ENABLE_COMPILATION_BENCHMARKS.
	"""

	import argparse
	import json
	import os
	import re
	import zipfile

	from dataclasses import asdict
	from typing import BinaryIO, Dict, Optional, TextIO
	from pathlib import PurePath

	from common.benchmark_definition import CompilationInfo, CompilationResults, CompilationStatistics, ModuleComponentSizes, get_git_commit_hash
	from common.benchmark_suite import BENCHMARK_SUITE_REL_PATH, BenchmarkSuite

	BENCHMARK_FLAGFILE = "flagfile"
	MODULE_DIR = "vmfb"
	MODULE_FILE_EXTENSION = ".vmfb"
	NINJA_LOG_HEADER = "ninja log v5"
	NINJA_BUILD_LOG = ".ninja_log"
	COMPILATION_STATS_MODULE_SUFFIX = "compile-stats"

	VM_COMPONENT_NAME = "module.fb"
	CONST_COMPONENT_NAME = "_const.bin"
	DISPATCH_COMPONENT_PATTERNS = [
	r".+_embedded_elf_.+\.so",
	r".+_vulkan_spirv_fb\.fb",
	r".+_cuda_nvptx_fb\.fb",
	r".+_vmvx_bytecode_fb\.bin",
	]


	def match_module_cmake_target(module_path: str) -> Optional[str]:
	# Get the last 4 parts of module path. They are expected to be:
	# benchmark_suites/<category>/vmfb/<module filename>.vmfb
	path_parts = PurePath(module_path).parts[-4:]
	if len(path_parts) < 4:
	return None
	if path_parts[0] != BENCHMARK_SUITE_REL_PATH:
	return None
	if path_parts[2] != MODULE_DIR:
	return None
	if os.path.splitext(path_parts[3])[1] != MODULE_FILE_EXTENSION:
	return None
	return os.path.join(*path_parts)


	def parse_compilation_time_from_ninja_log(log: TextIO) -> Dict[str, int]:
	"""Retrieve the compilation time (ms) from the Ninja build log.

	Returns:
	Map of target name and compilation time in ms.
	"""

	target_build_time_map = {}
	header = log.readline()
	if NINJA_LOG_HEADER not in header:
	raise NotImplementedError(f"Unsupported ninja log version: {header}")

	for line in log:
	start_time, end_time, _, target, _ = line.strip().split("\t")
	cmake_target = match_module_cmake_target(target)
	if cmake_target is None:
	continue

	start_time = int(start_time)
	end_time = int(end_time)
	target_build_time_map[cmake_target] = end_time - start_time

	return target_build_time_map


	def get_module_component_info(module: BinaryIO,
	module_file_bytes: int) -> ModuleComponentSizes:
	with zipfile.ZipFile(module) as module_zipfile:
	size_map = dict(
	(info.filename, info.file_size) for info in module_zipfile.infolist())

	vm_component_bytes = size_map[VM_COMPONENT_NAME]
	const_component_bytes = size_map[CONST_COMPONENT_NAME]
	identified_names = {VM_COMPONENT_NAME, CONST_COMPONENT_NAME}
	total_dispatch_component_bytes = 0
	for filename, size in size_map.items():
	for pattern in DISPATCH_COMPONENT_PATTERNS:
	if re.match(pattern, filename):
	total_dispatch_component_bytes += size
	identified_names.add(filename)
	break

	if identified_names != set(size_map.keys()):
	raise RuntimeError(
	f"Unrecognized components in the module: {size_map.keys()}.")

	return ModuleComponentSizes(
	file_bytes=module_file_bytes,
	vm_component_bytes=vm_component_bytes,
	const_component_bytes=const_component_bytes,
	total_dispatch_component_bytes=total_dispatch_component_bytes)


	def get_module_path(flag_file: TextIO) -> Optional[str]:
	"""Retrieve the module path for compilation statistics from the flag file."""

	module_path = None
	for line in flag_file:
	match = re.match("--module_file=(.+)", line.strip())
	if match:
	module_name, module_ext = os.path.splitext(match.group(1))
	module_path = f"{module_name}-{COMPILATION_STATS_MODULE_SUFFIX}{module_ext}"
	break

	return module_path


	def parse_arguments():
	"""Returns an argument parser with common options."""

	def check_dir_path(path):
	if os.path.isdir(path):
	return path
	else:
	raise argparse.ArgumentTypeError(path)

	parser = argparse.ArgumentParser()
	parser.add_argument("--output",
	required=True,
	help="Path to output JSON file.")
	parser.add_argument(
	"build_dir",
	metavar="<build-dir>",
	type=check_dir_path,
	help="Path to the build directory containing benchmark suites.")
	parser.add_argument("--verbose",
	action="store_true",
	help="Print internal information during execution.")

	return parser.parse_args()


	def main(args: argparse.Namespace):
	benchmark_suite_dir = os.path.join(args.build_dir, BENCHMARK_SUITE_REL_PATH)
	benchmark_suite = BenchmarkSuite.load_from_benchmark_suite_dir(
	benchmark_suite_dir)

	with open(os.path.join(args.build_dir, NINJA_BUILD_LOG), "r") as log_file:
	target_build_time_map = parse_compilation_time_from_ninja_log(log_file)

	compilation_statistics_list = []
	for category, _ in benchmark_suite.list_categories():
	benchmark_cases = benchmark_suite.filter_benchmarks_for_category(
	category=category)
	for benchmark_case in benchmark_cases:
	flag_file_path = os.path.join(benchmark_case.benchmark_case_dir,
	BENCHMARK_FLAGFILE)
	with open(flag_file_path, "r") as flag_file:
	module_path = get_module_path(flag_file)

	if module_path is None:
	raise RuntimeError(
	f"Can't find the module file in the flagfile: {flag_file_path}")
	module_path = os.path.abspath(
	os.path.join(benchmark_case.benchmark_case_dir, module_path))

	with open(module_path, "rb") as module_file:
	module_component_sizes = get_module_component_info(
	module_file,
	os.stat(module_path).st_size)

	cmake_target = match_module_cmake_target(module_path)
	if cmake_target is None:
	raise RuntimeError(
	f"Module path isn't a module cmake target: {module_path}")
	compilation_time_ms = target_build_time_map[cmake_target]

	compilation_info = CompilationInfo(model_name=benchmark_case.model_name,
	model_tags=benchmark_case.model_tags,
	model_source=category,
	target_arch=benchmark_case.target_arch,
	bench_mode=benchmark_case.bench_mode)
	compilation_statistics = CompilationStatistics(
	compilation_info=compilation_info,
	module_component_sizes=module_component_sizes,
	compilation_time_ms=compilation_time_ms)
	compilation_statistics_list.append(compilation_statistics)

	commit = get_git_commit_hash("HEAD")
	compilation_results = CompilationResults(
	commit=commit, compilation_statistics=compilation_statistics_list)

	json_object = asdict(compilation_results)
	with open(args.output, "w") as f:
	json.dump(json_object, f)

	if args.verbose:
	print(json.dumps(json_object, indent=4))


	if __name__ == "__main__":
	main(parse_arguments())