diff options
Diffstat (limited to 'pgo_tools/generate_pgo_profile.py')
| -rwxr-xr-x | pgo_tools/generate_pgo_profile.py | 477 |
1 files changed, 477 insertions, 0 deletions
diff --git a/pgo_tools/generate_pgo_profile.py b/pgo_tools/generate_pgo_profile.py new file mode 100755 index 00000000..e966ad42 --- /dev/null +++ b/pgo_tools/generate_pgo_profile.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +# Copyright 2023 The ChromiumOS Authors +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Generates a PGO profile for LLVM. + +**This script is meant to be run from inside of the chroot.** + +Note that this script has a few (perhaps surprising) side-effects: +1. The first time this is run in a chroot, it will pack up your existing llvm + and save it as a binpkg. +2. This script clobbers your llvm installation. If the script is run to + completion, your old installation will be restored. If it does not, it may + not be. +""" + +import argparse +import dataclasses +import logging +import os +from pathlib import Path +import shlex +import shutil +import subprocess +import sys +import tempfile +import textwrap +from typing import Dict, FrozenSet, List, Optional + +import pgo_tools + + +# This script runs `quickpkg` on LLVM. This file saves the version of LLVM that +# was quickpkg'ed. +SAVED_LLVM_BINPKG_STAMP = Path("/tmp/generate_pgo_profile_old_llvm.txt") + +# Triple to build with when not trying to get backend coverage. +HOST_TRIPLE = "x86_64-pc-linux-gnu" + +# List of triples we want coverage for. +IMPORTANT_TRIPLES = ( + HOST_TRIPLE, + "x86_64-cros-linux-gnu", + "armv7a-cros-linux-gnueabihf", + "aarch64-cros-linux-gnu", +) + +# Set of all of the cross-* libraries we need. +ALL_NEEDED_CROSS_LIBS = frozenset( + f"cross-{triple}/{package}" + for triple in IMPORTANT_TRIPLES + if triple != HOST_TRIPLE + for package in ("glibc", "libcxx", "llvm-libunwind", "linux-headers") +) + + +def ensure_llvm_binpkg_exists() -> bool: + """Verifies that we have an LLVM binpkg to fall back on. + + Returns: + True if this function actually created a binpkg, false if one already + existed. + """ + if SAVED_LLVM_BINPKG_STAMP.exists(): + pkg = Path(SAVED_LLVM_BINPKG_STAMP.read_text(encoding="utf-8")) + # Double-check this, since this package is considered a cache artifact + # by portage. Ergo, it can _technically_ be GC'ed at any time. + if pkg.exists(): + return False + + pkg = pgo_tools.quickpkg_llvm() + SAVED_LLVM_BINPKG_STAMP.write_text(str(pkg), encoding="utf-8") + return True + + +def restore_llvm_binpkg(): + """Installs the binpkg created by ensure_llvm_binpkg_exists.""" + logging.info("Restoring non-PGO'ed LLVM installation") + pkg = Path(SAVED_LLVM_BINPKG_STAMP.read_text(encoding="utf-8")) + assert ( + pkg.exists() + ), f"Non-PGO'ed binpkg at {pkg} does not exist. Can't restore" + pgo_tools.run(pgo_tools.generate_quickpkg_restoration_command(pkg)) + + +def find_missing_cross_libs() -> FrozenSet[str]: + """Returns cross-* libraries that need to be installed for workloads.""" + equery_result = pgo_tools.run( + ["equery", "l", "--format=$cp", "cross-*/*"], + check=False, + stdout=subprocess.PIPE, + ) + + # If no matching package is found, equery will exit with code 3. + if equery_result.returncode == 3: + return ALL_NEEDED_CROSS_LIBS + + equery_result.check_returncode() + has_packages = {x.strip() for x in equery_result.stdout.splitlines()} + return ALL_NEEDED_CROSS_LIBS - has_packages + + +def ensure_cross_libs_are_installed(): + """Ensures that we have cross-* libs for all `IMPORTANT_TRIPLES`.""" + missing_packages = find_missing_cross_libs() + if not missing_packages: + logging.info("All cross-compiler libraries are already installed") + return + + missing_packages = sorted(missing_packages) + logging.info("Installing cross-compiler libs: %s", missing_packages) + pgo_tools.run( + ["sudo", "emerge", "-j", "-G"] + missing_packages, + ) + + +def emerge_pgo_generate_llvm(): + """Emerges a sys-devel/llvm with PGO instrumentation enabled.""" + force_use = ( + "llvm_pgo_generate -llvm_pgo_use" + # Turn ThinLTO off, since doing so results in way faster builds. + # This is assumed to be OK, since: + # - ThinLTO should have no significant impact on where Clang puts + # instrprof counters. + # - In practice, both "PGO generated with ThinLTO enabled," and "PGO + # generated without ThinLTO enabled," were benchmarked, and the + # performance difference between the two was in the noise. + " -thinlto" + # Turn ccache off, since if there are valid ccache artifacts from prior + # runs of this script, ccache will lead to us not getting profdata from + # those. + " -wrapper_ccache" + ) + use = (os.environ.get("USE", "") + " " + force_use).strip() + + # Use FEATURES=ccache since it's not much of a CPU time penalty, and if a + # user runs this script repeatedly, they'll appreciate it. :) + force_features = "ccache" + features = (os.environ.get("FEATURES", "") + " " + force_features).strip() + logging.info("Building LLVM with USE=%s", shlex.quote(use)) + pgo_tools.run( + [ + "sudo", + f"FEATURES={features}", + f"USE={use}", + "emerge", + "sys-devel/llvm", + ] + ) + + +def build_profiling_env(profile_dir: Path) -> Dict[str, str]: + profile_pattern = str(profile_dir / "profile-%m.profraw") + return { + "LLVM_PROFILE_OUTPUT_FORMAT": "profraw", + "LLVM_PROFILE_FILE": profile_pattern, + } + + +def ensure_clang_invocations_generate_profiles(clang_bin: str, tmpdir: Path): + """Raises an exception if clang doesn't generate profraw files. + + Args: + clang_bin: the path to a clang binary. + tmpdir: a place where this function can put temporary files. + """ + tmpdir = tmpdir / "ensure_profiles_generated" + tmpdir.mkdir(parents=True) + pgo_tools.run( + [clang_bin, "--help"], + extra_env=build_profiling_env(tmpdir), + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + is_empty = next(tmpdir.iterdir(), None) is None + if is_empty: + raise ValueError( + f"The clang binary at {clang_bin} generated no profile" + ) + shutil.rmtree(tmpdir) + + +def write_unified_cmake_file( + into_dir: Path, absl_subdir: Path, gtest_subdir: Path +): + (into_dir / "CMakeLists.txt").write_text( + textwrap.dedent( + f"""\ + cmake_minimum_required(VERSION 3.10) + + project(generate_pgo) + + add_subdirectory({gtest_subdir}) + add_subdirectory({absl_subdir})""" + ), + encoding="utf-8", + ) + + +def fetch_workloads_into(target_dir: Path): + """Fetches PGO generation workloads into `target_dir`.""" + # The workload here is absl and gtest. The reasoning behind that selection + # was essentially a mix of: + # - absl is reasonably-written and self-contained + # - gtest is needed if tests are to be built; in order to have absl do much + # of any linking, gtest is necessary. + # + # Use the version of absl that's bundled with ChromeOS at the time of + # writing. + target_dir.mkdir(parents=True) + + def fetch_and_extract(gs_url: str, into_dir: Path): + tgz_full = target_dir / os.path.basename(gs_url) + pgo_tools.run( + [ + "gsutil", + "cp", + gs_url, + tgz_full, + ], + ) + into_dir.mkdir() + + pgo_tools.run( + ["tar", "xaf", tgz_full], + cwd=into_dir, + ) + + absl_dir = target_dir / "absl" + fetch_and_extract( + gs_url="gs://chromeos-localmirror/distfiles/" + "abseil-cpp-a86bb8a97e38bc1361289a786410c0eb5824099c.tar.bz2", + into_dir=absl_dir, + ) + + gtest_dir = target_dir / "gtest" + fetch_and_extract( + gs_url="gs://chromeos-mirror/gentoo/distfiles/" + "gtest-1b18723e874b256c1e39378c6774a90701d70f7a.tar.gz", + into_dir=gtest_dir, + ) + + unpacked_absl_dir = read_exactly_one_dirent(absl_dir) + unpacked_gtest_dir = read_exactly_one_dirent(gtest_dir) + write_unified_cmake_file( + into_dir=target_dir, + absl_subdir=unpacked_absl_dir.relative_to(target_dir), + gtest_subdir=unpacked_gtest_dir.relative_to(target_dir), + ) + + +@dataclasses.dataclass(frozen=True) +class WorkloadRunner: + """Runs benchmark workloads.""" + + profraw_dir: Path + target_dir: Path + out_dir: Path + + def run( + self, + triple: str, + extra_cflags: Optional[str] = None, + sysroot: Optional[str] = None, + ): + logging.info( + "Running workload for triple %s, extra cflags %r", + triple, + extra_cflags, + ) + if self.out_dir.exists(): + shutil.rmtree(self.out_dir) + self.out_dir.mkdir(parents=True) + + clang = triple + "-clang" + profiling_env = build_profiling_env(self.profraw_dir) + if sysroot: + profiling_env["SYSROOT"] = sysroot + + cmake_command: pgo_tools.Command = [ + "cmake", + "-G", + "Ninja", + "-DCMAKE_BUILD_TYPE=RelWithDebInfo", + f"-DCMAKE_C_COMPILER={clang}", + f"-DCMAKE_CXX_COMPILER={clang}++", + "-DABSL_BUILD_TESTING=ON", + "-DABSL_USE_EXTERNAL_GOOGLETEST=ON", + "-DABSL_USE_GOOGLETEST_HEAD=OFF", + "-DABSL_FIND_GOOGLETEST=OFF", + ] + + if extra_cflags: + cmake_command += ( + f"-DCMAKE_C_FLAGS={extra_cflags}", + f"-DCMAKE_CXX_FLAGS={extra_cflags}", + ) + + cmake_command.append(self.target_dir) + pgo_tools.run( + cmake_command, + extra_env=profiling_env, + cwd=self.out_dir, + ) + + pgo_tools.run( + ["ninja", "-v", "all"], + extra_env=profiling_env, + cwd=self.out_dir, + ) + + +def read_exactly_one_dirent(directory: Path) -> Path: + """Returns the single Path under the given directory. Raises otherwise.""" + ents = directory.iterdir() + ent = next(ents, None) + if ent is not None: + if next(ents, None) is None: + return ent + raise ValueError(f"Expected exactly one entry under {directory}") + + +def run_workloads(target_dir: Path) -> Path: + """Runs all of our workloads in target_dir. + + Args: + target_dir: a directory that already had `fetch_workloads_into` called + on it. + + Returns: + A directory in which profraw files from running the workloads are + saved. + """ + profraw_dir = target_dir / "profiles" + profraw_dir.mkdir() + + out_dir = target_dir / "out" + runner = WorkloadRunner( + profraw_dir=profraw_dir, + target_dir=target_dir, + out_dir=out_dir, + ) + + # Run the workload once per triple. + for triple in IMPORTANT_TRIPLES: + runner.run( + triple, sysroot=None if triple == HOST_TRIPLE else f"/usr/{triple}" + ) + + # Add a run of ThinLTO, so any ThinLTO-specific lld bits get exercised. + # Also, since CrOS uses -Os often, exercise that. + runner.run(HOST_TRIPLE, extra_cflags="-flto=thin -Os") + return profraw_dir + + +def convert_profraw_to_pgo_profile(profraw_dir: Path) -> Path: + """Creates a PGO profile from the profraw profiles in profraw_dir.""" + output = profraw_dir / "merged.prof" + profile_files = list(profraw_dir.glob("profile-*profraw")) + if not profile_files: + raise ValueError("No profraw files generated?") + + logging.info( + "Creating a PGO profile from %d profraw files", len(profile_files) + ) + generate_command = [ + "llvm-profdata", + "merge", + "--instr", + f"--output={output}", + ] + pgo_tools.run(generate_command + profile_files) + return output + + +def main(argv: List[str]): + logging.basicConfig( + format=">> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: " + "%(message)s", + level=logging.DEBUG, + ) + + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--output", + required=True, + type=Path, + help="Where to put the PGO profile", + ) + parser.add_argument( + "--use-old-binpkg", + action="store_true", + help=""" + This script saves your initial LLVM installation as a binpkg, so it may + restore that installation later in the build. Passing --use-old-binpkg + allows this script to use a binpkg from a prior invocation of this + script. + """, + ) + opts = parser.parse_args(argv) + + pgo_tools.exit_if_not_in_chroot() + + output = opts.output + + llvm_binpkg_is_fresh = ensure_llvm_binpkg_exists() + if not llvm_binpkg_is_fresh and not opts.use_old_binpkg: + sys.exit( + textwrap.dedent( + f"""\ + A LLVM binpkg packed by a previous run of this script is + available. If you intend this run to be another attempt at the + previous run, please pass --use-old-binpkg (so the old LLVM + binpkg is used as our 'baseline'). If you don't, please remove + the file referring to it at {SAVED_LLVM_BINPKG_STAMP}. + """ + ) + ) + + logging.info("Ensuring `cross-` libraries are installed") + ensure_cross_libs_are_installed() + tempdir = Path(tempfile.mkdtemp(prefix="generate_llvm_pgo_profile_")) + try: + workloads_path = tempdir / "workloads" + logging.info("Fetching workloads") + fetch_workloads_into(workloads_path) + + # If our binpkg is not fresh, we may be operating with a weird LLVM + # (e.g., a PGO'ed one ;) ). Ensure we always start with that binpkg as + # our baseline. + if not llvm_binpkg_is_fresh: + restore_llvm_binpkg() + + logging.info("Building PGO instrumented LLVM") + emerge_pgo_generate_llvm() + + logging.info("Ensuring instrumented compilers generate profiles") + for triple in IMPORTANT_TRIPLES: + ensure_clang_invocations_generate_profiles( + triple + "-clang", tempdir + ) + + logging.info("Running workloads") + profraw_dir = run_workloads(workloads_path) + + # This is a subtle but critical step. The LLVM we're currently working + # with was built by the LLVM represented _by our binpkg_, which may be + # a radically different version of LLVM than what was installed (e.g., + # it could be from our bootstrap SDK, which could be many months old). + # + # If our current LLVM's llvm-profdata is used to interpret the profraw + # files: + # 1. The profile generated will be for our new version of clang, and + # may therefore be too new for the older version that we still have + # to support. + # 2. There may be silent incompatibilities, as the stability guarantees + # of profraw files are not immediately apparent. + logging.info("Restoring LLVM's binpkg") + restore_llvm_binpkg() + pgo_profile = convert_profraw_to_pgo_profile(profraw_dir) + shutil.copyfile(pgo_profile, output) + except: + # Leave the tempdir, as it might help people debug. + logging.info("NOTE: Tempdir will remain at %s", tempdir) + raise + + logging.info("Removing now-obsolete tempdir") + shutil.rmtree(tempdir) + logging.info("PGO profile is available at %s.", output) + + +if __name__ == "__main__": + main(sys.argv[1:]) |