Faster sorting algorithms discovered using deep reinforcement learning (AlphaDev)
这是什么
排序是被调用最频繁的库函数之一,libc++ 里小规模排序(sort3/4/5)本质上是手工优化了几十年的 branchless sorting network。这篇 2023 年 6 月的 Nature 论文提出 AlphaDev:不在源码层做程序合成,而是直接在 x86 汇编指令层搜索——因为编译器输出的指令序列才是真正决定 latency 的东西,而且指令级还残留着人类没榨干的优化空间。
AlphaDev 是 AlphaZero 的直系后代(作者列表里有 Silver、Schrittwieser、Hubert)。它把"逐条追加汇编指令"当成游戏动作,以正确性 + latency 作为奖励。搜索空间比围棋还大(论文估计枚举量级 >10^700 同级),且一条错指令就让整个程序作废。最终发现的 sort3/4/5 被人工逆向成 C++ 合入 LLVM libc++(review D118029),是这些子程序十多年来第一次改动;论文还把方法迁移到 protobuf 的 VarInt 反序列化,得到 branchless 版本。
在"AI 优化 AI/软件基础设施"这条线上,它是 AlphaTensor(2022,矩阵乘法)之后、AlphaEvolve(2025,LLM 驱动)之前的关键一环:证明 RL 搜索的产物可以通过工业级 code review 落进所有软件依赖的底层库,而不只是停留在论文表格里。
机制与做法
AssemblyGame:把写程序变成单人游戏
状态 S_t = ⟨P_t, Z_t⟩:P_t 是当前已生成的程序,Z_t 是把程序跑在预定义输入上之后的寄存器/内存状态。每一步动作是追加一条合法 x86 指令(mov/cmp/cmovX/jX 等,AT&T 语法)。正确性奖励:对 sort3 用全部长度为 3 的未排序排列做输入,输出与期望排序结果逐一比对;latency 奖励:branchless 程序里长度与 latency 强相关,直接用程序长度做代理;有分支的 VarSort 则用实测 latency。
动作空间做了强剪枝:内存按递增顺序读、寄存器递增分配、禁止对内存做 cmp/cmov、每个内存位置只读写一次、不许用未初始化寄存器、不许连续两条 cmp。这些规则把天文数字的空间砍到可搜索。
AlphaZero 骨架 + 面向代码的表示网络 + 双 value head
学习算法就是 AlphaZero:网络输出 policy 和 value 引导 MCTS,policy 向 MCTS 访问计数回归。表示网络分两块:一个 MultiQuery Transformer 编码指令序列(opcode/operand 做 one-hot 再映射到 embedding),一个 MLP 编码 CPU 状态(寄存器和内存内容),让网络能预测程序对机器状态的作用。
针对 latency 优化的关键设计是双 value head:一个预测正确性,一个以实测 latency 为 Monte Carlo target 直接预测 latency。论文报告这比单 value head 明显更好——因为实测 latency 太贵,AlphaDev 只对不到 0.002% 的生成程序做了真实测量,其余全靠 latency head 泛化。训练用 TPU v3(每 core batch 1024,至多 16 core)+ 至多 512 个 TPU v4 actor,最坏情况约 2 天收敛。
latency 怎么测、和什么比
latency 测量用一套跨机器的 benchmark 服务:每次评估做 1000 次测量,每次在 10000 个随机输入上跑;计时用 CPU_CLK_UNHALTED.CORE 性能计数器;取第 5 百分位数作为结果(假设噪声——cache miss、抢占——都是单向使程序变慢的)。训练时用 10 台机器,最终评测用 100 台。
对照基线是把 SOTA 随机超优化器(Schkufza 风格 MCMC)改造成的 AlphaDev-S。冷启动版在固定长度 sort 上找不到最优解、在 VarSort 上全灭;用近似最优程序热启动后,在 branchless 的 sort3/4/5 上能追平 AlphaDev(而且算力更省),但在长度与 latency 脱钩的 VarSort 上仍然输——因为它无法承受对每次变异做实测 latency,只能优化长度。探索量对比:AlphaDev 最多探索约 1200 万个程序,AlphaDev-S 最多约 31 万亿个。
关键结果
- sort3 找到 17 条指令的程序(人类基准 18 条),并用约 10^32 规模的剪枝穷举(3 天多)证明 17 条是下界;sort5 也少 1 条,sort4 追平 SOTA;规模化后 sort6/7/8 分别省 3/2/1 条指令。
- 核心新颖点是两个可复用的局部优化:swap move(利用前序 comparator 保证 B≤C,把 min(A,B,C) 简化为 min(A,B),省 1 条指令)和 copy move(利用 D≥min(A,C) 的不变量用一次 copy 替代比较,省 1 条指令),模式在不同规模的 sorting network 中反复出现。
- VarSort4 的结构是质变而非省指令:不再按长度分派到不同 sorting network,而是先无条件调 sort3,长度大于 3 时再对剩余元素跑简化版 sort4,latency 收益主要来自这个新算法结构。
- 合入 LLVM libc++ 后:长度 5 的序列提升"最高 70%"(相对旧的基于插入排序的实现),超过 25 万元素的大序列约 1.7%,覆盖 uint32/uint64/float,在 ARMv8、Skylake、Zen 2 上验证。
- 迁移到 protobuf VarInt 反序列化:发现 branchless 且更短的方案,单值输入上比人类基准快约 3 倍;还学到把两个操作融合为单条指令的 assignment move。
- AlphaDev 对不到 0.002% 的候选程序做真实 latency 测量,其余靠 latency value head;t-SNE 显示随机超优化基线密集聚在种子程序附近,AlphaDev 覆盖的程序空间更分散。
实证核查
CC=clang bazel test :sort_functions_test 可本地验证;指令数与论文、README 逐一吻合。合入 LLVM 的 review 是 D118029,公开可查。与我们方向的关系
对课题组 ai4ai 方向,这篇的可借鉴点不在排序本身,而在"如何让 RL 优化一个昂贵的真实目标":双 value head(便宜代理 + 少量真实测量训练出的 latency 预测头)把真实评测压到 0.002%,这个模式可以平移到任何评测昂贵的搜索问题(kernel 调优、编译 pass 序列、超参搜索)。另一个值得记的教训是基线对比:热启动超优化器在 branchless 场景追平了 RL,说明"RL 赢"高度依赖任务结构(长度与目标脱钩、需要长程 credit assignment 时才赢)。
作为 AlphaTensor → AlphaDev → FunSearch/AlphaEvolve 谱系的一环,它标志着范式切换前夜:AlphaDev 用领域定制的 RL 环境 + 自训练网络,每个新问题都要重建环境;两年后 AlphaEvolve 用通用 LLM + 进化搜索覆盖了同类问题且不用训练。读它的价值是理解这条路线为什么被替代,以及哪些组件(真实 latency 反馈、剪枝规则、正确性测试即奖励)在 LLM 时代仍然保留。