拆解ml-compiler-opt的4步训练流水线：从默认轨迹采集到PPO强化学习

拆解ml-compiler-opt的4步训练流水线：从默认轨迹采集到PPO强化学习

【免费下载链接】ml-compiler-optInfrastructure for Machine Learning Guided Optimization (MLGO) in LLVM.项目地址: https://gitcode.com/gh_mirrors/mlc/ml-compiler-opt

ml-compiler-opt 是 Google 开源的机器学习引导编译器优化（MLGO，Machine Learning Guided Optimization）训练基础设施，用于把 LLVM 编译器中人工编写的优化启发式替换为机器学习模型。它目前已支持两大优化场景：面向代码体积的内联优化（inlining-for-size，对应-Oz）和面向性能的寄存器分配优化（regalloc-for-performance）。对想尝试"用 AI 训练编译器"的新手来说，本文带你完整拆解它的 4 步训练流水线：语料提取 → 默认轨迹采集 → 行为克隆热启动 → PPO 强化学习闭环训练。

ml-compiler-opt 如何把编译器变成"可训练"的环境？

在传统编译器里，"这个函数要不要内联？"这类决策由几十行精心调校的经验公式（启发式）决定。ml-compiler-opt 的思路是：

1. 观测（Observation）：编译器每次做决策时，把当前调用点周围几十个特征（如调用者/被调用者基本块数量、内联成本估计等）打包成观测向量；
2. 动作（Action）：由神经网络模型输出决策（内联/不内联）；
3. 奖励（Reward）：编译完成后，用最终二进制大小与启发式基线对比——变小是正奖励，变大是负奖励。

编译过程由此变成一个强化学习环境。项目根目录的 README.md 指出，LLVM 主仓提供"开发模式"clang（可通过 TFLite 从命令行热切换策略），而本仓库负责训练循环和相关工具。

四步流水线总览

第 1 步：提取训练语料——从真实代码中"挖矿"

强化学习需要海量训练样本。ml-compiler-opt 的语料来自真实项目：官方 Demo 以 Fuchsia 操作系统为例（任何能用 clang 构建、可生成compile_commands.json的项目都可行），模块越多越好。

关键机制是构建时开启clang_embed_bitcode=true：对象文件中会嵌入优化前的 LLVM 字节码和 clang 命令行，extract_ir工具即可从中把每个模块还原出来，形成语料库。语料加载逻辑见 compiler_opt/rl/corpus.py。

💡 新手提示：语料提取只需编译一次目标项目，是整个流水线中"一次性"的准备工作。

第 2 步：采集默认轨迹并生成特征词表

采集默认轨迹：用工具 compiler_opt/tools/generate_default_trace.py 驱动"启发式版"clang 重新编译整个语料库，把编译器每次决策时的观测值、动作、奖励记录成 tfrecord 轨迹文件。它支持--num_workers并行编译和--sampling_rate采样率（如 0.2 表示只处理 20% 的模块），并行调度基于 compiler_opt/distributed/ 下的本地工作池实现。

生成特征词表：观测特征大多是连续数值（如"被调用者有多少个基本块"），无法直接喂给网络。工具 compiler_opt/tools/generate_vocab.py 会对每个特征按其分布做1000 分位分桶（quantile bucketization），生成.buckets文件存入 compiler_opt/rl/inlining/vocab/——这里你能看到node_count.buckets、threshold.buckets等 30 多个特征的分桶文件。后续训练时，特征值会被映射为分桶序号，作为网络的稀疏特征输入。

⚠️ 词表在"特征集合或特征分布发生变化"时需要重新生成，官方 Demo 将其列为可选步骤。

第 3 步：行为克隆热启动——先学会"模仿启发式"

PPO 如果从随机策略开始，早期大量样本都是浪费。ml-compiler-opt 引入行为克隆（Behavioral Cloning）热启动：

• 训练入口：compiler_opt/rl/train_bc.py
• 配置文件：compiler_opt/rl/inlining/gin_configs/behavioral_cloning_nn_agent.gin，使用BCAgentConfig+ QNetwork 结构（隐藏层 40-40-20）

它直接读取第 2 步的默认轨迹，让网络模仿启发式的每一个内联决策。产出的模型就是"热启动模型"——它不要求比启发式更好，只要求足够接近，为 PPO 提供一个良好的初始策略，避免强化学习冷启动阶段的震荡。

第 4 步：PPO 强化学习——闭环训练循环

主训练入口是 compiler_opt/rl/train_locally.py，配合配置 compiler_opt/rl/inlining/gin_configs/ppo_nn_agent.gin（PPOAgentConfig+ ActorDistributionNetwork）。其核心循环在train_eval函数中非常直观：

1. 保存策略：将当前 PPO 策略落盘（saved_collect_policy）；
2. 采集数据：LocalDataCollector调度一批并行 worker，用该策略真实重编译语料中的模块，产出"观测-动作-奖励"轨迹，并维护每个模块的奖励统计（compiler_opt/rl/local_data_collector.py）；
3. PPO 训练：Trainer（compiler_opt/rl/trainer.py）在轨迹上做若干轮（num_iterations）策略梯度更新；
4. 重复：直到累计策略迭代达到num_policy_iterations上限。

配置文件中的关键超参数值得新手关注：

train_eval.num_policy_iterations = 3000 # 策略迭代轮数 train_eval.num_modules = 100 # 每轮采集的模块数 train_eval.num_iterations = 300 # 每轮 PPO 更新次数 PPOAgent.importance_ratio_clipping = 0.2 # PPO 经典的裁剪系数 PPOAgent.entropy_regularization = 0.003 # 熵正则，鼓励探索 PPOAgent.adaptive_kl_target = 0.01 # 自适应 KL 约束

如何判断训练好坏？启动 TensorBoard 观察reward_distribution：奖励均值和分位数的正负，代表模型相对启发式的大小改进/回退情况（正奖励 = 改进，负奖励 = 回退），该监控逻辑实现在 compiler_opt/rl/data_collector.py。官方提示完整 PPO 训练约需半天时间（建议 96 核左右的工作站）。

训练完成之后：把模型"装回"编译器

PPO 训练产出的是 TensorFlow SavedModel。部署时把它覆盖进 LLVM 源码树的llvm/lib/Analysis/models/inliner/，再以release 模式重新构建 clang（去掉 TFLite 依赖、策略静态嵌入），最终用-mllvm -enable-ml-inliner=release启用，再对比优化前后的体积报告即可验证收益。

快速上手：环境要求与完整 Demo

• 系统：Ubuntu（如 20.04）、Python 3.8/3.9/3.10；
• 依赖：pip3 install pipenv && pipenv sync --system（依赖声明在 Pipfile）；
• 需构建 TFLite 并以"开发模式"构建 LLVM（辅助脚本见 buildbot/build_tflite.sh 与 buildbot/buildbot_init.sh）。

项目提供了端到端教程：

• 内联策略训练 Demo：docs/inlining-demo/demo.md
• 寄存器分配策略训练 Demo：docs/regalloc-demo/demo.md

如果只想看流水线骨架而不跑完整 Demo，直接按本文 4 个步骤的顺序阅读对应源码，是最快的理解路径。

小结

ml-compiler-opt 的 4 步流水线可以概括为一句话：用真实代码造数据，用启发式轨迹打地基，用 PPO 在真实编译反馈上持续精进。除了内联和寄存器分配，它的框架设计（compiler_opt/rl/下的 env、env 配置、特征词表机制）也天然支持接入更多优化点；仓库中compiler_opt/es/目录还为计划中的进化策略（Evolution Strategies）训练预留了空间。想要给 LLVM"装上大脑"，这就是你的起点。

【免费下载链接】ml-compiler-optInfrastructure for Machine Learning Guided Optimization (MLGO) in LLVM.项目地址: https://gitcode.com/gh_mirrors/mlc/ml-compiler-opt