摘要

现代人工智能工作负载在训练与推理过程中高度依赖于经过优化的计算内核。这些AI内核遵循明确的数据流模式，例如在DRAM与SRAM之间搬运数据块（tiles），并在这些数据块上执行一系列计算操作。然而，尽管这些数据流模式清晰，编写高性能内核仍然极为复杂。要达到峰值性能，必须进行细致入微、面向硬件的优化，以充分挖掘现代加速器的潜力。尽管领域专用编译器试图减轻编写高性能内核的负担，但它们在易用性与表达能力方面仍存在显著差距。本文提出TileLang，一种面向AI内核编程的通用分块编程模型，旨在提升AI内核开发的效率。TileLang将调度空间（包括线程绑定、内存布局、张量化与流水线化）与数据流逻辑解耦，并将其封装为一组可定制的注解与原语。这一设计使开发者能够专注于内核本身的数据流逻辑，而将绝大多数优化任务交由编译器自动完成。我们在多种常用设备上开展了全面的实验，覆盖大量不同场景。评估结果表明，TileLang在关键内核上达到了当前最先进的性能水平。这充分证明，其统一的块（block）与线程（thread）编程范式，以及透明的调度能力，能够同时满足现代AI系统开发对性能与灵活性的双重需求。