[论文] 瓦片设计：在基于瓦片的多PE加速器上自动化GEMM部署

Source: arXiv - 2512.13638v1

请提供您希望翻译的具体文本（例如摘要、章节内容或其他段落），我将按照要求保留源链接并将其翻译成简体中文。谢谢！

概述

本文介绍了 Design in Tiles (DiT)，一个自动化框架，用于弥合高级 GEMM（通用矩阵乘法）代码与驱动当今 AI 芯片的低层、基于瓦片的多处理元件（PE）加速器之间的差距。通过将部署工具链与可配置的执行模型相结合，DiT 能够生成接近最优的 GEMM 内核，而无需当前硬件特定库所需的繁琐手动调优。

关键贡献

• 端到端自动化：一个工具链，接受 GEMM 描述并为任意基于 tile 的加速器配置生成完全调优的实现。
• 可配置可执行模型：一个参数化性能模型，捕获 PE 数量、tile 尺寸、存储层次结构和 tile 间带宽，实现快速设计空间探索。
• 可扩展映射策略：一种系统化的方法，将矩阵划分到二维 tile 网格（例如 32 × 32 tiles）上，以最大化 PE 利用率并隐藏通信延迟。
• 性能持平及超越：在模拟的 1979 TFLOPS FP8 加速器（带宽 4 TB/s）上，展示相较于 NVIDIA GH200 专家调优 GEMM 库的 1.2–2.0× 加速。
• 开源原型：作者发布 DiT 框架及一套基准脚本，鼓励社区采纳和进一步研究。

方法论

1. 硬件抽象：作者将基于瓦片的加速器建模为一个由相同处理单元（PE）组成的网格，每个 PE 具有本地寄存器、一个小型 Scratchpad（局部存储）以及高带宽互连。瓦片大小、PE 数量和内存带宽等参数被公开为可调节的 knob（调节钮）。
2. 瓦片级调度：GEMM 的三个嵌套循环 (i, j, k) 被重新排序并划分为瓦片，以匹配硬件网格。DiT 自动决定每个瓦片应计算输出矩阵的多少行/列，以及归约维度 (k) 如何在时间步上拆分。
3. 数据移动规划：利用可执行模型，DiT 预测从共享内存层次加载瓦片的开销，安排双缓冲，并插入预取操作以保持 PE 的忙碌状态。
4. 代码生成：高级调度被降级为 C++/CUDA 类的内核，面向特定的 PE ISA（例如基于 RISC‑V 或自定义 SIMD）。生成的代码使用厂商的后端进行编译，生成可在目标加速器上运行的二进制文件。
5. 迭代优化：轻量级自动调优器在模型性能估计的指导下，探索一小组瓦片大小候选方案，以挑选出能够实现最高预测 PE 利用率的配置。

结果与发现

• PE 利用率: DiT 始终保持 > 95 % 的 PE 处于活跃状态，而在不规则形状上手工调优的 GH200 内核约为 ~70 % 。
• 内存带宽: 基于模型的预取调度实现了 > 90 % 的理论 4 TB/s 带宽，消除了因朴素数据移动导致的停顿。
• 编译时间: 端到端生成（包括自动调优）在标准工作站上耗时不到 5 分钟，远快于定制 ASIC 库常见的数周手动调优周期。

实际意义

• Accelerator Vendors：DiT 提供了一条在新 AI 芯片上交付“即插即用” GEMM 库的路径，降低了对昂贵且仅限专家的软件团队的需求。
• Framework Integrators：深度学习堆栈（例如 PyTorch、TensorFlow）可以将 DiT 作为自定义硬件的后端调用，在无需重写 kernel 的情况下获得性能提升。
• Edge & Cloud Deployments：能够自动适配不同矩阵形状的 GEMM 意味着对经常涉及非方形张量（例如 transformer 注意力头）的推理工作负载实现更好的利用率。
• Hardware‑Software Co‑Design：设计者可以在模型中实验瓦片数量、互连带宽和内存容量，立即看到对 GEMM 性能的影响，从而在设计周期的早期做出更好的硅片决策。

局限性与未来工作

• 模型保真度：虽然可执行模型捕获了主要的延迟来源，但它抽象掉了细粒度的影响，例如缓存行冲突和电压‑频率缩放，这些可能影响真实硅片的结果。
• GEMM 之外的范围：DiT 目前专注于密集矩阵乘法；将框架扩展到卷积、稀疏核或混合精度工作负载仍是一个未解决的挑战。
• 硬件多样性：原型针对特定类别的基于 tile 的多 PE 加速器；将其适配到异构设计（例如将 tile 与向量单元结合）将需要额外的建模工作。
• 自动调优开销：虽然轻量，但自动调优器仍可能在极大设计空间中错过全局最优的 tile 大小；集成更复杂的搜索算法是计划中的方向。

Design in Tiles 表明，借助合适的抽象和工具，传统上繁琐的 GEMM 映射到前沿 AI 加速器的任务可以实现自动化，既能提升性能，又能加快下一代硬件的上市时间。

作者

• Aofeng Shen
• Chi Zhang
• Yakup Budanaz
• Alexandru Calotoiu
• Torsten Hoefler
• Luca Benini

论文信息

• arXiv ID: 2512.13638v1
• 分类: cs.DC, cs.AR
• 发布日期: 2025年12月15日
• PDF: Download PDF