用少得多的代码实现高性能计算

许多公司在招聘人才方面投入巨大，以创建支撑现代人工智能系统的高性能库代码。例如，英伟达开发了一些最先进的高性能计算（HPC）库，打造了一条难以被他人突破的竞争护城河。

但如果有几个学生，在几个月内，仅用几百行代码，而非几万或几十万行，就能与最先进的高性能计算库相媲美，那又会怎样呢？

这就是麻省理工学院计算机科学与人工智能实验室（CSAIL）的研究人员用一种名为Exo 2的新编程语言所展示的成果。

Exo 2属于一类新的编程语言，麻省理工学院教授乔纳森·拉根 - 凯利称之为“用户可调度语言”（USLs）。USLs并不指望不透明的编译器自动生成尽可能快的代码，而是让程序员掌握主动权，允许他们编写“调度”，明确控制编译器如何生成代码。这使得性能工程师能够将指定计算内容的简单程序，转换为与原始规范执行相同任务但速度快得多的复杂程序。

现有USL（如最初的Exo）的局限性之一是其调度操作相对固定，这使得在不同的“内核”（高性能库中的各个组件）间复用调度代码变得困难。

相比之下，Exo 2使用户能够在编译器外部定义新的调度操作，有助于创建可复用的调度库。论文的第一作者Yuka Ikarashi是麻省理工学院电气工程与计算机科学专业的博士生，也是计算机科学与人工智能实验室（CSAIL）的成员，她表示Exo 2可以将调度代码总量减少100倍，并在多个不同平台上实现与最先进的实现方案相媲美的性能，其中包括为许多机器学习应用提供支持的基本线性代数子程序（BLAS）。对于专注于针对不同操作、数据类型和目标架构优化内核的高性能计算（HPC）工程师来说，这使其成为一个颇具吸引力的选择。

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

Exo 2的一个主要优势在于，它通过在不同应用程序和硬件目标间复用调度代码，减少了任何时刻所需的编码工作量。研究人员在Exo 2中实现了一个约2000行代码的调度库，封装了针对线性代数和特定目标（AVX512、AVX2、Neon和Gemmini硬件加速器）的可复用优化。该库整合了80多个高性能内核的调度工作，每个内核最多只需十几行代码，其性能与MKL、OpenBLAS、BLIS和Halide相当，甚至更优。

《Exo 2》包含一种名为“游标”的新颖机制，它提供了一种在整个调度过程中指向目标代码的 “稳定引用”。Ikarashi 表示，稳定引用对于用户在库函数中封装调度至关重要，因为它使调度代码独立于目标代码转换。

“我们认为，通用规范语言（USL）应设计为可由用户扩展，而非具有一组固定的操作，” 井良说。“这样，一种语言就能通过实现适应不同优化需求和应用领域的库，发展到足以支持大型项目。”

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

井良（Ikarashi）和拉根 - 凯利（Ragan-Kelley）与研究生凯文·钱（Kevin Qian）、萨米尔·德鲁比（Samir Droubi）、Adobe公司的亚历克斯·莱金克（Alex Reinking）以及前计算机科学与人工智能实验室（CSAIL）博士后吉尔伯特·伯恩斯坦（Gilbert Bernstein，现任华盛顿大学教授）共同撰写了这篇论文。这项研究部分由美国国防高级研究计划局（DARPA）和美国国家科学基金会资助，而第一作者还得到了正村（Masason）、船井（Funai）和夸德（Quad）奖学金的支持。