recipes/contrib/clang_ml_training.py

# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Recipe for training ML inliner model for Clang."""

from urllib.parse import urlparse

from PB.go.chromium.org.luci.buildbucket.proto import build as build_pb2
from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb2
from PB.recipes.fuchsia.contrib.clang_ml_training import InputProperties

DEPS = [
    "fuchsia/build",
    "fuchsia/buildbucket_util",
    "fuchsia/cas_util",
    "fuchsia/checkout",
    "fuchsia/cipd_util",
    "fuchsia/docker",
    "fuchsia/git",
    "fuchsia/git_checkout",
    "fuchsia/jiri",
    "fuchsia/sso",
    "fuchsia/tensorflow",
    "recipe_engine/buildbucket",
    "recipe_engine/cipd",
    "recipe_engine/context",
    "recipe_engine/file",
    "recipe_engine/json",
    "recipe_engine/path",
    "recipe_engine/properties",
    "recipe_engine/raw_io",
    "recipe_engine/step",
]

PROPERTIES = InputProperties

PLATFORM_TO_TARGET = {
    "linux-amd64": "x86_64-unknown-linux-gnu",
    "linux-arm64": "aarch64-unknown-linux-gnu",
    "mac-amd64": "x86_64-apple-darwin",
    "mac-arm64": "arm64-apple-darwin",
    "windows-amd64": "x86_64-pc-windows-msvc",
    "windows-arm64": "aarch64-pc-windows-msvc",
}

ML_COMPILER_OPT_TAG = "4120b973b3843a2b604533a40a0febcfba8e85ee"

REGISTRY_URL = "us-central1-docker.pkg.dev/fuchsia-artifacts/fuchsia-clang-mlgo"

IMAGE_NAME = "ml-compiler-opt"

CPUINFO_REPOSITORY = (
    "https://fuchsia.googlesource.com/third_party/github.com/pytorch/cpuinfo/"
)
CPUINFO_TAG = "e77879a708ba3cb7e6966a93fbbc0ec044300038"

RUY_REPOSITORY = "https://fuchsia.googlesource.com/third_party/github.com/google/ruy"
RUY_TAG = "3286a34cc8de6149ac6844107dfdffac91531e72"

ABSEIL_REPOSITORY = "https://fuchsia.googlesource.com/third_party/abseil-cpp"
ABSEIL_TAG = "273292d1cfc0a94a65082ee350509af1d113344d"

EIGEN_REPOSITORY = "https://third-party-mirror.googlesource.com/eigen"
EIGEN_TAG = "cefe4a0684b2d4165e775e1a4e8f0eaee8311cfe"

NEON_2_SSE_REPOSITORY = (
    "https://fuchsia.googlesource.com/third_party/github.com/intel/ARM_NEON_2_x86_SSE"
)
NEON_2_SSE_TAG = "a15b489e1222b2087007546b4912e21293ea86ff"

FLATBUFFERS_REPOSITORY = "https://fuchsia.googlesource.com/third_party/flatbuffers"
FLATBUFFERS_TAG = "615616cb5549a34bdf288c04bc1b94bd7a65c396"

PROTOBUF_REPOSITORY = "https://fuchsia.googlesource.com/third_party/protobuf"
PROTOBUF_TAG = "c9869dc7803eb0a21d7e589c40ff4f9288cd34ae"

# TODO: switch to our mirrors

TENSORFLOW_REPOSITORY = (
    "https://fuchsia.googlesource.com/third_party/github.com/tensorflow/tensorflow"
)
TENSORFLOW_TAG = "d74d591e8f9e42be640dd93cfa0584914e25e98e"

GEMMLOWP_REPOSITORY = (
    "https://fuchsia.googlesource.com/third_party/github.com/google/gemmlowp"
)
GEMMLOWP_TAG = "08e4bb339e34017a0835269d4a37c4ea04d15a69"

FFT_REPOSITORY = (
    "https://fuchsia.googlesource.com/third_party/github.com/petewarden/OouraFFT"
)
FFT_TAG = "c6fd2dd6d21397baa6653139d31d84540d5449a2"

FARMHASH_REPOSITORY = (
    "https://fuchsia.googlesource.com/third_party/github.com/google/farmhash"
)
FARMHASH_TAG = "f5996d61e7d7352c06cc0598b7ae98b6e1adc75a"


def RunSteps(api, props):
    props.revision = props.revision or api.git.get_remote_branch_head(
        api.sso.sso_to_https(props.remote),
        "refs/heads/releases/canary",
    )
    # Incremental cache is not used, checkout.root_dir will be located
    # under api.path.start_dir
    checkout = api.checkout.fuchsia_with_options(
        build_input=build_pb2.Build.Input(
            gitiles_commit=common_pb2.GitilesCommit(
                id=props.revision,
                project=props.project,
                host=urlparse(api.sso.sso_to_https(props.remote)).hostname,
            ),
        ),
        project=props.project,
        manifest=props.manifest,
        remote=props.remote,
    )

    with api.context(cwd=checkout.root_dir):
        revision = checkout.project("fuchsia")["revision"]

        package_data = api.jiri.package(
            ["fuchsia/third_party/clang/${platform}"],
            test_data=[
                {
                    "name": "fuchsia/third_party/clang/${platform}",
                    "manifest": str(
                        checkout.root_dir.joinpath("integration", "prebuilts")
                    ),
                    "path": str(
                        checkout.root_dir.joinpath(
                            "prebuilt", "third_party", "clang", "linux-x64"
                        )
                    ),
                    "platforms": [
                        "linux-amd64",
                        "linux-arm64",
                        "mac-amd64",
                        "windows-amd64",
                    ],
                    "version": "git_revision:f52666985d7011b539f26f54e09a5c89b62dad56",
                }
            ],
        )
        assert len(package_data) == 1
        clang_version = package_data[0]["version"].split(":", 1)[1]

    with api.step.nest("prerequisites"):
        cipd_dir = api.path.start_dir / "cipd"
        pkgs = api.cipd.EnsureFile()
        pkgs.add_package("fuchsia/sdk/core/linux-amd64", "latest", "sdk")
        pkgs.add_package(
            "fuchsia/third_party/sysroot/linux",
            "integration",
            "linux",
        )
        pkgs.add_package(
            "fuchsia/third_party/cmake/${platform}",
            "integration",
        )
        pkgs.add_package(
            "fuchsia/third_party/ninja/${platform}",
            "integration",
        )
        api.cipd.ensure(cipd_dir, pkgs)
        sdk_dir = cipd_dir / "sdk"
        sysroot_dir = cipd_dir / "linux"

    # cpuinfo
    with api.step.nest("cpuinfo"):
        cpuinfo_src_dir, _ = api.git_checkout(
            CPUINFO_REPOSITORY,
            revision=CPUINFO_TAG,
        )
        cpuinfo_build_dir = api.path.start_dir / "cpuinfo_build"
        cpuinfo_install_dir = api.path.start_dir / "cpuinfo_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                cpuinfo_src_dir,
                "-B",
                cpuinfo_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={cpuinfo_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                "-DCPUINFO_BUILD_UNIT_TESTS:BOOL=OFF",
                "-DCPUINFO_BUILD_MOCK_TESTS:BOOL=OFF",
                "-DCPUINFO_BUILD_BENCHMARKS:BOOL=OFF",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                cpuinfo_build_dir,
                "install",
            ],
        )

    # ruy
    with api.step.nest("ruy"):
        ruy_src_dir, _ = api.git_checkout(
            RUY_REPOSITORY,
            revision=RUY_TAG,
        )
        ruy_build_dir = api.path.start_dir / "ruy_build"
        ruy_install_dir = api.path.start_dir / "ruy_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                ruy_src_dir,
                "-B",
                ruy_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={ruy_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                f"-Dcpuinfo_DIR:PATH={cpuinfo_install_dir.joinpath('share', 'cpuinfo')}",
                "-DRUY_MINIMAL_BUILD:BOOL=ON",
                "-DRUY_ENABLE_INSTALL:BOOL=ON",
                "-DRUY_FIND_CPUINFO:BOOL=ON",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                ruy_build_dir,
                "install",
            ],
        )

    # absl
    with api.step.nest("absl"):
        absl_src_dir, _ = api.git_checkout(
            ABSEIL_REPOSITORY,
            revision=ABSEIL_TAG,
        )
        absl_build_dir = api.path.start_dir / "absl_build"
        absl_install_dir = api.path.start_dir / "absl_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                absl_src_dir,
                "-B",
                absl_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={absl_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                "-DABSL_BUILD_TESTING:BOOL=OFF",
                "-DABSL_ENABLE_INSTALL:BOOL=ON",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                absl_build_dir,
                "install",
            ],
        )

    # eigen
    with api.step.nest("eigen"):
        eigen_src_dir, _ = api.git_checkout(
            EIGEN_REPOSITORY,
            revision=EIGEN_TAG,
        )
        eigen_build_dir = api.path.start_dir / "eigen_build"
        eigen_install_dir = api.path.start_dir / "eigen_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                eigen_src_dir,
                "-B",
                eigen_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={eigen_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                "-DEIGEN_BUILD_DOC:BOOL=OFF",
                "-DEIGEN_BUILD_TESTING:BOOL=OFF",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                eigen_build_dir,
                "install",
            ],
        )

    # ARM_NEON_2_x86_SSE
    with api.step.nest("neon2sse"):
        neon2sse_src_dir, _ = api.git_checkout(
            NEON_2_SSE_REPOSITORY,
            revision=NEON_2_SSE_TAG,
        )
        neon2sse_build_dir = api.path.start_dir / "neon2sse_build"
        neon2sse_install_dir = api.path.start_dir / "neon2sse_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                neon2sse_src_dir,
                "-B",
                neon2sse_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={neon2sse_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                neon2sse_build_dir,
                "install",
            ],
        )

    # flatbuffers
    with api.step.nest("flatbuffers"):
        flatbuffers_src_dir, _ = api.git_checkout(
            FLATBUFFERS_REPOSITORY,
            revision=FLATBUFFERS_TAG,
        )
        flatbuffers_build_dir = api.path.start_dir / "flatbuffers_build"
        flatbuffers_install_dir = api.path.start_dir / "flatbuffers_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                flatbuffers_src_dir,
                "-B",
                flatbuffers_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={flatbuffers_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                "-DFLATBUFFERS_BUILD_TESTS:BOOL=OFF",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                flatbuffers_build_dir,
                "install",
            ],
        )

    # tensorflow
    with api.step.nest("tensorflow"):
        # TODO: extend tensorflow to allow using our own checkouts
        # gemmlowp_src_dir, _ = api.git_checkout(
        #     GEMMLOWP_REPOSITORY,
        #     revision=GEMMLOWP_TAG,
        # )
        # fft_src_dir, _ = api.git_checkout(
        #     FFT_REPOSITORY,
        #     revision=FFT_TAG,
        # )
        # farmhash_src_dir, _ = api.git_checkout(
        #     FARMHASH_REPOSITORY,
        #     revision=FARMHASH_TAG,
        # )

        tensorflow_src_dir, _ = api.git_checkout(
            TENSORFLOW_REPOSITORY,
            revision=TENSORFLOW_TAG,
        )
        tensorflow_build_dir = api.path.start_dir / "tensorflow_build"
        tensorflow_install_dir = api.path.start_dir / "tensorflow_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                tensorflow_src_dir.joinpath("tensorflow", "lite"),
                "-B",
                tensorflow_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={tensorflow_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                "-DTFLITE_ENABLE_INSTALL:BOOL=ON",
                "-DTFLITE_ENABLE_XNNPACK:BOOL=OFF",
                f"-Dcpuinfo_DIR:PATH={cpuinfo_install_dir.joinpath('share', 'cpuinfo')}",
                f"-Druy_DIR:PATH={ruy_install_dir.joinpath('lib', 'cmake', 'ruy')}",
                f"-Dabsl_DIR:PATH={absl_install_dir.joinpath('lib', 'cmake', 'absl')}",
                f"-DEigen3_DIR:PATH={eigen_install_dir.joinpath('share', 'eigen3', 'cmake')}",
                f"-DNEON_2_SSE_DIR:PATH={neon2sse_install_dir.joinpath('lib', 'cmake', 'NEON_2_SSE')}",
                f"-DFlatbuffers_DIR:PATH={flatbuffers_install_dir.joinpath('lib', 'cmake', 'flatbuffers')}",
                f"-DOVERRIDABLE_FETCH_CONTENT_gemmlowp_GIT_REPOSITORY={GEMMLOWP_REPOSITORY}",
                f"-DOVERRIDABLE_FETCH_CONTENT_gemmlowp_GIT_TAG={GEMMLOWP_TAG}",
                f"-DOVERRIDABLE_FETCH_CONTENT_fft2d_GIT_REPOSITORY={FFT_REPOSITORY}",
                f"-DOVERRIDABLE_FETCH_CONTENT_fft2d_GIT_TAG={FFT_TAG}",
                f"-DOVERRIDABLE_FETCH_CONTENT_farmhash_GIT_REPOSITORY={FARMHASH_REPOSITORY}",
                f"-DOVERRIDABLE_FETCH_CONTENT_farmhash_GIT_TAG={FARMHASH_TAG}",
                "-DOVERRIDABLE_FETCH_CONTENT_USE_GIT=ON",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                tensorflow_build_dir,
                "install",
            ],
        )

    # protobuf
    with api.step.nest("protobuf"):
        protobuf_src_dir, _ = api.git_checkout(
            PROTOBUF_REPOSITORY,
            revision=PROTOBUF_TAG,
        )
        protobuf_build_dir = api.path.start_dir / "protobuf_build"
        protobuf_install_dir = api.path.start_dir / "protobuf_install"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                protobuf_src_dir / "cmake",
                "-B",
                protobuf_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM:PATH={cipd_dir / 'ninja'}",
                "-DCMAKE_FIND_PACKAGE_PREFER_CONFIG:BOOL=ON",
                f"-DCMAKE_INSTALL_PREFIX:PATH={protobuf_install_dir}",
                "-DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                "-Dprotobuf_BUILD_TESTS:BOOL=OFF",
            ],
        )
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                protobuf_build_dir,
                "install",
            ],
        )

    # llvm
    with api.step.nest("llvm"):
        llvm_src_dir, _ = api.git_checkout(
            "https://llvm.googlesource.com/llvm-project",
            revision=clang_version,
        )
        llvm_build_dir = api.path.start_dir / "llvm_build"
        api.step(
            "configure",
            [
                cipd_dir.joinpath("bin", "cmake"),
                "-S",
                llvm_src_dir / "llvm",
                "-B",
                llvm_build_dir,
                "-G",
                "Ninja",
                f"-DCMAKE_MAKE_PROGRAM={cipd_dir / 'ninja'}",
                "-DCMAKE_INSTALL_PREFIX=",
                "-DLLVM_ENABLE_LTO=OFF",
                f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                f"-DLINUX_x86_64-unknown-linux-gnu_SYSROOT={sysroot_dir}",
                f"-DLINUX_aarch64-unknown-linux-gnu_SYSROOT={sysroot_dir}",
                f"-DFUCHSIA_SDK={cipd_dir / 'sdk'}",
                f"-Dcpuinfo_DIR:PATH={cpuinfo_install_dir.joinpath('share', 'cpuinfo')}",
                f"-Druy_DIR:PATH={ruy_install_dir.joinpath('lib', 'cmake', 'ruy')}",
                f"-Dabsl_DIR:PATH={absl_install_dir.joinpath('lib', 'cmake', 'absl')}",
                f"-DEigen3_DIR:PATH={eigen_install_dir.joinpath('share', 'eigen3', 'cmake')}",
                f"-DNEON_2_SSE_DIR:PATH={neon2sse_install_dir.joinpath('lib', 'cmake', 'NEON_2_SSE')}",
                f"-DFlatbuffers_DIR:PATH={flatbuffers_install_dir.joinpath('lib', 'cmake', 'flatbuffers')}",
                f"-Dprotobuf_DIR:PATH={protobuf_install_dir.joinpath('lib', 'cmake', 'protobuf')}",
                f"-Dtensorflow-lite_DIR:PATH={tensorflow_install_dir.joinpath('lib', 'cmake', 'tensorflow-lite')}",
                f"-DTENSORFLOW_SRC_DIR:PATH={tensorflow_src_dir}",
                "-DLLVM_HAVE_TFLITE:BOOL=ON",
                "-C",
                llvm_src_dir.joinpath(
                    "clang",
                    "cmake",
                    "caches",
                    "Fuchsia-stage2.cmake",
                ),
            ],
        )

        # Build the Clang distribution.
        api.step(
            "build",
            [
                cipd_dir / "ninja",
                "-C",
                llvm_build_dir,
                "distribution",
            ],
        )

        # Create the Clang toolchain.
        clang_dir = api.path.start_dir / "clang"
        with api.context(env={"DESTDIR": clang_dir}):
            api.step(
                "install",
                [
                    cipd_dir / "ninja",
                    "-C",
                    llvm_build_dir,
                    "install-distribution",
                ],
            )
        api.step(
            "generate runtimes.json",
            cmd=[
                "vpython3",
                checkout.root_dir.joinpath("scripts", "clang", "generate_runtimes.py"),
                f"--clang-prefix={clang_dir}",
                f"--sdk-dir={sdk_dir}",
                f"--build-id-dir={clang_dir.joinpath('lib', '.build-id')}",
            ],
            stdout=api.raw_io.output_text(
                leak_to=clang_dir.joinpath("lib", "runtime.json"), add_output_log=True
            ),
        )

        # Upload the Clang binaries to CAS.
        api.cas_util.upload(clang_dir, step_name="archive clang")

        # Use the just built Clang toolchain to build Fuchsia.
        checkout.clang_toolchain_dir = clang_dir

    with api.step.nest("corpus"):
        corpus_dir = api.path.start_dir / "corpus"

        for target_arch, fint_params_path in props.fint_params_paths.items():
            with api.step.nest(target_arch):
                fint_params = api.file.read_text(
                    "read fint params",
                    checkout.root_dir / fint_params_path,
                    test_data='field: "value"',
                )
                fint_params += 'gn_args: "clang_embed_bitcode = true"\n'
                fint_params += 'gn_args: "clang_ml_inliner = false"\n'
                api.file.write_text(
                    "write fint params",
                    checkout.root_dir / fint_params_path,
                    fint_params,
                )

                # Build Fuchsia for each target architecture...
                build_result = api.build.with_options(
                    checkout=checkout,
                    fint_params_path=fint_params_path,
                    build_dir=checkout.root_dir.joinpath("out", target_arch),
                    clean_check=False,
                )

                with api.step.nest("check binary sizes") as presentation:
                    presentation.logs["binary sizes JSON output"] = api.json.dumps(
                        build_result.binary_sizes, indent=2
                    )

                # ...and extract IR from the generated object files.
                docker_run(
                    api,
                    pyscript="compiler_opt/tools/extract_ir.py",
                    step_name="extract ir",
                    args=[
                        "--cmd_filter=^-O2|-Os|-Oz$",
                        f"--input={build_result.compdb_path}",
                        "--input_type=json",
                        f"--llvm_objcopy_path={build_result.tool('llvm-objcopy')}",
                        f"--output_dir={corpus_dir / target_arch}",
                    ],
                )

        # Combine IR from both builds into a single training corpus.
        docker_run(
            api,
            pyscript="compiler_opt/tools/combine_training_corpus.py",
            step_name="combine training corpus",
            args=[
                f"--root_dir={corpus_dir}",
            ],
        )

        # Upload the corpus to CAS.
        api.cas_util.upload(corpus_dir, step_name="archive corpus")

    with api.step.nest("trace"):
        vocab_dir = api.path.start_dir / "vocab"
        api.file.ensure_directory("ensure vocab directory", vocab_dir)

        default_trace_dir = api.path.start_dir / "default_trace"

        docker_run(
            api,
            pyscript="compiler_opt/tools/generate_default_trace.py",
            step_name="generate default trace",
            args=[
                f"--data_path={corpus_dir}",
                f"--output_path={default_trace_dir}",
                f"--gin_files={'/ml-compiler-opt/compiler_opt/rl/inlining/gin_configs/common.gin'}",
                f"--gin_bindings=clang_path='{clang_dir.joinpath('bin', 'clang')}'",
                f"--gin_bindings=llvm_size_path='{clang_dir.joinpath('bin', 'llvm-size')}'",
                "--sampling_rate=0.2",
            ],
        )

        api.cas_util.upload(default_trace_dir, step_name="archive default trace")

        docker_run(
            api,
            pyscript="compiler_opt/tools/sparse_bucket_generator.py",
            step_name="generate default vocab",
            args=[
                f"--gin_files={'/ml-compiler-opt/compiler_opt/rl/inlining/gin_configs/common.gin'}",
                f"--input={default_trace_dir}",
                f"--output_dir={vocab_dir}",
            ],
        )

        api.cas_util.upload(vocab_dir, step_name="archive default vocab")

    with api.step.nest("warmstart"):
        # Train a behavioral cloning model based on the above trace, that mimics
        # default inlining behavior. This is the 'warmstart' model.
        warmstart_dir = api.path.start_dir / "warmstart"

        docker_run(
            api,
            pyscript="compiler_opt/rl/train_bc.py",
            step_name="train bc",
            args=[
                f"--root_dir={warmstart_dir}",
                f"--data_path={default_trace_dir}",
                f"--gin_files={'/ml-compiler-opt/compiler_opt/rl/inlining/gin_configs/behavioral_cloning_nn_agent.gin'}",
            ],
        )

        api.cas_util.upload(warmstart_dir, step_name="archive warmstart")

    with api.step.nest("train"):
        # Starting from the 'warmstart' model, train the optimized model.
        model_dir = api.path.start_dir / "model"

        args = [
            f"--root_dir={model_dir}",
            f"--data_path={corpus_dir}",
            f"--gin_files={'/ml-compiler-opt/compiler_opt/rl/inlining/gin_configs/ppo_nn_agent.gin'}",
            f"--gin_bindings=clang_path='{clang_dir.joinpath('bin', 'clang')}'",
            f"--gin_bindings=llvm_size_path='{clang_dir.joinpath('bin', 'llvm-size')}'",
            f"--gin_bindings=train_eval.warmstart_policy_dir='{warmstart_dir / 'saved_policy'}'",
        ]
        if props.num_iterations:
            args.append(
                f"--gin_bindings=train_eval.num_policy_iterations={int(props.num_iterations)}"
            )
        # TODO(phosek): Re-enable this once the issue we've seen on bots is addressed.
        args.append("--gin_bindings=train_eval.use_random_network_distillation=False")

        docker_run(
            api,
            pyscript="compiler_opt/rl/train_locally.py",
            step_name="train locally",
            args=args,
            env={"TF_CPP_MIN_LOG_LEVEL": 2},
        )

        api.cas_util.upload(model_dir, step_name="archive model")

    saved_policy_dir = model_dir / "saved_policy"

    if api.buildbucket.build.builder.bucket == "prod":
        api.cipd_util.upload_package(
            "fuchsia/model/inlining",
            saved_policy_dir,
            search_tag={"git_revision": revision},
        )

    # TODO(phosek): We would ideally read the targets from Jiri manifest as:
    # for platform in package_data[0]["platforms"]:
    #     if not platform in PLATFORM_TO_TARGET:
    #         continue
    with api.step.nest("generate"):
        for platform in sorted(PLATFORM_TO_TARGET.keys()):
            with api.step.nest(platform):
                inliner_model_dir = api.path.start_dir.joinpath(
                    "inliner_model", platform
                )
                api.file.ensure_directory("ensure dir", inliner_model_dir)

                docker_run(
                    api,
                    pyscript="../usr/local/bin/saved_model_cli",
                    step_name="aot compile model",
                    args=[
                        "aot_compile_cpu",
                        "--multithreading=false",
                        f"--dir={saved_policy_dir}",
                        "--tag_set=serve",
                        "--signature_def_key=action",
                        f"--output_prefix={inliner_model_dir / 'InlinerSizeModel'}",
                        "--cpp_class=llvm::InlinerSizeModel",
                        f"--target_triple={PLATFORM_TO_TARGET[platform]}",
                    ],
                )

                with api.docker.create(
                    REGISTRY_URL + "/" + IMAGE_NAME + ":" + ML_COMPILER_OPT_TAG,
                    [
                        "python3",
                        "-c",
                        "import importlib.util, os; print(os.path.dirname(importlib.util.find_spec('tensorflow').origin))",
                    ],
                ) as container:
                    tensorflow_aot_path = api.docker.start(
                        container, attach=True, stdout=api.raw_io.output_text()
                    ).stdout.strip()
                    api.docker.copy(
                        container,
                        [
                            f"{tensorflow_aot_path}/include",
                            f"{tensorflow_aot_path}/xla_aot_runtime_src",
                        ],
                        inliner_model_dir / "tensorflow",
                    )

                api.cas_util.upload(
                    inliner_model_dir, step_name="archive inliner model"
                )

                if api.buildbucket.build.builder.bucket == "prod":
                    api.cipd_util.upload_package(
                        f"fuchsia/model/inlining/{platform}",
                        inliner_model_dir,
                        search_tag={"git_revision": revision},
                    )

    with api.step.nest("evaluate"):
        # Evaluate trained policy on a corpus.
        performance_report = api.path.start_dir / "performance_report"
        docker_run(
            api,
            pyscript="compiler_opt/tools/generate_default_trace.py",
            step_name="evaluate trained policy",
            args=[
                f"--data_path={corpus_dir}",
                f"--policy_path={model_dir / 'saved_policy'}",
                f"--output_performance_path={performance_report}",
                f"--gin_files={'/ml-compiler-opt/compiler_opt/rl/inlining/gin_configs/common.gin'}",
                f"--gin_bindings=clang_path='{clang_dir.joinpath('bin', 'clang')}'",
                f"--gin_bindings=llvm_size_path='{clang_dir.joinpath('bin', 'llvm-size')}'",
                "--sampling_rate=0.2",
            ],
            env={"TF_CPP_MIN_LOG_LEVEL": 2},
        )
        api.cas_util.upload(
            api.path.start_dir,
            [performance_report],
            step_name="archive performance report",
        )

        # Evaluate trained policy on Fuchsia.

        # TODO(phosek): This logic is largely the same as the one during the corpus
        # extraction, we should deduplicate it by extracting the code into a subroutine.
        with api.step.nest("llvm"):
            inliner_model_dir = api.path.start_dir.joinpath(
                "inliner_model", "linux-amd64"
            )

            # Run CMake to configure Clang build.
            build_dir = api.path.start_dir / "llvm_build"
            api.step(
                "configure",
                [
                    cipd_dir.joinpath("bin", "cmake"),
                    "-G",
                    "Ninja",
                    f"-DCMAKE_MAKE_PROGRAM={cipd_dir / 'ninja'}",
                    "-DCMAKE_INSTALL_PREFIX=",
                    "-DLLVM_ENABLE_LTO=OFF",
                    f"-DCMAKE_TOOLCHAIN_FILE={checkout.root_dir.joinpath('scripts', 'clang', 'ToolChain.cmake')}",
                    f"-DLINUX_x86_64-unknown-linux-gnu_SYSROOT={sysroot_dir}",
                    f"-DLINUX_aarch64-unknown-linux-gnu_SYSROOT={sysroot_dir}",
                    f"-DFUCHSIA_SDK={cipd_dir / 'sdk'}",
                    f"-DTENSORFLOW_AOT_PATH={inliner_model_dir / 'tensorflow'}",
                    f"-DLLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL={inliner_model_dir / 'InlinerSizeModel.h'}",
                    f"-DLLVM_OVERRIDE_MODEL_OBJECT_INLINERSIZEMODEL={inliner_model_dir / 'InlinerSizeModel.o'}",
                    "-DLLVM_RAEVICT_MODEL_PATH=none",
                    "-DCMAKE_INSTALL_RPATH_USE_LINK_PATH=ON",
                    f"-DCMAKE_SYSROOT={sysroot_dir}",
                    "-C",
                    llvm_src_dir.joinpath(
                        "clang",
                        "cmake",
                        "caches",
                        "Fuchsia-stage2.cmake",
                    ),
                    "-S",
                    llvm_src_dir / "llvm",
                    "-B",
                    build_dir,
                ],
            )

            # Build the Clang distribution.
            api.step(
                "build",
                [
                    cipd_dir / "ninja",
                    "-C",
                    build_dir,
                    "distribution",
                ],
            )

            # Create the Clang toolchain.
            clang_dir = api.path.start_dir / "clang"
            with api.context(env={"DESTDIR": clang_dir}):
                api.step(
                    "install",
                    [
                        cipd_dir / "ninja",
                        "-C",
                        build_dir,
                        "install-distribution",
                    ],
                )
            api.step(
                "generate runtimes.json",
                cmd=[
                    "vpython3",
                    "-vpython-spec",
                    api.tensorflow.vpython_spec,
                    checkout.root_dir.joinpath(
                        "scripts", "clang", "generate_runtimes.py"
                    ),
                    f"--clang-prefix={clang_dir}",
                    f"--sdk-dir={sdk_dir}",
                    f"--build-id-dir={clang_dir.joinpath('lib', '.build-id')}",
                ],
                stdout=api.raw_io.output_text(
                    leak_to=clang_dir.joinpath("lib", "runtime.json"),
                    add_output_log=True,
                ),
            )

            # Use the just built Clang toolchain to build Fuchsia.
            checkout.clang_toolchain_dir = clang_dir

            # Upload the Clang binaries to CAS.
            api.cas_util.upload(clang_dir, step_name="archive clang")

        with api.step.nest("fuchsia"):
            corpus_dir = api.path.start_dir / "corpus"

            for target_arch, fint_params_path in props.fint_params_paths.items():
                with api.step.nest(target_arch):
                    # Build Fuchsia for each target architecture...
                    build_result = api.build.with_options(
                        checkout=checkout,
                        fint_params_path=fint_params_path,
                        build_dir=checkout.root_dir.joinpath("out", target_arch),
                        clean_check=False,
                    )

                    with api.step.nest("check binary sizes") as presentation:
                        presentation.logs["binary sizes JSON output"] = api.json.dumps(
                            build_result.binary_sizes, indent=2
                        )


def docker_run(api, pyscript, args=None, env=None, **kwargs):
    cmd = [
        "python3",
        "/ml-compiler-opt/" + pyscript,
    ]
    if args:
        cmd.extend(args)
    updated_env = {"PYTHONPATH": "/ml-compiler-opt"}
    if env:
        updated_env.update(env)
    api.docker.run(
        image=REGISTRY_URL + "/" + IMAGE_NAME + ":" + ML_COMPILER_OPT_TAG,
        cmd_args=cmd,
        dir_mapping=[
            (api.path.start_dir, api.path.start_dir),
            (
                api.path.start_dir / "vocab",
                "/ml-compiler-opt/compiler_opt/rl/inlining/vocab",
            ),
        ],
        cwd="/ml-compiler-opt",
        env=updated_env,
        **kwargs,
    )


def GenTests(api):
    yield (
        api.buildbucket_util.test("default", bucket="prod")
        + api.properties(
            manifest="fuchsia",
            remote="https://fuchsia.googlesource.com/fuchsia",
            fint_params_paths={
                "arm64": "specs/clang-ml-training-arm64.fint.textproto",
                "x64": "specs/clang-ml-training-x64.fint.textproto",
            },
        )
    )
    yield (
        api.buildbucket_util.test("num_iterations", bucket="prod")
        + api.properties(
            manifest="fuchsia",
            remote="https://fuchsia.googlesource.com/fuchsia",
            fint_params_paths={
                "arm64": "specs/clang-ml-training-arm64.fint.textproto",
                "x64": "specs/clang-ml-training-x64.fint.textproto",
            },
            num_iterations=5000,
        )
    )