PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks

本文介绍了 PolyBlocks，这是一个基于 MLIR 的模块化编译器基础设施，通过组合轻量级仿射分析与启发式成本模型自动执行多级分块、融合及算子映射等优化，实现了从高层框架到特定 AI 芯片的高效代码生成，并在 NVIDIA GPU 上的实验表明其性能可媲美甚至超越 Torch Inductor 和 XLA。

要点

这篇论文介绍了一个名为 PolyBlocks 的新系统。为了让你轻松理解，我们可以把AI 芯片（如 NVIDIA 显卡）想象成一座超级繁忙的现代化厨房，把AI 模型（如 PyTorch 写的代码）想象成复杂的食谱。

1. 核心问题：为什么需要 PolyBlocks？

现状：大厨与传菜员的困境
目前，当我们用 PyTorch 或 JAX 写 AI 程序时，就像是一个只会写“食谱”的大厨（高级框架）。

• 传统做法（Eager 模式）： 大厨每写一步（比如“切菜”），就喊一声，传菜员（编译器）立刻去厨房找现成的工具（厂商库，如 CuDNN）把菜切好，端上来，再喊下一句。这导致厨房进进出出非常频繁，效率低，而且很多步骤是重复的。
• 现有编译器（如 Torch Inductor, XLA）： 它们试图把整张食谱看一遍，然后优化。但它们有个大毛病：它们太依赖“现成的工具包”了。如果食谱里有个特殊的菜，工具包里没有，它们就束手无策，或者只能生硬地拼凑。而且，如果厨房换了新设备（新芯片），这些编译器往往需要大改，因为它们和特定的“工具包”绑定太死。

PolyBlocks 的解决方案：全自动的“超级管家”
PolyBlocks 就像是一个拥有上帝视角的全自动超级管家。

• 它不依赖任何现成的“工具包”（厂商库）。
• 它拿到食谱后，会从头到尾重新设计整个烹饪流程。
• 它能自动决定：哪些步骤可以合并？哪些食材可以提前切好放在手边（显存优化）？怎么利用厨房里的特殊机器（矩阵计算单元）？
• 最重要的是，如果厨房换成了新的设备（比如新的 AI 芯片），只要告诉管家新设备的特性，它就能利用同一套逻辑，重新生成一套完美的烹饪流程，而不需要重新发明轮子。

2. PolyBlocks 是怎么工作的？（核心魔法）

PolyBlocks 的核心在于它把复杂的优化过程拆解成了几个像乐高积木一样的步骤（Pass Pipelines）：

A. 切片与融合（Slicing-based Fusion）：把“分步做”变成“一气呵成”

• 比喻： 假设食谱说“先切洋葱，再切土豆，然后把它们炒在一起”。
  • 普通做法： 切好洋葱放盘子里，切好土豆放另一个盘子，最后倒进锅里炒。中间多洗了两次盘子，还多走了几步路。
  • PolyBlocks 的做法： 它发现“切洋葱”和“切土豆”其实可以同时进行，或者切完直接下锅。它把这两个步骤融合成一个动作，甚至把中间那个“盘子”（临时内存）都省掉了。
  • 创新点： 它的“融合”非常聪明，甚至允许为了省时间，稍微多切一点点洋葱（冗余计算），只要整体速度更快就行。

B. 分块与搬运（Tiling & Packing）：把“大仓库”变成“手边小抽屉”

• 比喻： 厨房的冰箱（显存）很大，但拿东西慢；灶台旁边的抽屉（片上内存/寄存器）很小，但伸手就能拿到。
• PolyBlocks 的做法： 它不会一次性把整袋面粉（大数据）从冰箱搬到灶台。它会计算好，每次只搬刚好够炒一盘菜的量（分块/Tiling）到抽屉里。
• 动态打包： 对于卷积（一种特殊的图像处理），它甚至能像变魔术一样，在搬运的过程中就把食材“切好、摆好”（On-the-fly packing），让机器一拿起来就能直接加工，不需要额外的整理时间。

C. 自动适配“超级机器”（Mapping to Matrix Units）

• 比喻： 现代厨房里有专门的“自动切肉机”（GPU 的 Tensor Core），一次能切 16x16 的肉块。
• PolyBlocks 的做法： 它会自动把食谱里的普通切菜指令，转换成“自动切肉机”能听懂的语言。它不需要人工去写复杂的机器代码，而是自动把数据排列成机器最喜欢的形状。

D. 注意力机制的融合（Attention Layer Fusion）

• 比喻： 在 Transformer 模型（如大语言模型）中，有一个叫“注意力”的步骤，需要计算大量的数据关系，通常非常慢。
• PolyBlocks 的做法： 它把原本需要分好几步、来回搬运数据的复杂过程，压缩成一个超级流畅的流水线。它甚至能在计算过程中直接修正数据，不需要把中间结果写回硬盘再读出来。这让大模型跑得飞快。

3. 效果如何？（实战表现）

论文在 NVIDIA 的 A100 和 A10 显卡上做了测试：

• 对比对象： 目前业界最强的两个编译器（Torch Inductor 和 XLA），它们通常依赖厂商提供的“现成工具库”（如 CuDNN）。
• 结果：
  • 速度： 在大多数情况下，PolyBlocks 生成的代码比 Inductor 和 XLA 更快，或者至少一样快。
  • 独立性： 最惊人的是，PolyBlocks 完全不需要依赖那些昂贵的厂商库。它自己生成的代码，在单个任务（如矩阵乘法）上，竟然能和那些由顶级工程师手写的、经过几十年优化的库（如 CuBLAS）打得有来有回，甚至有时更快！
  • 通用性： 它不仅能跑 PyTorch，也能跑 JAX。而且，如果未来出了新的 AI 芯片，PolyBlocks 只需要花很少的精力就能适配，不需要像其他编译器那样推倒重来。

4. 总结：为什么这很重要？

想象一下，以前我们要去一个新的城市（新芯片）旅行，必须找当地导游（厂商库）带路，如果导游不在，我们就寸步难行。

PolyBlocks 就像是一个自带 GPS 和全能翻译的旅行机器人。

• 不管你去哪个城市（新芯片），它都能自己规划路线。
• 不管你说什么语言（PyTorch, JAX），它都能听懂并执行。
• 它不需要依赖当地的导游，自己就能把路走得比导游还快。

这篇论文证明了：完全由编译器自动生成的代码，不仅可以和人工优化的库相媲美，而且在面对未来不断变化的 AI 硬件时，具有更强的生命力和适应性。 这是通往未来 AI 硬件自动适配的关键一步。

技术摘要

PolyBlocks：面向 AI 芯片与编程框架的编译器基础设施

1. 研究背景与问题 (Problem)

随着人工智能（AI）领域的快速发展，出现了大量专用加速器芯片（如 NVIDIA GPU、TPU 等）以及高层编程框架（如 PyTorch, JAX, TensorFlow）。然而，现有的编译器基础设施在连接高层抽象与底层硬件性能时面临以下挑战：

• 性能瓶颈与库依赖：当前主流编译器（如 Torch Inductor, XLA）为了追求性能，严重依赖厂商手动优化的库（如 CuDNN, CuBLAS, FlashAttention）。这导致编译器生成的代码难以超越这些库，且缺乏跨算子的优化能力（如算子融合）。
• 可移植性与扩展性差：现有编译器通常针对特定硬件（CPU/GPU）构建，使用了多种中间表示（IR）和不同的后端。当需要支持新型 AI 芯片时，往往需要重新开发大量优化逻辑，难以复用。
• 手动调优的局限性：虽然像 Triton 这样的中层框架提供了更好的抽象，但仍需要程序员手动处理分块（tiling）和融合（fusion），难以自动处理复杂的计算图。
• 融合能力不足：现有的融合技术（如标准循环融合）难以处理复杂的依赖关系，特别是针对 Attention 层等包含大规模中间结果（如 $Q \times K^T$）的计算，难以实现全图融合以消除全局内存访问。

2. 方法论 (Methodology)

PolyBlocks 是一个基于 MLIR (Multi-Level Intermediate Representation) 的模块化、可复用的编译器基础设施。其核心设计理念是完全代码生成（Fully Code-Generating），即不依赖任何预编译的厂商库，而是从高层框架直接生成针对特定目标硬件的低级指令和 intrinsics。

2.1 核心架构

PolyBlocks 采用五阶段的 Pass Pipeline 架构，将编译过程解耦为：

1. 前端 (S1-S2)：将来自 PyTorch (Torch-MLIR), JAX (StableHLO), 或 TensorFlow (MHLO) 的 IR 转换为 affine（仿射）和 memref（内存引用）dialect。此阶段主要引入缓冲区语义，几乎与目标硬件无关。
2. 中端优化器 (S3)：这是核心优化阶段。基于 affine 和 memref dialect，执行多级别分块、算子融合、数据移动代码生成、向量化以及映射到矩阵单元（Matrix Units/Tensor Cores）。此阶段利用轻量级的仿射分析（Affine Analysis）而非昂贵的多面体（Polyhedral）整数集运算，以平衡编译速度与优化能力。
3. 后端 (S4-S5)：将优化后的 IR 转换为特定目标（如 NVIDIA GPU 的 NVVM, AMD GPU 的 ROCm, 或 CPU 的 OpenMP）的并行抽象，最终生成 LLVM IR 和机器码。
4. 编译器驱动 (C4)：用 Python 编写，负责编排 Pass 流程、处理 JIT 编译时的输入输出数据转换（Tensor 到 Memref ABI）以及内核启动。

2.2 关键优化技术

• 基于切片的仿射融合 (Slicing-based Affine Fusion)：
  • 超越了传统的循环融合。通过计算生产者（Producer）需要提供给消费者（Consumer）的数据切片，将生产者“拉入”消费者循环中。
  • 允许引入少量的冗余计算（通过成本模型控制），以消除中间缓冲区（Intermediate Buffers）和全局内存访问。
  • 支持将非完美嵌套（Imperfect Nests）的算子（如深度卷积）进行融合。
• 两阶段分块与融合策略：
  • 首先对关键算子（如 MatMul, Conv）进行分块，生成快速缓冲区（Fast Buffers，如 Shared Memory）。
  • 然后将其他算子融合进这些分块后的循环中，确保数据在片上内存中的最大化复用。
• Attention 层自动融合：
  • 针对 Transformer 中的 Attention 层，设计了 reduce-reduce-fusion 和 wmma-fusion 两个 Pass。
  • 能够自动融合 $Q \times K^T$ 的矩阵乘法、Softmax 和 $Softmax \times V$ 的乘法，甚至处理其中的修正计算（Correction），从而完全消除 $Q \times K^T$ 结果写入 DRAM 的需求，实现类似 FlashAttention 的性能。
• 卷积的即时打包 (On-the-fly Packing)：
  • 将卷积转换为矩阵乘法时，不预先在内存中展开 Im2Col 操作（避免巨大的内存开销），而是在快速内存（Shared Memory）中按需加载并打包数据块，直接映射到矩阵单元。
• 矩阵单元映射：
  • 自动将仿射循环映射到硬件特定的矩阵单元（如 NVIDIA Tensor Cores 的 WMMA），无需手动编写复杂的底层代码。

3. 主要贡献 (Key Contributions)

1. 首个全代码生成的 MLIR 基础设施：PolyBlocks 证明了在不依赖 CuDNN/CuBLAS 等厂商库的情况下，仅通过编译器自动生成的代码，即可在 NVIDIA GPU 上达到甚至超越现有最佳库的性能。
2. 强大的跨算子融合能力：通过基于切片的融合技术，实现了跨不同算子（如卷积与点积、Attention 层内部）的深度融合，显著减少了全局内存访问。
3. 模块化与可移植性：其架构设计使得为新型 AI 芯片构建编译器变得容易。只需定制 S3 阶段的中端优化和 S4/S5 的后端代码，即可复用大部分基础设施。
4. 自动化的 Attention 优化：无需手动编写 FlashAttention 内核，PolyBlocks 能自动将 Attention 层优化为寄存器级复用的高效代码。
5. 广泛的框架支持：统一支持 PyTorch, JAX, TensorFlow，并能无缝处理混合精度（FP16/FP32）计算。

4. 实验结果 (Results)

实验在 NVIDIA A100 和 A10 GPU 上进行，对比了 Torch Inductor、XLA、TensorRT 以及厂商库（CuDNN/CuBLAS）。

• 端到端性能：
  • PyTorch (Batch Size 1)：PolyBlocks 平均比 Eager 模式快 2.15 倍，比 Inductor 快 1.4 倍，比 TensorRT 快 2.4 倍。
  • PyTorch (Batch Size 8)：PolyBlocks 平均比 Eager 快 1.8 倍，与 Inductor 性能相当（0.97x），但在小 Batch 场景下优势更明显。
  • JAX：PolyBlocks 比 JAX Eager 快 2.12 倍，比 XLA 快 1.15 倍。
• 单算子性能：
  • 卷积 (Convolutions)：在数百个卷积算子测试中，PolyBlocks 生成的代码在大量情况下与 CuDNN 性能相当，甚至在近 50 个案例中超过 CuDNN 2 倍以上。
  • 矩阵乘法 (MatMul)：性能与 CuBLAS 和手写 Triton 内核非常接近。
• 消融实验：
  • Tensor Cores：开启 Tensor Cores 带来约 17 倍 的加速。
  • 跨算子融合：相比未融合但其他优化已开启的代码，融合带来 2.87 倍 的额外提升，证明了减少全局内存访问的重要性。
  • Attention 优化：在 Attention 层配置中，结合 reduce-reduce 和 wmma 融合，相比 Inductor 实现了显著的加速（部分场景下高达 0.69x 的时间，即速度提升约 1.4 倍，具体取决于配置）。

5. 意义与影响 (Significance)

• 打破厂商库垄断：PolyBlocks 展示了编译器自动代码生成的巨大潜力，表明未来的 AI 硬件优化可以不再完全依赖厂商手动优化的黑盒库，从而加速新硬件的生态建设。
• 加速新硬件落地：其模块化架构使得为新型 AI 芯片（如 NPU, 专用 ASIC）开发高性能编译器变得可行且高效，无需从零开始构建复杂的优化逻辑。
• 提升开发效率：通过“一行指令”（One-line directive）即可将 PyTorch/JAX 模型编译为高性能代码，降低了开发者优化硬件性能的门槛。
• 推动编译器技术演进：PolyBlocks 提出的基于轻量级仿射分析和切片融合的优化策略，为处理大规模、复杂依赖的计算图提供了新的范式，解决了传统多面体编译在扩展性上的瓶颈。

综上所述，PolyBlocks 不仅是一个高性能的编译器后端，更是一个旨在解决 AI 芯片与软件生态之间“最后一公里”问题的通用基础设施，为未来 AI 硬件的灵活部署和性能挖掘提供了强有力的技术支撑。