tools/boltifyer.py - toolchain/android_rust - Git at Google

 #!/usr/bin/env python3
 #
 # Copyright (C) 2022 The Android Open Source Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.

 import argparse
 from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
 import shutil
 import subprocess
 import sys
 from tempfile import TemporaryDirectory
 from typing import Optional, TextIO

 import context

 from android_rust import build_platform
 from android_rust.paths import *
 from android_rust.toolchains import CLANG_TOOLCHAIN_HOST, ClangToolchain, RustToolchain
 from android_rust.utils import (
     TERM_RED,
     ExtantPath,
     ResolvedPath,
     ScriptException,
     archive_create,
     archive_extract,
     extend_suffix,
     get_prebuilt_binary_paths,
     print_colored,
     print_fs_tree,
     is_archive,
     reify_singleton_patterns)

 from test_compiler import initialize_clang_prebuilt

 #
 # Constants
 #

 BOLT_INSTRUMENTATION_SUBJECTS: list[str] = [
     "lib/librustc_driver-*.so",
     "lib/libstd-*.so",
     "lib/libLLVM.so.*",
 ]

 BOLT_SKIP_FUNCS: dict[str, list[str]] = {
     "librustc_driver": ["blake3_hash_many_sse2/1",],
 }

 #
 # Program logic
 #


 def parse_args() -> argparse.Namespace:
     """Parses arguments and returns the parsed structure."""
     parser = argparse.ArgumentParser("Run BOLT on a Rust toolchain archive")
     parser.add_argument(
         "archive_path", type=ExtantPath, help="Path to the prebuilt archive to instrument/optimize")

     parser.add_argument(
         "--build-name",
         default="bolted",
         help="Desired filename for the output archive (w/o extension)")
     parser.add_argument(
         "--dist",
         "-d",
         dest="dist_path",
         type=ResolvedPath,
         default=DIST_PATH_DEFAULT,
         help="Where to place distributable artifacts")
     parser.add_argument(
         "--clang-prebuilt",
         type=ExtantPath,
         default=CLANG_TOOLCHAIN_HOST,
         help="Path to Clang toolchain that will be used to merge profile data")

     parser.add_argument(
         "--strip-only", action="store_true", help="Don't invoke BOLT, just strip the binaries")

     action_group = parser.add_mutually_exclusive_group()
     action_group.add_argument(
         "--profile-generate",
         type=ExtantPath,
         nargs="?",
         const=OUT_PATH_PROFILES,
         help="Instrument the toolchain for BOLT profile generation with an absolute path")
     action_group.add_argument(
         "--profile-generate-relative",
         type=Path,
         nargs="?",
         const=Path("out/profiles"),
         help="Instrument the toolchain for BOLT profile generation with a relative path")
     action_group.add_argument(
         "--profile-use",
         type=ExtantPath,
         nargs="?",
         const=OUT_PATH_PROFILES,
         help="Path to archive or directory with a bolt/ subdir containing BOLT profiles")

     return parser.parse_args()


 @contextmanager
 def handle_input_path(input_path: Path) -> Iterator[Path]:
     if is_archive(input_path):
         with TemporaryDirectory() as temp_dir:
             temp_dir_path = Path(temp_dir)
             print(f"Unpacking archive into {temp_dir}")
             archive_extract(input_path, temp_dir_path)
             yield temp_dir_path
     else:
         yield input_path


 @contextmanager
 def expand_profiles(profile_path: Path | None) -> Iterator[Path | None]:
     if profile_path is not None:
         if is_archive(profile_path):
             with TemporaryDirectory() as temp_dir:
                 temp_dir_path = Path(temp_dir)
                 print(f"Unpacking profiles into {temp_dir}")
                 archive_extract(profile_path, temp_dir_path)
                 yield temp_dir_path
         else:
             yield profile_path
     else:
         yield None


 def invoke_bolt(
         toolchain: ClangToolchain,
         obj_path: Path,
         bolt_log: TextIO,
         options: list[str] = []) -> None:
     bolt_path = toolchain.bolt()
     assert (bolt_path is not None)

     obj_path_bolted = extend_suffix(obj_path, ".bolt")

     bolt_log.write(f"BOLTing {str(obj_path)}\n")
     print(f"BOLTing {str(obj_path)}\n")
     bolt_log.flush()

     for (lib_prefix, func_names) in BOLT_SKIP_FUNCS.items():
         if obj_path.name.startswith(lib_prefix):
             options = options.copy()
             options.append(f"--skip-funcs={','.join(func_names)}")

     try:
         subprocess.run([bolt_path] + options + ["-o", obj_path_bolted, obj_path],
                        check=True,
                        stdout=bolt_log,
                        stderr=bolt_log)
     except subprocess.CalledProcessError as err:
         print(f"Failed to process object {str(obj_path)}")
         raise err

     bolt_log.write("\n")
     bolt_log.flush()

     obj_path.unlink()
     obj_path_bolted.rename(obj_path)


 def process_objects(
         toolchain: ClangToolchain,
         root: Path,
         strip_only: bool,
         profile_generate: Optional[Path],
         profile_use: Optional[Path],
         bolt_log: TextIO) -> None:

     print("Processing objects")

     instrumentation_subjects = reify_singleton_patterns(root, BOLT_INSTRUMENTATION_SUBJECTS)
     optimization_subjects = get_prebuilt_binary_paths(root, toplevel_only=True)

     for obj_path in get_prebuilt_binary_paths(root, subdirs=["bin", "lib"]):
         print(f"Boltifyer is looking at object {str(obj_path)}")

         if not strip_only and obj_path in optimization_subjects:
             print(f"Object {str(obj_path)} is an optimization subject")
             options = ["--peepholes=all"]
             if obj_path in instrumentation_subjects:
                 if profile_generate:
                     fdata_path = profile_generate / PROFILE_SUBDIR_BOLT / obj_path.name
                     options += [
                         "--instrument",
                         f"--instrumentation-file={fdata_path}",
                         "--instrumentation-file-append-pid",
                     ]

                 elif profile_use:
                     fdata_path = profile_use / PROFILE_SUBDIR_BOLT / (obj_path.name + ".fdata")

                     if fdata_path.exists():
                         options += [
                             f"--data={fdata_path}",
                             "--reorder-blocks=ext-tsp",
                             "--reorder-functions=hfsort",
                             "--split-functions",
                             "--split-all-cold",
                             "--split-eh",
                             "--dyno-stats",
                         ]
                     else:
                         print(f"Profile data missing for {obj_path.relative_to(root)}")

             invoke_bolt(toolchain, obj_path, bolt_log, options)

         else:
             print(f"Object {str(obj_path)} is NOT an optimization subject\n")
             toolchain.strip_symbols(obj_path)


 def boltify_toolchain(
         input_path: Path,
         dist_path: Path,
         build_name: str,
         strip_only: bool = False,
         profile_generate: Path | None = None,
         profile_use: Path | None = None,
         toolchain: ClangToolchain = CLANG_TOOLCHAIN_HOST) -> None:
     with (dist_path / BOLT_LOG_NAME).open("w") as bolt_log:
         with handle_input_path(input_path) as input_dir_path:
             print_fs_tree(input_dir_path)

             with expand_profiles(profile_use) as profile_use_path:
                 if profile_use_path is not None:
                     print_fs_tree(profile_use_path)

                     version_file_path = profile_use_path / PROFILE_SUBDIR_BOLT / PROFILED_VERSION_FILENAME
                     if version_file_path.exists():
                         with open(version_file_path) as version_file:
                             profiled_version = version_file.read()
                             input_toolchain = RustToolchain(
                                 input_dir_path, host=build_platform.system())
                             if input_toolchain.get_version() != profiled_version:
                                 raise ScriptException(
                                     "Profiled compiler and input compiler versions do not match")

                 try:
                     process_objects(
                         toolchain,
                         input_dir_path,
                         strip_only,
                         profile_generate,
                         profile_use_path,
                         bolt_log)
                     bolt_log.flush()
                 except subprocess.CalledProcessError as err:
                     print("Error BOLTing prebuilt objects")
                     raise err

                 print(f"Creating BOLTed archive")
                 try:
                     archive_create(dist_path / f"rust-{build_name}", input_dir_path, overwrite=True)
                 except subprocess.CalledProcessError as err:
                     print("Error creating the BOLTed archive")
                     raise err


 def main() -> None:
     if not build_platform.is_linux():
         raise ScriptException("BOLT tools are only available on Linux hosts")

     args = parse_args()

     initialize_clang_prebuilt(args)

     boltify_toolchain(
         args.archive_path,
         args.dist_path,
         args.build_name,
         args.strip_only,
         args.profile_generate or args.profile_generate_relative,
         args.profile_use)

     if args.profile_use is not None and args.profile_use.name.endswith(
             ".tar.xz") and args.profile_use.parent != args.dist_path:
         shutil.copy2(args.profile_use, args.dist_path)


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

     except (ScriptException, argparse.ArgumentTypeError) as err:
         print_colored(str(err), TERM_RED)
         sys.exit(1)
	#!/usr/bin/env python3
	#
	# Copyright (C) 2022 The Android Open Source Project
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	import argparse
	from collections.abc import Iterator
	from contextlib import contextmanager
	from pathlib import Path
	import shutil
	import subprocess
	import sys
	from tempfile import TemporaryDirectory
	from typing import Optional, TextIO

	import context

	from android_rust import build_platform
	from android_rust.paths import *
	from android_rust.toolchains import CLANG_TOOLCHAIN_HOST, ClangToolchain, RustToolchain
	from android_rust.utils import (
	TERM_RED,
	ExtantPath,
	ResolvedPath,
	ScriptException,
	archive_create,
	archive_extract,
	extend_suffix,
	get_prebuilt_binary_paths,
	print_colored,
	print_fs_tree,
	is_archive,
	reify_singleton_patterns)

	from test_compiler import initialize_clang_prebuilt

	#
	# Constants
	#

	BOLT_INSTRUMENTATION_SUBJECTS: list[str] = [
	"lib/librustc_driver-*.so",
	"lib/libstd-*.so",
	"lib/libLLVM.so.*",
	]

	BOLT_SKIP_FUNCS: dict[str, list[str]] = {
	"librustc_driver": ["blake3_hash_many_sse2/1",],
	}

	#
	# Program logic
	#


	def parse_args() -> argparse.Namespace:
	"""Parses arguments and returns the parsed structure."""
	parser = argparse.ArgumentParser("Run BOLT on a Rust toolchain archive")
	parser.add_argument(
	"archive_path", type=ExtantPath, help="Path to the prebuilt archive to instrument/optimize")

	parser.add_argument(
	"--build-name",
	default="bolted",
	help="Desired filename for the output archive (w/o extension)")
	parser.add_argument(
	"--dist",
	"-d",
	dest="dist_path",
	type=ResolvedPath,
	default=DIST_PATH_DEFAULT,
	help="Where to place distributable artifacts")
	parser.add_argument(
	"--clang-prebuilt",
	type=ExtantPath,
	default=CLANG_TOOLCHAIN_HOST,
	help="Path to Clang toolchain that will be used to merge profile data")

	parser.add_argument(
	"--strip-only", action="store_true", help="Don't invoke BOLT, just strip the binaries")

	action_group = parser.add_mutually_exclusive_group()
	action_group.add_argument(
	"--profile-generate",
	type=ExtantPath,
	nargs="?",
	const=OUT_PATH_PROFILES,
	help="Instrument the toolchain for BOLT profile generation with an absolute path")
	action_group.add_argument(
	"--profile-generate-relative",
	type=Path,
	nargs="?",
	const=Path("out/profiles"),
	help="Instrument the toolchain for BOLT profile generation with a relative path")
	action_group.add_argument(
	"--profile-use",
	type=ExtantPath,
	nargs="?",
	const=OUT_PATH_PROFILES,
	help="Path to archive or directory with a bolt/ subdir containing BOLT profiles")

	return parser.parse_args()


	@contextmanager
	def handle_input_path(input_path: Path) -> Iterator[Path]:
	if is_archive(input_path):
	with TemporaryDirectory() as temp_dir:
	temp_dir_path = Path(temp_dir)
	print(f"Unpacking archive into {temp_dir}")
	archive_extract(input_path, temp_dir_path)
	yield temp_dir_path
	else:
	yield input_path


	@contextmanager
	def expand_profiles(profile_path: Path \| None) -> Iterator[Path \| None]:
	if profile_path is not None:
	if is_archive(profile_path):
	with TemporaryDirectory() as temp_dir:
	temp_dir_path = Path(temp_dir)
	print(f"Unpacking profiles into {temp_dir}")
	archive_extract(profile_path, temp_dir_path)
	yield temp_dir_path
	else:
	yield profile_path
	else:
	yield None


	def invoke_bolt(
	toolchain: ClangToolchain,
	obj_path: Path,
	bolt_log: TextIO,
	options: list[str] = []) -> None:
	bolt_path = toolchain.bolt()
	assert (bolt_path is not None)

	obj_path_bolted = extend_suffix(obj_path, ".bolt")

	bolt_log.write(f"BOLTing {str(obj_path)}\n")
	print(f"BOLTing {str(obj_path)}\n")
	bolt_log.flush()

	for (lib_prefix, func_names) in BOLT_SKIP_FUNCS.items():
	if obj_path.name.startswith(lib_prefix):
	options = options.copy()
	options.append(f"--skip-funcs={','.join(func_names)}")

	try:
	subprocess.run([bolt_path] + options + ["-o", obj_path_bolted, obj_path],
	check=True,
	stdout=bolt_log,
	stderr=bolt_log)
	except subprocess.CalledProcessError as err:
	print(f"Failed to process object {str(obj_path)}")
	raise err

	bolt_log.write("\n")
	bolt_log.flush()

	obj_path.unlink()
	obj_path_bolted.rename(obj_path)


	def process_objects(
	toolchain: ClangToolchain,
	root: Path,
	strip_only: bool,
	profile_generate: Optional[Path],
	profile_use: Optional[Path],
	bolt_log: TextIO) -> None:

	print("Processing objects")

	instrumentation_subjects = reify_singleton_patterns(root, BOLT_INSTRUMENTATION_SUBJECTS)
	optimization_subjects = get_prebuilt_binary_paths(root, toplevel_only=True)

	for obj_path in get_prebuilt_binary_paths(root, subdirs=["bin", "lib"]):
	print(f"Boltifyer is looking at object {str(obj_path)}")

	if not strip_only and obj_path in optimization_subjects:
	print(f"Object {str(obj_path)} is an optimization subject")
	options = ["--peepholes=all"]
	if obj_path in instrumentation_subjects:
	if profile_generate:
	fdata_path = profile_generate / PROFILE_SUBDIR_BOLT / obj_path.name
	options += [
	"--instrument",
	f"--instrumentation-file={fdata_path}",
	"--instrumentation-file-append-pid",
	]

	elif profile_use:
	fdata_path = profile_use / PROFILE_SUBDIR_BOLT / (obj_path.name + ".fdata")

	if fdata_path.exists():
	options += [
	f"--data={fdata_path}",
	"--reorder-blocks=ext-tsp",
	"--reorder-functions=hfsort",
	"--split-functions",
	"--split-all-cold",
	"--split-eh",
	"--dyno-stats",
	]
	else:
	print(f"Profile data missing for {obj_path.relative_to(root)}")

	invoke_bolt(toolchain, obj_path, bolt_log, options)

	else:
	print(f"Object {str(obj_path)} is NOT an optimization subject\n")
	toolchain.strip_symbols(obj_path)


	def boltify_toolchain(
	input_path: Path,
	dist_path: Path,
	build_name: str,
	strip_only: bool = False,
	profile_generate: Path \| None = None,
	profile_use: Path \| None = None,
	toolchain: ClangToolchain = CLANG_TOOLCHAIN_HOST) -> None:
	with (dist_path / BOLT_LOG_NAME).open("w") as bolt_log:
	with handle_input_path(input_path) as input_dir_path:
	print_fs_tree(input_dir_path)

	with expand_profiles(profile_use) as profile_use_path:
	if profile_use_path is not None:
	print_fs_tree(profile_use_path)

	version_file_path = profile_use_path / PROFILE_SUBDIR_BOLT / PROFILED_VERSION_FILENAME
	if version_file_path.exists():
	with open(version_file_path) as version_file:
	profiled_version = version_file.read()
	input_toolchain = RustToolchain(
	input_dir_path, host=build_platform.system())
	if input_toolchain.get_version() != profiled_version:
	raise ScriptException(
	"Profiled compiler and input compiler versions do not match")

	try:
	process_objects(
	toolchain,
	input_dir_path,
	strip_only,
	profile_generate,
	profile_use_path,
	bolt_log)
	bolt_log.flush()
	except subprocess.CalledProcessError as err:
	print("Error BOLTing prebuilt objects")
	raise err

	print(f"Creating BOLTed archive")
	try:
	archive_create(dist_path / f"rust-{build_name}", input_dir_path, overwrite=True)
	except subprocess.CalledProcessError as err:
	print("Error creating the BOLTed archive")
	raise err


	def main() -> None:
	if not build_platform.is_linux():
	raise ScriptException("BOLT tools are only available on Linux hosts")

	args = parse_args()

	initialize_clang_prebuilt(args)

	boltify_toolchain(
	args.archive_path,
	args.dist_path,
	args.build_name,
	args.strip_only,
	args.profile_generate or args.profile_generate_relative,
	args.profile_use)

	if args.profile_use is not None and args.profile_use.name.endswith(
	".tar.xz") and args.profile_use.parent != args.dist_path:
	shutil.copy2(args.profile_use, args.dist_path)


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

	except (ScriptException, argparse.ArgumentTypeError) as err:
	print_colored(str(err), TERM_RED)
	sys.exit(1)