MLGO 的基础设施 - 一种基于机器学习的编译器优化框架。

MLGO 是一个在 LLVM 中系统性集成机器学习技术的框架。它用机器学习模型替代了 LLVM 中人工编写的优化启发式规则。目前，MLGO 框架支持两种优化：

1. 针对代码大小的内联（LLVM RFC）；
2. 针对性能的寄存器分配（LLVM RFC）。

这些编译器组件均已合并到 LLVM 主仓库中。本仓库则包含了 MLGO 的训练基础设施及相关工具。

我们目前使用两种不同的机器学习算法来训练策略：策略梯度和进化策略。现阶段，本仓库仅支持策略梯度训练；进化策略训练的实现已在我们的开发计划中。

请参阅此 演示，了解如何从零开始使用策略梯度训练自定义的针对代码大小的内联策略；或参阅此 演示，了解如何训练自定义的针对性能的寄存器分配策略。

如需了解更多关于 MLGO 的信息，请查阅我们的论文《MLGO：一种基于机器学习的编译器优化框架》（arXiv 链接）。

有关如何参与本项目的更多信息，请参阅 贡献指南。

预训练模型

我们偶尔会发布可直接与 LLVM 配合使用的预训练模型。这些模型以 GitHub 发布的形式提供，并按照 [任务]-[主版本].[次版本] 的命名方式。版本号遵循语义化规范：主版本对应于 LLVM/编译器端的重大变更，而次版本则对应于与编译器无关的模型更新。

在构建 LLVM 时，可以设置 -DLLVM_INLINER_MODEL_PATH 标志，指定内联模型的路径。若将该标志设为 download，CMake 将自动从 GitHub 下载最新且兼容的模型。其他可能的取值如下：

# 模型位于 /tmp/model 目录下，即存在 /tmp/model/saved_model.pb 文件以及训练过程中生成的其他 TensorFlow SavedModel 文件。
-DLLVM_INLINER_MODEL_PATH=/tmp/model

# 下载最新且兼容的模型。
-DLLVM_INLINER_MODEL_PATH=download

先决条件

当前系统的假设条件如下：

• 最新的 Ubuntu 发行版，例如 22.04；
• Python 3.10.x 或 3.11.x；
• 对于目前唯一支持的本地训练模式，建议使用高性能工作站（例如具有 96 个硬件线程的设备）。

训练过程假定 Clang 已以 ML“开发模式”构建。请参考以下资源：

• LLVM 官方文档；
• 构建脚本 buildbot_init.sh。

模型训练的具体先决条件如下：

使用 Pipenv：

pip3 install pipenv

安装实际依赖项：

./versioned_pipenv sync --system --categories "packages dev-packages ci"

请注意，上述命令必须在仓库根目录下执行，因为运行时需要当前工作目录中存在 Pipfile.lock 文件。

该命令还将安装所有依赖包，包括开发包（dev-packages 类别）以及仅在 CI 环境中使用的包（ci 类别）。如果您不需要这些内容，可以在 categories 参数中将其排除。

此外，若要运行测试脚本 run_tests.sh，还需安装：

sudo apt-get install virtualenv

需要注意的是，构建 LLVM 的“发布模式”同样需要相同的 TensorFlow 包。

文档

一篇使用 Fuchsia 作为代码库、从中提取语料并训练模型的端到端 演示。

如何添加新功能指南。 可扩展性模型。

ml-compiler-opt 快速上手指南

ml-compiler-opt 是 MLGO（机器学习引导的编译器优化框架）的训练基础设施。它用于在 LLVM 中训练机器学习模型，以替代传统的人工优化启发式规则，目前主要支持“基于大小的内联优化”和“基于性能的寄存器分配优化”。

环境准备

系统要求

• 操作系统：推荐近期版本的 Ubuntu（如 22.04）。
• Python 版本：3.10.x 或 3.11.x。
• 硬件建议：本地训练需要高性能工作站（推荐 96 个硬件线程或以上），以获得合理的训练速度。

前置依赖

1. LLVM/Clang 构建： 需要构建开启 ML "development-mode" 的 Clang。具体配置请参考 LLVM 官方文档 或项目中的 buildbot 初始化脚本。

2. 基础工具： 安装 pipenv 和 virtualenv（用于运行测试）：

   pip3 install pipenv
   sudo apt-get install virtualenv
   

3. Python 依赖包： 在项目根目录下执行以下命令同步依赖（包含开发包和 CI 包）。如需仅安装核心包，可调整 --categories 参数。 注：国内用户若遇网络问题，可配置 pip 使用清华或阿里镜像源。

   ./versioned_pipenv sync --system --categories "packages dev-packages ci"
   

安装步骤

本仓库主要提供训练基础设施，核心编译器组件已集成在主 LLVM 仓库中。以下是获取并配置环境的流程：

1. 克隆代码库：

   git clone https://github.com/google/ml-compiler-opt.git
   cd ml-compiler-opt
   

2. 安装 Python 环境： 确保已在项目根目录执行上述“前置依赖”中的 versioned_pipenv sync 命令。

3. 配置 LLVM 构建（可选 - 使用预训练模型）： 如果您不想从头训练，可以直接在构建 LLVM 时集成预训练模型。

   • 自动下载最新兼容模型：

     cmake -DLLVM_INLINER_MODEL_PATH=download ...
     

   • 指定本地模型路径：

     cmake -DLLVM_INLINER_MODEL_PATH=/tmp/model ...
     

     (注：路径下需包含 saved_model.pb 等 TensorFlow 生成的文件)

基本使用

最简单的使用方式是通过端到端演示来训练一个策略模型。项目提供了两个主要方向的 Demo：

示例：训练“基于大小的内联”策略 (Inlining-for-size)

该示例展示如何从零开始，使用策略梯度（Policy Gradient）算法，基于 Fuchsia 代码库提取语料并训练模型。

1. 查看演示文档： 详细步骤请参阅项目内的演示文档：

   docs/inlining-demo/demo.md
   

2. 核心流程概览：

   • 数据收集：使用特定的 clang 构建编译目标代码库，收集优化决策数据。
   • 特征提取：从编译过程中提取特征。
   • 模型训练：运行训练脚本，使用 Policy Gradient 算法更新模型权重。
   • 评估与应用：将训练好的模型导出，并在 LLVM 构建中通过 -DLLVM_INLINER_MODEL_PATH 加载验证。

示例：训练“基于性能的寄存器分配”策略 (Regalloc-for-performance)

类似地，针对寄存器分配的优化训练流程如下：

1. 查看演示文档：

   docs/regalloc-demo/demo.md
   

2. 执行训练： 按照文档指引，配置相应的优化任务类型，复用上述数据收集和训练管线即可。

提示：更多关于如何添加新特征或扩展系统的细节，请参考 docs/adding_features.md 和 docs/extensibility.md。