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.

 See docs/developers/developing_iree/benchmark_suites.md for how to build the
 benchmark suites.
 """

 import pathlib
 import sys

 # Add build_tools python dir to the search path.
 sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

 import argparse
 import json
 import os
 import re
 import zipfile

 from dataclasses import asdict, dataclass
 from typing import BinaryIO, Dict, List, Optional, TextIO

 from common import benchmark_definition
 from common.benchmark_definition import CompilationInfo, CompilationResults, CompilationStatistics, ModuleComponentSizes, get_git_commit_hash
 from common import benchmark_config
 from e2e_test_artifacts import iree_artifacts
 from e2e_test_framework import serialization
 from e2e_test_framework.definitions import iree_definitions

 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"
 E2E_TEST_ARTIFACTS_REL_PATH = "e2e_test_artifacts"

 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\.fb",
 ]


 @dataclass(frozen=True)
 class ModuleInfo(object):
   module_path: pathlib.Path
   stream_stats_path: pathlib.Path


 def match_module_cmake_target(module_path: pathlib.PurePath) -> Optional[str]:
   if module_path.match(f"{E2E_TEST_ARTIFACTS_REL_PATH}/iree_*/"
                        f"{iree_artifacts.MODULE_FILENAME}"):
     # <e2e test artifacts dir>/iree_<module dir>/<module filename>
     path_parts = module_path.parts[-3:]
     # Join to get the CMake target name. This is *not* a filesystem path, so we
     # don't want \ separators on Windows that we would get with os.path.join().
     return '/'.join(path_parts)

   return None


 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(pathlib.PurePath(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())

   identified_names = set()
   if VM_COMPONENT_NAME in size_map:
     vm_component_bytes = size_map[VM_COMPONENT_NAME]
     identified_names.add(VM_COMPONENT_NAME)
   else:
     vm_component_bytes = 0

   if CONST_COMPONENT_NAME in size_map:
     const_component_bytes = size_map[CONST_COMPONENT_NAME]
     identified_names.add(CONST_COMPONENT_NAME)
   else:
     const_component_bytes = 0

   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_map_from_compilation_benchmark_config(
     compilation_benchmark_config_data: TextIO,
     e2e_test_artifacts_dir: pathlib.PurePath
 ) -> Dict[CompilationInfo, ModuleInfo]:
   benchmark_config = json.load(compilation_benchmark_config_data)
   gen_configs = serialization.unpack_and_deserialize(
       data=benchmark_config["generation_configs"],
       root_type=List[iree_definitions.ModuleGenerationConfig])
   module_map = {}
   for gen_config in gen_configs:
     model = gen_config.imported_model.model
     compile_config = gen_config.compile_config
     target_archs = []
     for compile_target in compile_config.compile_targets:
       arch = compile_target.target_architecture
       target_archs.append(
           (f"{arch.type.value}-{arch.architecture}-{arch.microarchitecture}-"
            f"{compile_target.target_abi.value}"))
     compilation_info = CompilationInfo(
         name=gen_config.name,
         model_name=model.name,
         model_tags=tuple(model.tags),
         model_source=model.source_type.value,
         target_arch=f"[{','.join(target_archs)}]",
         compile_tags=tuple(compile_config.tags),
         gen_config_id=gen_config.composite_id)
     module_dir_path = pathlib.Path(
         iree_artifacts.get_module_dir_path(module_generation_config=gen_config,
                                            root_path=e2e_test_artifacts_dir))
     module_path = module_dir_path / iree_artifacts.MODULE_FILENAME
     stream_stats_path = (module_dir_path /
                          iree_artifacts.SCHEDULING_STATS_FILENAME)
     module_map[compilation_info] = ModuleInfo(
         module_path=module_path, stream_stats_path=stream_stats_path)

   return module_map


 def _check_dir_path(path_str: str) -> pathlib.Path:
   path = pathlib.Path(path_str)
   if not path.is_dir():
     raise argparse.ArgumentTypeError(f"{path} is not a directory.")
   return path


 def _check_file_path(path_str: str) -> pathlib.Path:
   path = pathlib.Path(path_str)
   if not path.is_file():
     raise argparse.ArgumentTypeError(f"{path} is not a file.")
   return path


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

   parser = argparse.ArgumentParser(
       description="Collect compilation statistics from benchmark suites.")
   parser.add_argument(
       "--compilation_benchmark_config",
       type=_check_file_path,
       required=True,
       help="Exported compilation benchmark config of e2e test artifacts.")
   parser.add_argument("--build_log",
                       type=_check_file_path,
                       required=True,
                       help="Path to the ninja build log.")
   parser.add_argument("--e2e_test_artifacts_dir",
                       type=_check_dir_path,
                       required=True,
                       help="Path to the e2e test artifacts directory.")
   parser.add_argument("--output",
                       type=pathlib.Path,
                       help="Path to output JSON file.")

   return parser.parse_args()


 def main(args: argparse.Namespace):
   config_data = args.compilation_benchmark_config.open("r")
   module_map = get_module_map_from_compilation_benchmark_config(
       compilation_benchmark_config_data=config_data,
       e2e_test_artifacts_dir=args.e2e_test_artifacts_dir)
   build_log_path = args.build_log

   with build_log_path.open("r") as log_file:
     target_build_time_map = parse_compilation_time_from_ninja_log(log_file)

   compilation_statistics_list = []
   for compilation_info, module_info in module_map.items():
     module_path = module_info.module_path
     with module_path.open("rb") as module_file:
       module_component_sizes = get_module_component_info(
           module_file,
           module_path.stat().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]

     stream_stats_json = json.loads(module_info.stream_stats_path.read_text())
     exec_stats_json = stream_stats_json["stream-aggregate"]["execution"]
     ir_stats = benchmark_definition.IRStatistics(
         stream_dispatch_count=exec_stats_json["dispatch-count"])

     compilation_statistics = CompilationStatistics(
         compilation_info=compilation_info,
         module_component_sizes=module_component_sizes,
         compilation_time_ms=compilation_time_ms,
         ir_stats=ir_stats)
     compilation_statistics_list.append(compilation_statistics)

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

   json_output = json.dumps(asdict(compilation_results), indent=2)
   if args.output is None:
     print(json_output)
   else:
     args.output.write_text(json_output)


 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.

	See docs/developers/developing_iree/benchmark_suites.md for how to build the
	benchmark suites.
	"""

	import pathlib
	import sys

	# Add build_tools python dir to the search path.
	sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

	import argparse
	import json
	import os
	import re
	import zipfile

	from dataclasses import asdict, dataclass
	from typing import BinaryIO, Dict, List, Optional, TextIO

	from common import benchmark_definition
	from common.benchmark_definition import CompilationInfo, CompilationResults, CompilationStatistics, ModuleComponentSizes, get_git_commit_hash
	from common import benchmark_config
	from e2e_test_artifacts import iree_artifacts
	from e2e_test_framework import serialization
	from e2e_test_framework.definitions import iree_definitions

	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"
	E2E_TEST_ARTIFACTS_REL_PATH = "e2e_test_artifacts"

	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\.fb",
	]


	@dataclass(frozen=True)
	class ModuleInfo(object):
	module_path: pathlib.Path
	stream_stats_path: pathlib.Path


	def match_module_cmake_target(module_path: pathlib.PurePath) -> Optional[str]:
	if module_path.match(f"{E2E_TEST_ARTIFACTS_REL_PATH}/iree_*/"
	f"{iree_artifacts.MODULE_FILENAME}"):
	# <e2e test artifacts dir>/iree_<module dir>/<module filename>
	path_parts = module_path.parts[-3:]
	# Join to get the CMake target name. This is not a filesystem path, so we
	# don't want \ separators on Windows that we would get with os.path.join().
	return '/'.join(path_parts)

	return None


	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(pathlib.PurePath(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())

	identified_names = set()
	if VM_COMPONENT_NAME in size_map:
	vm_component_bytes = size_map[VM_COMPONENT_NAME]
	identified_names.add(VM_COMPONENT_NAME)
	else:
	vm_component_bytes = 0

	if CONST_COMPONENT_NAME in size_map:
	const_component_bytes = size_map[CONST_COMPONENT_NAME]
	identified_names.add(CONST_COMPONENT_NAME)
	else:
	const_component_bytes = 0

	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_map_from_compilation_benchmark_config(
	compilation_benchmark_config_data: TextIO,
	e2e_test_artifacts_dir: pathlib.PurePath
	) -> Dict[CompilationInfo, ModuleInfo]:
	benchmark_config = json.load(compilation_benchmark_config_data)
	gen_configs = serialization.unpack_and_deserialize(
	data=benchmark_config["generation_configs"],
	root_type=List[iree_definitions.ModuleGenerationConfig])
	module_map = {}
	for gen_config in gen_configs:
	model = gen_config.imported_model.model
	compile_config = gen_config.compile_config
	target_archs = []
	for compile_target in compile_config.compile_targets:
	arch = compile_target.target_architecture
	target_archs.append(
	(f"{arch.type.value}-{arch.architecture}-{arch.microarchitecture}-"
	f"{compile_target.target_abi.value}"))
	compilation_info = CompilationInfo(
	name=gen_config.name,
	model_name=model.name,
	model_tags=tuple(model.tags),
	model_source=model.source_type.value,
	target_arch=f"[{','.join(target_archs)}]",
	compile_tags=tuple(compile_config.tags),
	gen_config_id=gen_config.composite_id)
	module_dir_path = pathlib.Path(
	iree_artifacts.get_module_dir_path(module_generation_config=gen_config,
	root_path=e2e_test_artifacts_dir))
	module_path = module_dir_path / iree_artifacts.MODULE_FILENAME
	stream_stats_path = (module_dir_path /
	iree_artifacts.SCHEDULING_STATS_FILENAME)
	module_map[compilation_info] = ModuleInfo(
	module_path=module_path, stream_stats_path=stream_stats_path)

	return module_map


	def _check_dir_path(path_str: str) -> pathlib.Path:
	path = pathlib.Path(path_str)
	if not path.is_dir():
	raise argparse.ArgumentTypeError(f"{path} is not a directory.")
	return path


	def _check_file_path(path_str: str) -> pathlib.Path:
	path = pathlib.Path(path_str)
	if not path.is_file():
	raise argparse.ArgumentTypeError(f"{path} is not a file.")
	return path


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

	parser = argparse.ArgumentParser(
	description="Collect compilation statistics from benchmark suites.")
	parser.add_argument(
	"--compilation_benchmark_config",
	type=_check_file_path,
	required=True,
	help="Exported compilation benchmark config of e2e test artifacts.")
	parser.add_argument("--build_log",
	type=_check_file_path,
	required=True,
	help="Path to the ninja build log.")
	parser.add_argument("--e2e_test_artifacts_dir",
	type=_check_dir_path,
	required=True,
	help="Path to the e2e test artifacts directory.")
	parser.add_argument("--output",
	type=pathlib.Path,
	help="Path to output JSON file.")

	return parser.parse_args()


	def main(args: argparse.Namespace):
	config_data = args.compilation_benchmark_config.open("r")
	module_map = get_module_map_from_compilation_benchmark_config(
	compilation_benchmark_config_data=config_data,
	e2e_test_artifacts_dir=args.e2e_test_artifacts_dir)
	build_log_path = args.build_log

	with build_log_path.open("r") as log_file:
	target_build_time_map = parse_compilation_time_from_ninja_log(log_file)

	compilation_statistics_list = []
	for compilation_info, module_info in module_map.items():
	module_path = module_info.module_path
	with module_path.open("rb") as module_file:
	module_component_sizes = get_module_component_info(
	module_file,
	module_path.stat().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]

	stream_stats_json = json.loads(module_info.stream_stats_path.read_text())
	exec_stats_json = stream_stats_json["stream-aggregate"]["execution"]
	ir_stats = benchmark_definition.IRStatistics(
	stream_dispatch_count=exec_stats_json["dispatch-count"])

	compilation_statistics = CompilationStatistics(
	compilation_info=compilation_info,
	module_component_sizes=module_component_sizes,
	compilation_time_ms=compilation_time_ms,
	ir_stats=ir_stats)
	compilation_statistics_list.append(compilation_statistics)

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

	json_output = json.dumps(asdict(compilation_results), indent=2)
	if args.output is None:
	print(json_output)
	else:
	args.output.write_text(json_output)


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