高性能计算,代码更少

由 emer 发布

许多公司投入巨资聘请人才来创建支撑现代人工智能系统的高性能库代码。例如,NVIDIA 开发了一些最先进的高性能计算 (HPC) 库,建立了一条竞争壁垒,事实证明其他公司很难突破。

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

麻省理工学院计算机科学与人工智能实验室 (CSAIL) 的研究人员通过一种名为 Exo 2的新编程语言展示了这一点。

Exo 2 属于麻省理工学院教授 Jonathan Ragan-Kelley 称之为“用户可调度语言”(USL)的新型编程语言。USL 让程序员掌控全局,让他们可以编写“调度”来明确控制编译器生成代码的方式,而不是指望一个不透明的编译器自动生成最快的代码。这使得性能工程师能够将指定要计算内容的简单程序转换为复杂程序,这些复杂程序可以执行与原始规范相同的操作,但速度要快得多。

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

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

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

Exo 2 的一个主要优势是,它通过在应用程序和硬件目标之间重用调度代码,减少了每次所需的编码工作量。研究人员在 Exo 2 中实现了一个包含大约 2,000 行代码的调度库,封装了线性代数专用和目标专用(AVX512、AVX2、Neon 和 Gemmini 硬件加速器)的可重用优化。该库整合了 80 多个高性能内核的调度工作,每个内核最多包含十几行代码,性能可与 MKL、OpenBLAS、BLIS 和 Halide 相媲美甚至更好。

Exo 2 包含一种名为“游标”的新型机制,它提供了所谓的“稳定引用”,用于在整个调度过程中指向目标代码。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 的支持。

参考: https://news.mit.edu/2025/high-performance-computing-with-much-less-code-0313


0条评论

发表评论


验证码