pgo_tools: update llvm pgo generation

- use `workon_llvm.sh` instead of hand-installing things
- fix cross-* installation installation logic to install all
  cross-compilers, rather than trying to install (a potentially
  incorrect set of) binpkgs for some

BUG=b:408157467
TEST=create_chroot_and_generate_pgo_profile passes with a bootstrap
TEST=SDK update (crrev.com/c/6430206)

Change-Id: I0f1775a03633d919b9a6256df167f044caa732ae
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/6430096
Reviewed-by: Bob Haarman <[email protected]>
Commit-Queue: George Burgess <[email protected]>
Tested-by: George Burgess <[email protected]>

Author: gburgessiv

pgo_tools/create_chroot_and_generate_pgo_profile.py

@@ -99,16 +99,23 @@ def generate_pgo_profile(
     else:
         git_utils.checkout(git_dir=llvm_project, ref=sha)
 
+    logging.info("Working on all toolchain projects")
+    workon_llvm = (
+        cros_paths.CHROOT_SOURCE_ROOT
+        / cros_paths.TOOLCHAIN_UTILS
+        / "llvm_tools"
+        / "workon_llvm.sh"
+    )
     pgo_utils.run(
         cros_sdk
         + [
-            "cros-workon",
-            "--host",
-            "start",
-            "sys-devel/llvm",
+            workon_llvm,
+            # No need to set any boards up here.
+            "-",
         ],
         cwd=repo_root,
     )
+    logging.info("Running PGO profile generation script inside of the chroot")
     generate_pgo_profile_path = (
         cros_paths.CHROOT_SOURCE_ROOT
         / cros_paths.TOOLCHAIN_UTILS_PYBIN

pgo_tools/generate_pgo_profile.py

@@ -26,7 +26,7 @@
 import sys
 import tempfile
 import textwrap
-from typing import Dict, FrozenSet, Generator, List, Optional
+from typing import Dict, Generator, List, Optional
 
 from pgo_tools import pgo_utils
 
@@ -132,34 +132,17 @@ def restore_llvm_binpkg() -> None:
         pgo_utils.run(pgo_utils.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_utils.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)
+def ensure_toolchain_is_up_to_date():
+    """Ensures that all toolchain/cross-compiler packages are up-to-date."""
+    logging.info("Updating all toolchain packages to the new toolchain")
+    # Using the command from go/crostc-mage-misc#using-the-new-toolchain
     pgo_utils.run(
-        ["sudo", "emerge", "-j", "-G"] + missing_packages,
+        [
+            "sudo",
+            "cros_setup_toolchains",
+            "--include-boards=amd64-generic",
+            "--nousepkg",
+        ],
     )
 
 
@@ -501,19 +484,21 @@ def main(argv: List[str]):
             )
         )
 
-    logging.info("Ensuring `cross-` libraries are installed")
-    ensure_cross_libs_are_installed()
+    logging.info("Installing the new toolchain...")
+    ensure_toolchain_is_up_to_date()
     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()
+        # Subtle: ensure we always start with the initial LLVM. The new one
+        # might emit a PGO profile version that the old one can't parse, which
+        # makes it impossible to build the new with the old (w/ PGO enabled).
+        #
+        # The cross-compiler libs being out of sync is fine; `sys-devel/llvm`
+        # ships with the cross-compiler libs it needs for host builds.
+        restore_llvm_binpkg()
 
         logging.info("Building PGO instrumented LLVM")
         emerge_pgo_generate_llvm()

pgo_tools/generate_pgo_profile_test.py

@@ -1,64 +1,19 @@
-#!/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.
 
 """Tests for generate_pgo_profile."""
 
-from pathlib import Path
-import shutil
-import tempfile
-import unittest
 from unittest import mock
 
+from llvm_tools import test_helpers
 from pgo_tools import generate_pgo_profile
 from pgo_tools import pgo_utils
 
 
-class Test(unittest.TestCase):
+class Test(test_helpers.TempDirTestCase):
     """Tests for generate_pgo_profile."""
 
-    @mock.patch.object(pgo_utils, "run")
-    def test_find_missing_cross_libs_works_for_empty_results(self, mock_run):
-        mock_run.return_value.returncode = 3
-        mock_run.return_value.stdout = ""
-        self.assertEqual(
-            generate_pgo_profile.find_missing_cross_libs(),
-            generate_pgo_profile.ALL_NEEDED_CROSS_LIBS,
-        )
-
-        mock_run.return_value.returncode = 0
-        self.assertEqual(
-            generate_pgo_profile.find_missing_cross_libs(),
-            generate_pgo_profile.ALL_NEEDED_CROSS_LIBS,
-        )
-
-    @mock.patch.object(pgo_utils, "run")
-    def test_find_missing_cross_libs_filters_results_properly(self, mock_run):
-        mock_run.return_value.returncode = 0
-        mock_run.return_value.stdout = "\n".join(
-            generate_pgo_profile.ALL_NEEDED_CROSS_LIBS
-        )
-        self.assertEqual(generate_pgo_profile.find_missing_cross_libs(), set())
-
-        some_cross_libs = sorted(generate_pgo_profile.ALL_NEEDED_CROSS_LIBS)
-        del some_cross_libs[len(some_cross_libs) // 3 :]
-        mock_run.return_value.stdout = "\n".join(
-            some_cross_libs + ["cross-foo/bar"]
-        )
-
-        expected_result = generate_pgo_profile.ALL_NEEDED_CROSS_LIBS - set(
-            some_cross_libs
-        )
-        self.assertEqual(
-            generate_pgo_profile.find_missing_cross_libs(), expected_result
-        )
-
-    def make_tempdir(self) -> Path:
-        tempdir = Path(tempfile.mkdtemp(prefix="generate_pgo_profile_test_"))
-        self.addCleanup(lambda: shutil.rmtree(tempdir))
-        return tempdir
-
     def test_read_exactly_one_dirent_works(self):
         tempdir = self.make_tempdir()
         ent = tempdir / "one-ent"