| #!/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 pathlib import Path |
| import shutil |
| import subprocess |
| from tempfile import TemporaryDirectory |
| from typing import Optional, TextIO |
| |
| from paths import * |
| from utils import ( |
| ResolvedPath, |
| extend_suffix, |
| get_prebuilt_binary_paths, |
| reify_singleton_patterns, |
| run_and_exit_on_failure, |
| strip_symbols) |
| |
| # |
| # Constants |
| # |
| |
| BOLT_INSTRUMENTATION_SUBJECTS: list[str] = [ |
| "lib/librustc_driver-*.so", |
| "lib/libstd-*.so", |
| "lib/libLLVM-*.so", |
| ] |
| |
| # |
| # 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=ResolvedPath, |
| help="Release name for the dist result") |
| |
| parser.add_argument( |
| "outname", |
| 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") |
| |
| action_group = parser.add_mutually_exclusive_group() |
| action_group.add_argument( |
| "--profile-generate", type=ResolvedPath, nargs="?", const=OUT_PATH_PROFILES, |
| help="Instrument the toolchain for BOLT profile generation") |
| action_group.add_argument( |
| "--profile-use", type=ResolvedPath, nargs="?", const=OUT_PATH_PROFILES, |
| help="Use the provided profile when optimizing the toolchain") |
| |
| return parser.parse_args() |
| |
| |
| def unpack_archive(archive_path: Path) -> TemporaryDirectory[str]: |
| tmpdir = TemporaryDirectory() |
| print(f"Unpacking archive into {tmpdir.name}") |
| shutil.unpack_archive(archive_path.as_posix(), tmpdir.name) |
| return tmpdir |
| |
| |
| def invoke_bolt(obj_path: Path, bolt_log: TextIO, options: str = "") -> None: |
| obj_path_bolted = extend_suffix(obj_path, ".bolt") |
| |
| bolt_log.write(f"BOLTing {obj_path.as_posix()}\n") |
| bolt_log.flush() |
| run_and_exit_on_failure( |
| f"{BOLT_PATH} {options} -o {obj_path_bolted} {obj_path}", |
| f"Failed to generate BOLTed binary for {obj_path}", |
| stdout=bolt_log, stderr=subprocess.STDOUT) |
| |
| bolt_log.write("\n") |
| bolt_log.flush() |
| |
| obj_path.unlink() |
| obj_path_bolted.rename(obj_path) |
| |
| |
| def process_objects(root: Path, 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): |
| if obj_path in optimization_subjects: |
| 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") |
| options += [ |
| f"--data={fdata_path}", |
| "--reorder-blocks=cache+", |
| "--reorder-functions=hfsort", |
| "--split-functions=2", |
| "--split-all-cold", |
| "--split-eh", |
| "--dyno-stats", |
| ] |
| |
| invoke_bolt(obj_path, bolt_log, " ".join(options)) |
| |
| else: |
| strip_symbols(obj_path) |
| |
| |
| def pack_archive(srcdir: str, outpath: Path) -> None: |
| print(f"Creating BOLTed archive") |
| shutil.make_archive(outpath.as_posix(), "gztar", srcdir) |
| |
| |
| def main() -> None: |
| args = parse_args() |
| |
| with BOLT_LOG_PATH.open("w") as bolt_log: |
| with unpack_archive(args.archive_path) as tmpdir: |
| process_objects(Path(tmpdir), args.profile_generate, args.profile_use, bolt_log) |
| pack_archive(tmpdir, args.dist_path / f"rust-{args.outname}") |
| |
| |
| if __name__ == "__main__": |
| main() |