Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using $\mathbb{F}_2$

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Zhou, Keren, Lezcano, Mario, Goucher, Adam, Rakhmati, Akhmed, Niu, Jeff, Lebar, Justin, Szczerbuk, Pawel, Bell, Peter, Tillet, Phil, Raoux, Thomas, Moudallal, Zahi
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866911491077177344
author Zhou, Keren
Lezcano, Mario
Goucher, Adam
Rakhmati, Akhmed
Niu, Jeff
Lebar, Justin
Szczerbuk, Pawel
Bell, Peter
Tillet, Phil
Raoux, Thomas
Moudallal, Zahi
author_facet Zhou, Keren
Lezcano, Mario
Goucher, Adam
Rakhmati, Akhmed
Niu, Jeff
Lebar, Justin
Szczerbuk, Pawel
Bell, Peter
Tillet, Phil
Raoux, Thomas
Moudallal, Zahi
contents 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 $\mathbb{F}_2$. 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_23819
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using $\mathbb{F}_2$
Zhou, Keren
Lezcano, Mario
Goucher, Adam
Rakhmati, Akhmed
Niu, Jeff
Lebar, Justin
Szczerbuk, Pawel
Bell, Peter
Tillet, Phil
Raoux, Thomas
Moudallal, Zahi
Programming Languages
Hardware Architecture
Distributed, Parallel, and Cluster Computing
Performance
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 $\mathbb{F}_2$. 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.
title Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using $\mathbb{F}_2$
topic Programming Languages
Hardware Architecture
Distributed, Parallel, and Cluster Computing
Performance
url https://arxiv.org/abs/2505.23819