自动优化张量程序,让代码跑得更快更省内存。
Axon: A Synthesizing Superoptimizer for Tensor Programs

- 用程序合成自动生成高效指令,无需手动调优。
- 通过代数变换和硬件约束,找到性能最优的计算方案。
- 适合需要极致性能的AI加速器开发人员使用。
编写高性能AI加速器内核需要深入掌握分块、指令选择、数据布局和算子融合等技术,对程序员要求极高。本文聚焦基于分块的AI加速器程序,提出Axon——一种用于张量程序的合成式超优化器:它利用程序合成从语义规范自动生成目标指令,并探索语义等价的程序变体以经验性地选择最佳性能内核。Axon通过在计算图中传播算子发现代数变换,并使用带无界张量的SMT验证所有变换保持语义不变,无需手工编写重写规则。随后将张量操作降低为目标ISA指令,结合硬件描述约束探索分块配置,并融合算子与指令以最小化内存流量。
原文摘要 · Abstract (English)
Writing high performance kernels for AI accelerators requires deep expertise in tiling, instruction selection, data layout, and operator fusion placing a significant burden on programmers. In this paper, we focus on tile based AI accelerator programs and present Axon, a synthesizing superoptimizer for tensor programs: it uses program synthesis to automatically generate target instructions from semantics specifications, and explores semantically equivalent program variants to select the best performing kernel empirically. Axon discovers algebraic transformations by propagating operators through computation graphs and uses SMT over unbounded tensors to guarantee that all transformations preserve semantics without requiring hand crafted rewrite rules. It then lowers tensor operations to target ISA instructions, explores tiling configurations constrained by hardware descriptions, and fuses operators and instructions to minimize memory traffic.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。