一句话:用 F_2 上的线性代数统一建模张量布局,实现通用布局转换并集成到 Triton,提升代码生成效率与鲁棒性。
高效张量计算是现代深度学习负载的基石,但现有方案难以灵活且高性能地设计与实现张量布局,即逻辑张量与硬件资源之间的映射。随着深度学习算法与硬件复杂度上升,布局处理亟需通用而系统化的方法。本文提出 Linear Layouts,用 F_2 上的线性代数建模张量布局:将布局表示为作用于硬件表示比特的二元矩阵,从而支持通用布局定义而非传统的逐一特判,并实现通用布局间转换,消除现有方案中的二次爆炸问题。作者将其集成到 Triton 编译器中,在优化单个 Triton 算子及完整 Kernel 上验证了有效性,同时表明该方案降低了编译器后端工程负担,并修复了 Triton 遗留布局系统中的若干缺陷。
Efficient tensor computation is a cornerstone of modern deep learning (DL) workloads, yet existing approaches struggle to achieve flexible and performant design and implementation of tensor layouts -- mappings between logical tensors and hardware resources. The increasing complexity of DL algorithms and hardware demands a generic and systematic approach to handling tensor layouts. In this work, we introduce Linear Layouts, a novel approach that models tensor layouts using linear algebra over F2. By representing tensor layouts as binary matrices acting on the bits of the hardware representation, our approach enables a generic layout definition—as opposed to the classical case-by-case approach -- and allows for generic layout-to-layout conversions, eliminating the quadratic explosion that plagues existing solutions. We integrate linear layouts with Triton and demonstrate their effectiveness in optimizing individual Triton operators as well as kernels written in Triton. We also show that linear layouts reduce engineering effort in the compiler backend while fixing several bugs in Triton's legacy layout system.
DOI 原文 · arXiv · PDF(开放获取) · DBLP
@inproceedings{0001CGRNLS0TRM26,
title = {Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using F_2},
author = {Keren Zhou and Mario Lezcano-Casado and Adam P. Goucher and Akhmed Rakhmati and Jeff Niu and Justin Lebar and Pawel Szczerbuk and Peter Bell and Phil Tillet and Thomas Raoux and Zahi Moudallal},
booktitle = {ASPLOS},
year = {2025},
doi = {10.1145/3760250.3762221}
}