1#!/usr/bin/env python3 2# Copyright 2023 The ChromiumOS Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6"""Runs benchmarks, given potentially multiple PGO profiles. 7 8**This script is meant to be run from inside of the chroot.** 9 10This script overwrites your chroot's LLVM temporarily, but if it runs to 11completion, it will restore you to your previous version of LLVM. Care is taken 12so that the same baseline LLVM is used to build all LLVM versions this script 13benchmarks. 14""" 15 16import argparse 17import dataclasses 18import enum 19import json 20import logging 21from pathlib import Path 22import shlex 23import shutil 24import subprocess 25import sys 26from typing import IO, List, Optional, Union 27 28import pgo_tools 29 30 31# The full path to where `sys-devel/llvm` expects local profiles to be if 32# `USE=llvm_pgo_use_local` is specified. 33LOCAL_PROFILE_LOCATION = Path( 34 "/mnt/host/source/src/third_party/chromiumos-overlay", 35 "sys-devel/llvm/files/llvm-local.profdata", 36).resolve() 37 38CHROOT_HYPERFINE = Path.home() / ".cargo/bin/hyperfine" 39 40 41class SpecialProfile(enum.Enum): 42 """An enum representing a 'special' (non-Path) profile.""" 43 44 REMOTE = enum.auto() 45 NONE = enum.auto() 46 47 def __str__(self) -> str: 48 if self is self.REMOTE: 49 return "@remote" 50 if self is self.NONE: 51 return "@none" 52 raise ValueError(f"Unknown SpecialProfile value: {repr(self)}") 53 54 55@dataclasses.dataclass(frozen=True, eq=True) 56class RunData: 57 """Data describing the results of one hyperfine run.""" 58 59 tag: str 60 user_time: float 61 system_time: float 62 63 @staticmethod 64 def from_json(tag: str, json_contents: IO) -> "RunData": 65 """Converts a hyperfine JSON file's contents into a RunData.""" 66 results = json.load(json_contents)["results"] 67 if len(results) != 1: 68 raise ValueError(f"Expected one run result; got {results}") 69 return RunData( 70 tag=tag, 71 user_time=results[0]["user"], 72 system_time=results[0]["system"], 73 ) 74 75 76ProfilePath = Union[SpecialProfile, Path] 77 78 79def parse_profile_path(path: str) -> ProfilePath: 80 for p in SpecialProfile: 81 if path == str(p): 82 return p 83 return Path(path).resolve() 84 85 86def ensure_hyperfine_is_installed(): 87 if CHROOT_HYPERFINE.exists(): 88 return 89 90 logging.info("Installing hyperfine for benchmarking...") 91 pgo_tools.run( 92 [ 93 "cargo", 94 "install", 95 "hyperfine", 96 ] 97 ) 98 assert ( 99 CHROOT_HYPERFINE.exists() 100 ), f"hyperfine was installed, but wasn't at {CHROOT_HYPERFINE}" 101 102 103def construct_hyperfine_cmd( 104 llvm_ebuild: Path, 105 profile: ProfilePath, 106 llvm_binpkg: Path, 107 use_thinlto: bool, 108 export_json: Optional[Path] = None, 109) -> pgo_tools.Command: 110 if isinstance(profile, Path): 111 if profile != LOCAL_PROFILE_LOCATION: 112 shutil.copyfile(profile, LOCAL_PROFILE_LOCATION) 113 use_flags = "-llvm_pgo_use llvm_pgo_use_local" 114 elif profile is SpecialProfile.NONE: 115 use_flags = "-llvm_pgo_use" 116 elif profile is SpecialProfile.REMOTE: 117 use_flags = "llvm_pgo_use" 118 else: 119 raise ValueError(f"Unknown profile type: {type(profile)}") 120 121 quickpkg_restore = shlex.join( 122 str(x) 123 for x in pgo_tools.generate_quickpkg_restoration_command(llvm_binpkg) 124 ) 125 126 setup_cmd = ( 127 f"{quickpkg_restore} && " 128 f"sudo FEATURES=ccache USE={shlex.quote(use_flags)}" 129 # Use buildpkg-exclude so our existing llvm binpackage isn't 130 # overwritten. 131 " emerge sys-devel/llvm --buildpkg-exclude=sys-devel/llvm" 132 ) 133 134 if use_thinlto: 135 benchmark_use = "thinlto" 136 else: 137 benchmark_use = "-thinlto" 138 139 ebuild_llvm = ( 140 f"sudo USE={shlex.quote(benchmark_use)} " 141 f"ebuild {shlex.quote(str(llvm_ebuild))}" 142 ) 143 cmd: pgo_tools.Command = [ 144 CHROOT_HYPERFINE, 145 "--max-runs=3", 146 f"--setup={setup_cmd}", 147 f"--prepare={ebuild_llvm} clean prepare", 148 ] 149 150 if export_json: 151 cmd.append(f"--export-json={export_json}") 152 153 cmd += ( 154 "--", 155 # At the moment, building LLVM seems to be an OK benchmark. It has some 156 # C in it, some C++, and each pass on Cloudtops takes no more than 7 157 # minutes. 158 f"{ebuild_llvm} compile", 159 ) 160 return cmd 161 162 163def validate_profiles( 164 parser: argparse.ArgumentParser, profiles: List[ProfilePath] 165): 166 number_of_path_profiles = 0 167 nonexistent_profiles = [] 168 seen_profile_at_local_profile_location = False 169 for profile in profiles: 170 if not isinstance(profile, Path): 171 continue 172 173 if not profile.exists(): 174 nonexistent_profiles.append(profile) 175 176 number_of_path_profiles += 1 177 if profile == LOCAL_PROFILE_LOCATION: 178 seen_profile_at_local_profile_location = True 179 180 if number_of_path_profiles > 1 and seen_profile_at_local_profile_location: 181 parser.error( 182 f"Cannot use the path {LOCAL_PROFILE_LOCATION} as a profile if " 183 "there are other profiles specified by path." 184 ) 185 186 if nonexistent_profiles: 187 nonexistent_profiles.sort() 188 parser.error( 189 "One or more profiles do not exist: " f"{nonexistent_profiles}" 190 ) 191 192 193def run_benchmark( 194 use_thinlto: bool, 195 profiles: List[ProfilePath], 196) -> List[RunData]: 197 """Runs the PGO benchmark with the given parameters. 198 199 Args: 200 use_thinlto: whether to benchmark the use of ThinLTO 201 profiles: profiles to benchmark with 202 collect_run_data: whether to return a CombinedRunData 203 204 Returns: 205 A CombinedRunData instance capturing the performance of the benchmark 206 runs. 207 """ 208 ensure_hyperfine_is_installed() 209 210 llvm_ebuild_path = Path( 211 pgo_tools.run( 212 ["equery", "w", "sys-devel/llvm"], stdout=subprocess.PIPE 213 ).stdout.strip() 214 ) 215 216 baseline_llvm_binpkg = pgo_tools.quickpkg_llvm() 217 accumulated_run_data = [] 218 with pgo_tools.temporary_file( 219 prefix="benchmark_pgo_profile" 220 ) as tmp_json_file: 221 for profile in profiles: 222 cmd = construct_hyperfine_cmd( 223 llvm_ebuild_path, 224 profile, 225 baseline_llvm_binpkg, 226 use_thinlto=use_thinlto, 227 export_json=tmp_json_file, 228 ) 229 # Format the profile with `repr(str(profile))` so that we always 230 # get a quoted, but human-friendly, representation of the profile. 231 logging.info( 232 "Profile %r: Running %s", 233 str(profile), 234 shlex.join(str(x) for x in cmd), 235 ) 236 pgo_tools.run(cmd) 237 238 with tmp_json_file.open(encoding="utf-8") as f: 239 accumulated_run_data.append(RunData.from_json(str(profile), f)) 240 241 logging.info("Restoring original LLVM...") 242 pgo_tools.run( 243 pgo_tools.generate_quickpkg_restoration_command(baseline_llvm_binpkg) 244 ) 245 return accumulated_run_data 246 247 248def main(argv: List[str]): 249 logging.basicConfig( 250 format=">> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: " 251 "%(message)s", 252 level=logging.INFO, 253 ) 254 255 parser = argparse.ArgumentParser( 256 description=__doc__, 257 formatter_class=argparse.RawDescriptionHelpFormatter, 258 ) 259 parser.add_argument( 260 "--thinlto", 261 action="store_true", 262 help="If specified, this will benchmark builds with ThinLTO enabled.", 263 ) 264 parser.add_argument( 265 "profile", 266 nargs="+", 267 type=parse_profile_path, 268 help=f""" 269 The path to a profile to benchmark. There are two special values here: 270 '{SpecialProfile.REMOTE}' and '{SpecialProfile.NONE}'. For 271 '{SpecialProfile.REMOTE}', this will just use the default LLVM PGO 272 profile for a benchmark run. For '{SpecialProfile.NONE}', all PGO will 273 be disabled for a benchmark run. 274 """, 275 ) 276 opts = parser.parse_args(argv) 277 278 pgo_tools.exit_if_not_in_chroot() 279 280 profiles = opts.profile 281 validate_profiles(parser, profiles) 282 283 run_benchmark(opts.thinlto, profiles) 284 285 286if __name__ == "__main__": 287 main(sys.argv[1:]) 288