boltifyer.py

#!/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

from paths import *
from utils import (
    ExtantPath,
    ResolvedPath,
    archive_create,
    archive_extract,
    extend_suffix,
    get_prebuilt_binary_paths,
    print_context,
    print_fs_tree,
    is_archive,
    reify_singleton_patterns,
    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=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")

    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")
    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 unpack_prebuilt_archive(archive_path: Path) -> Iterator[Path]:
    with TemporaryDirectory() as temp_dir:
        temp_dir_path = Path(temp_dir)
        print(f"Unpacking archive into {temp_dir}")
        if archive_extract(archive_path, temp_dir_path) != 0:
            raise RuntimeError("Failed to unpack prebuilt archive")
        yield temp_dir_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}")
                if archive_extract(profile_path, temp_dir_path) != 0:
                    raise RuntimeError("Failed to unpack profiles")
                yield temp_dir_path
        else:
            yield profile_path
    else:
        yield None


def invoke_bolt(obj_path: Path, bolt_log: TextIO, options: list[str] = []) -> int:
    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()
    retcode = subprocess.run(
        [BOLT_PATH] + options + ["-o", obj_path_bolted, obj_path],
        stdout=bolt_log, stderr=bolt_log).returncode

    if retcode == 0:
        bolt_log.write("\n")
        bolt_log.flush()

        obj_path.unlink()
        obj_path_bolted.rename(obj_path)

    else:
        print(f"Failed to process object {str(obj_path)}")

    return retcode


def process_objects(root: Path, profile_generate: Optional[Path], profile_use: Optional[Path], bolt_log: TextIO) -> int:
    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):
        print(f"Boltifyer is looking at object {str(obj_path)}")

        if 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")
                    options += [
                        f"--data={fdata_path}",
                        "--reorder-blocks=ext-tsp",
                        "--reorder-functions=hfsort",
                        "--split-functions",
                        "--split-all-cold",
                        "--split-eh",
                        "--dyno-stats",
                    ]

            retcode = invoke_bolt(obj_path, bolt_log, options)
            if retcode != 0:
                print(f"Failed to run BOLT on {obj_path.as_posix()}")
                return retcode

        else:
            print(f"Object {str(obj_path)} is NOT an optimization subject")
            retcode = strip_symbols(obj_path)
            if retcode != 0:
                print(f"Failed to strip symbols from {obj_path.as_posix()}")
                return retcode

    return 0


def boltify_archive(input_path: Path, dist_path: Path, build_name: str, profile_generate: Path | None = None, profile_use: Path | None = None) -> int:
    with (dist_path / BOLT_LOG_NAME).open("w") as bolt_log:
        with unpack_prebuilt_archive(input_path) as toolchain_dir_path:
            print_fs_tree(toolchain_dir_path)

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

                retcode = process_objects(toolchain_dir_path, profile_generate, profile_use_path, bolt_log)
                bolt_log.flush()
                if retcode != 0:
                    print("Error BOLTing prebuilt objects")
                    return retcode

                print(f"Creating BOLTed archive")
                retcode = archive_create(dist_path / f"rust-{build_name}", toolchain_dir_path)
                if retcode != 0:
                    print("Error creating the BOLTed archive")
                return retcode


def main() -> int:
    print_context()

    args   = parse_args()
    with (args.dist_path / BOLT_LOG_NAME).open("w") as bolt_log:
        bolt_log.write(f"Skipping BOLT on modified branch\n")
    retval = 0

    #retval = boltify_archive(
    #    args.archive_path,
    #    args.dist_path,
    #    args.build_name,
    #    args.profile_generate,
    #    args.profile_use)

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

    return retval


if __name__ == "__main__":
    sys.exit(main())