MIT News | Massachusetts Institute of Technology

更少代码实现高性能计算:探索 Exo 2 编程语言

Exo 2 编程语言实现了编译器外部可重用的调度库。

Adam Conner-Simons | MIT CSAIL

发布日期:2025年3月13日

计算机芯片的特写照片,主要由绿色和黄色组件构成

说明:

一种名为 “Exo 2” 的新型编程语言能够实现高性能编码,只需几百行代码就能与最先进的库竞争,而无需数万甚至数十万行代码。

图片来源:

Photo: David Carron/Wikimedia Commons

许多公司投入大量资金聘请人才来创建支撑现代人工智能系统的高性能库代码。例如,NVIDIA 开发了一些最先进的 HPC (High-performance computing) 库,构建了一个竞争壁垒,其他公司很难突破。

但是,如果几个学生在几个月内,就能用几百行代码与最先进的 HPC 库竞争,而不是数万甚至数十万行代码呢?

这正是 MIT 计算机科学与人工智能实验室 (CSAIL) 的研究人员通过一种名为 Exo 2 的新编程语言所展示的。

Exo 2 属于 MIT 教授 Jonathan Ragan-Kelley 称之为 “用户可调度语言 (USL)” 的一类新型编程语言。USL 不是寄希望于不透明的编译器自动生成尽可能快的代码,而是让程序员掌握主动权,允许他们编写 “调度”,从而明确控制编译器生成代码的方式。这使得性能工程师能够将指定计算内容的简单程序,转换为与原始规范执行相同操作,但速度快得多的复杂程序。

现有 USL (例如最初的 Exo) 的一个局限性在于它们相对固定的调度操作集,这使得在不同的 “内核” (高性能库中的各个组件) 之间重用调度代码变得困难。

相比之下,Exo 2 允许用户在编译器外部定义新的调度操作,从而促进可重用的调度库的创建。第一作者 Yuka Ikarashi 是 MIT 电气工程与计算机科学专业的博士生,也是 CSAIL 的附属研究员。她表示,Exo 2 可以将总调度代码减少 100 倍,并在包括 Basic Linear Algebra Subprograms (BLAS) 在内的多个不同平台上提供与最先进实现相媲美的性能,而 BLAS 为许多机器学习应用提供支持。这使其成为 HPC 工程师的一个有吸引力的选择,他们专注于跨不同操作、数据类型和目标架构优化内核。

“这是一种自下而上的自动化方法,而不是在高性能代码上进行 ML/AI 搜索,” Ikarashi 说。“这意味着性能工程师和硬件实现者可以编写自己的调度库,这是一组应用于其硬件以达到峰值性能的优化技术。”

Exo 2 的一个主要优点是,它通过在应用程序和硬件目标之间重用调度代码,减少了任何时候所需的编码工作量。研究人员在 Exo 2 中实现了一个大约 2000 行代码的调度库,其中封装了线性代数特定的和目标特定的 (AVX512, AVX2, Neon, 和 Gemmini 硬件加速器) 可重用优化。该库整合了超过 80 个高性能内核的调度工作,每个内核最多十几行代码,所提供的性能与 MKL, OpenBLAS, BLIS, 和 Halide 相当或更好。

Exo 2 包含一种名为 “Cursors” 的新颖机制,它提供了他们在整个调度过程中所谓的指向对象代码的 “稳定引用”。Ikarashi 说,稳定引用对于用户将调度封装在库函数中至关重要,因为它使得调度代码独立于对象代码转换。

“我们认为,USL 应该被设计为用户可扩展的,而不是具有固定的操作集,” Ikarashi 说。“通过这种方式,一种语言可以通过实现库来支持大型项目,这些库可以适应不同的优化要求和应用领域。”

Exo 2 的设计允许性能工程师专注于高级优化策略,同时通过使用安全的原语来确保底层对象代码在功能上保持等效。未来,该团队希望扩展 Exo 2 对不同类型硬件加速器 (例如 GPU) 的支持。一些正在进行的项目旨在从正确性、编译时间和表达性方面改进编译器分析本身。

Ikarashi 和 Ragan-Kelley 与研究生 Kevin Qian 和 Samir Droubi、Adobe 的 Alex Reinking 以及前 CSAIL 博士后 Gilbert Bernstein (现为华盛顿大学教授) 共同撰写了论文。这项研究部分由美国国防高级研究计划局 (DARPA) 和美国国家科学基金会资助,而第一作者也得到了 Masason、Funai 和 Quad Fellowships 的支持。

Paper

Paper: “Exo 2: Growing a Scheduling Language.”

Related Links