用哪些硬件来考 DARPA MOCHA 的编译器
MOCHA 想让编译器快速适配从没见过的异构加速器,那测试基准怎么选?SEI 列出五个筛选维度与九项候选硬件,并请各团队自己推举目标,最终落在 AMD RDNA 4 与 Tenstorrent Tensix 上。
中文
复制

现代计算机不再围绕单一处理器构建。今天一台像样的机器是一个异构组合体:CPU、GPU,以及一个不断扩张的专用加速器家族,分别面向机器学习、信号处理和网络。想从这个组合体里榨出好性能并不容易,而要在编译器从未见过的硬件上快速做到这一点,就更难。DARPA MOCHA——即 DARPA 的异构架构机器学习与优化引导编译器项目——正是为弥合这一差距而设。SEI 人工智能部门的先进计算实验室过去几个月一直在思考一个将决定该项目未来两年走向的问题:项目的编译器应该拿哪些硬件来做测试?本文介绍我们处理这个问题的思路:MOCHA 想做什么,SEI 在其中扮演什么角色,我们如何筛选候选硬件,候选清单以及我们如何对每个候选逐一施压测试,最终的选择又落在了哪里。
自动化计算机优化
MOCHA 是 DARPA 信息处理技术办公室的项目,由 Dr. Howard Shrobe 管理。推动这项工作的,是一个众所周知的痛点:传统编译器在设计时并没有考虑为 CPU、GPU 和应用加速器构成的异构混合体生成高效机器码。要利用一款新加速器,开发者通常得手写专用代码,并依赖厂商调优过的库。这条路走得通,但既慢又贵,还悄悄助长了厂商锁定:一旦应用是针对某个专有库写的,想换到别的硬件上就意味着重写。把编译器扩展成支持一种真正全新的计算单元,是只有编译器专家才能完成的手工活。它耗时、易错,而且跟不上新硅片出现的节奏。MOCHA 的假设是,数据驱动方法、机器学习和先进优化可以加速这一过程,让编译器能够以最少的人工介入迅速适配新硬件。一个关键洞见是:目标硬件的性能模型驱动着编译的每一步,而手工构建这些模型正是核心瓶颈。如果这些模型可以改为通过在真实硬件上测量生成的代码、并挖掘架构文档来自动生成,那么支持一款新设备的成本就会大幅下降。DARPA MOCHA 的一条指导原则是 ALARA——把人的参与保持在合理可及的最低水平。ALARA 体现了 MOCHA 的双重侧重:既要快速让编译器支持新硬件,也要让编译能够跨越异构硬件。在某一款新芯片上跑得快还不够;这个项目关心的是,要同时覆盖一大类各不相同的计算单元,需要投入多少人力。
SEI 的角色
SEI 的高级计算实验室隶属于 AI Division,为政府团队提供支持,重点在测试与评估。政府团队由 DARPA 项目经理和若干系统工程技术助理(SETA)组成。具体来说,我们负责收集和评估各种选项,把项目经理做决定所需的信息交到他手上,由他决定哪些计算单元进入项目。随后我们搭建并维护评估机器,把选定的计算单元集成进去,同时为 MOCHA 建立测量方法,以便公平比较各参与团队的结果。参与团队负责开发编译器技术;我们的任务是给他们一组特征明确、有代表性且难度合适的目标去攻,并尽可能保证评估本身的有效性。DARPA 计划在项目结束时让 MOCHA 涵盖六种不同的计算类型。计算类型的界定不只看硬件架构,还要看指令集和编程模型是否不同。按这个定义,数据中心 GPU、把现场可编程门阵列(FPGA)结构与空间 AI 引擎阵列融合在一起的器件、长向量处理器,以及 RISC-V 加数据流的 AI 加速器,属于四种不同类型——尽管随手一看,后三种可能都会被归为“加速器”。这个项目的要点是展示跨架构和指令集架构(ISA)边界进行快速、低成本的重新定向,因此目标集合中架构的多样性必不可少。还有一个现实约束:所选硬件必须能实际装进评估机器。这台机器是一台工作站级塔式机,围绕 Intel Core Ultra 9 285K 搭建,它本身就带来几种计算类型:CPU 上的 AVX2 SIMD、集成的 Xe GPU,以及一个神经处理单元;此外还有 PCIe 5.0 连接,并已装好一块 NVIDIA RTX 4500 Ada 卡作为基线参考。因此,候选加速器需要以 PCIe 卡的形式提供,并且要适配单台塔式机的机箱、功耗和散热条件。
选择候选加速器的五个因素
五个因素决定了我们的候选名单:可获得性、成熟度、价格可承受性、可编程性,以及能否把设备装进评测机器。可获得性和装机条件淘汰了一批本身很有意思的选项,包括只能以整机服务器形式出货的晶圆级引擎和可重构数据流系统,以及买不到、装不上的纯云端加速器。价格可承受性则让我们对一些比整机其余部分加起来还贵的芯片保持清醒。影响最大的因素是可编程性,具体来说是每个目标平台在编译器和多级中间表示(MLIR)支持上的现状。MOCHA 的参与者绝大多数构建在 LLVM 和 MLIR 生态之上。LLVM 的关键创新是可扩展的 IR,以及让开发者与之交互的工具链。MLIR 是更晚近的创新,它通过在不同抽象层次上定义 IR 扩展了这一能力。它已成为现代编译器基础设施的连接组织,让编译器可以在多个抽象层次上表达计算,并逐步将其下沉到具体设备。如果一个目标平台已经有 MLIR 或 LLVM 路径,参与者就有可能接入,然后把精力放在真正有意思的研究上:可重定向的代码生成、学习型代价模型、优化选择,最终是在异构单元之间划分工作。如果目标平台是一个封闭的黑盒,只能通过厂商预先调优的高层推理栈访问,那么无论投入多少机器学习,MOCHA 编译器能施展的空间都可能非常有限。
开放的低层访问与封闭的厂商库之间的张力贯穿整个项目。针对专有库编程很方便,但这与快速支持新硬件的目标从根本上相冲突,因为库只为厂商已经决定支持的硬件而存在。在评估每个候选时,我们都仔细考察了它的编程模型有多开放,以及是否存在或现实上有可能打通一条基于 MLIR 的、直达硬件的路径。
候选清单
按这些标准筛选后,我们的目标短名单覆盖了 GPU、FPGA 加 AI 引擎的空间架构器件、一款向量处理器、几款各不相同的 AI 加速器,以及网络芯片——足以支撑六种左右真正不同的计算类型:
- AMD Instinct MI350P —— 数据中心 GPU(CDNA 4),通过 ROCm/HIP 编程,经由 rocMLIR 和 Triton 路径拥有成熟的 MLIR 支持。值得注意的是,这是新发布的 PCIe 形态产品,把 OAM 级别的 Instinct 算力带进了标准插槽。
- AMD Versal ACAP —— 集 FPGA 逻辑、空间架构 AI 引擎阵列和 ARM 核于一体的异构器件。开放的 mlir-aie/IRON 工具链及其 Peano LLVM 后端,使 AI 引擎阵列成为一个真正有趣、贴近硬件的 MLIR 目标。
- Intel Data Center GPU Max 1100 —— Xe-HPC GPU,通过 oneAPI/SYCL 编程,可经 Intel Triton XPU 后端和 SPIR-V 接入 MLIR。它是 Max 系列中唯一以 PCIe 卡形式提供的型号。
- Intel Gaudi 3 —— 配备矩阵引擎和 VLIW-SIMD 张量引擎的 AI 加速器。自定义算子通过开放的 TPC-LLVM 编译器用 TPC-C 编写,图编译器会构建基于 MLIR 的融合算子,不过图编译器本身仍是专有的。
- Intel Xe iGPU —— 评测机 CPU 中已集成的核显。它与 Max 1100 共用 oneAPI/SYCL 接口和 MLIR 路径,因此是一个零成本的移植目标。
- NEC SX-Aurora TSUBASA —— PCIe 卡上的经典长向量处理器,拥有上游 LLVM 后端。它为向量化、自动调优和带宽受限的 HPC 算子提供了架构多样性。
- Qualcomm Cloud AI 100 —— 推理加速器,主要路径是 ONNX/PyTorch,但与它“封闭”的名声相反,它还附带一个基于上游 LLVM 的开放编译器,并支持注册底层自定义算子。
- Tenstorrent Blackhole(p100a/p150a) —— 价格亲民且异常开放的 AI 加速器,将 Tensix 核与 RISC-V 结合,配有原生、完全开源的 MLIR 编译器(tt-mlir/tt-forge),可接收来自 PyTorch、JAX 和 ONNX 的模型。
- MangoBoost BoostX DPU 和 GPUBoost RNIC —— 面向网络的部件(一块 SmartNIC/DPU 和一块 RDMA 网卡),为求完整而列入,但没有通用计算的 MLIR 或 LLVM 算子路径。
向执行团队征询意见,以及我们给他们出的难题
在敲定上面的候选清单之前,我们把草稿发给各执行团队,问了两个问题:我们漏掉了什么本该列入的目标?其中哪些是你们的工具链根本支持不了的?回答很坦率,也很有用。各团队指出了哪些部分有真正的 LLVM/MLIR 后端、哪些没有,对价值完全依赖闭源推理栈的目标提出了反对,也直说网络部分不是他们会去做的计算目标。候选目标的筛选始终遵循一条原则:目标要有难度,但不能做不到。太容易的目标——技术栈成熟完善、厂商已经优化好——并不能真正检验 MOCHA 的核心主张,即快速、低成本的适配,因为难做的部分厂商早就做完了。太难的目标——没有文档的黑盒,没有底层编程接口,微架构也无从建模——只会让工作卡死,执行团队把时间耗在跟工具链较劲上,而不是推进研究。理想的目标是:足够开放,能够触及并分析;又和执行团队已经熟悉的东西有足够差异,使得重新定向能真正考验他们的编译器、成本模型和内核生成器。
硬件选型:AMD RDNA 4 GPU 与 Tenstorrent Tensix 核心
在起草那份候选清单和撰写本文之间,进一步筛选出可实际获取的硬件,使范围收窄,项目的首批采购聚焦于两款卡:AMD Radeon AI PRO R9700 和 Tenstorrent Blackhole p150a。基于 AMD ROCm/HIP 技术栈的 GPU 从一开始就注定要占据一席之地。在各执行团队中,它是公认的首选:一款严肃的、非 NVIDIA 的 GPU,拥有真正的 MLIR 路径(经由 rocMLIR 和 Triton),并且与多个执行提案核心处的张量、稀疏、图与成本建模工作直接相关。我们最初选的是新发布的 MI350P,即 AMD 旗舰 Instinct 产品的 PCIe 形态。但 AMD Research 的工程师告知我们,他们自己也在等待 MI350P 芯片,预计要到 2027 年春季才能拿到,远远晚于我们启动第二年评估所需的窗口。于是我们换成了 Radeon AI PRO R9700,一款现在就能买到的 workstation 卡。它能插进标准 PCIe 插槽,使用同一套 ROCm/HIP 编程模型,也能走通同样的 MLIR 和 Triton 路径。它保留了我们想要的 AMD GPU 计算类型,同时又是执行团队今年真能装进机器里的东西。R9700 出人意料地成了一个很好的 MOCHA 目标,原因恰恰触及该项目的核心。由于它基于新架构(RDNA 4),AMD 自家手工调优的汇编库尚未完全覆盖它。其中若干库带有硬编码的支持架构列表,直接把这枚卡排除在外,遇到它时便悄悄回退到慢速路径。真正有效的是编译器路线:Triton 和 MLIR 路径会在运行时为新架构即时生成原生 kernel,而这正是预编译的厂商库失效之处。这就是 MOCHA 论点的微缩版——编译器生成的代码能在手工调优的库无能为力之处重新适配新芯片——也意味着 MOCHA 编译器有真实、可测量的性能空间可以争取,而不是面对一个已被厂商打磨到接近最优的基线。为了在架构上形成对照,我们选择了 Tenstorrent Blackhole p150a。R9700 是运行在成熟 LLVM 后端上的 GPU,而 Blackhole 则截然不同:一组 Tensix 核心搭配通用 RISC-V 核心,通过 Tenstorrent 完全开放、以 MLIR 为核心的编译器栈(tt-mlir 和 tt-forge)编程,该栈经由 StableHLO 接收来自 PyTorch、JAX 和 ONNX 的模型。它的 MLIR 故事可以说是我们评估过的所有方案中最强的。Blackhole 是一个开源编译器的具名目标,其开发完全公开进行,并有文档化的 StableHLO 入口点供执行团队的编译器接入。这里的困难不在于进门,而在于要学习一座定制的 dialect 高塔,而非熟悉的 LLVM 目标模型。
这正是 MOCHA 想要计时、并最终实现自动化的那类重定向挑战。我们原本希望在这第一批中纳入第三种架构类型:AMD Versal。它的 AI engine 阵列是一种空间数据流结构,与 GPU 或 Tensix 阵列都很不一样,而且它在 mlir-aie 中有一套颇具吸引力的开放 MLIR 工具链。最终我们没能找到一款 Versal 器件,既对外暴露 AI engine,又能以 PCIe 卡的形式出货、装进评测机,所以 AI engine 这一类型只能留到项目后期。两块卡给项目带来的是互补的重定向问题,而不是重复的问题。R9700 检验的是在一个成熟生态内做重定向——只是这个生态的库恰好对这款全新的具体架构还不成熟;Blackhole 检验的则是重定向到一个全新的架构类别,其技术栈年轻但完全开放。再加上评测机本身已有的算力,即 AVX2 SIMD、Xe iGPU 和 NPU,以及 NVIDIA 基线,它们让项目朝着六种计算类型的目标迈出了实质性的一步。三块 PCIe 卡无法同时装进这台机器。至于还要同时加入哪些卡才能达到六种计算类型的指标,我们会在流程后期再决定。
下一代硬件:从数月的调优到数天的测量
挑选第一批硬件只是开局的第一步。接下来还有若干问题会一路伴随整个项目。一个关键问题关乎抽象层级。对于足够不透明的设备,我们在 ISA 层面的编程效率可能永远比不上厂商自己的高层工具。可依赖这些专有工具,又与快速支持新硬件的目标相冲突。找准该瞄准的层级、提升对黑盒微架构建模的能力,将决定未来哪些设备值得纳入。前两个目标已经选定,工作重心也从选择转向执行:把 Radeon AI PRO R9700 和 Tenstorrent Blackhole p150a 集成进评测机器,摸清它们的特性,并搭起测量流水线,用来比较各家实现者的编译器能在它们身上做到什么。此外,我们还要定义能够受益于六种计算类型的工作负载,针对这六种计算类型手工实现这些负载,然后让各家实现者去挑战自动完成分解。MOCHA 成功之处带来的回报,是一个采用下一款新型加速器只需数天测量、而非数月专家手工调优的世界。把硬件选对,是弄清这一切的第一步。
其他资源
下载 DARPA MOCHA(面向异构架构的机器学习与优化引导编译器)- https://www.darpa.mil/research/programs/mocha-machine-learning