Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers

TL;DR

Axe Layout introduces a hardware-aware abstraction for mapping tensor coordinates to physical device space, unifying various layout strategies to optimize deep learning workloads across heterogeneous hardware.

Contribution

We propose Axe Layout, a unified, hardware-aware layout abstraction that simplifies and optimizes data placement and collective operations in machine learning compilers.

Findings

Achieves performance close to hand-tuned kernels on GPUs and multi-device setups.

Unifies tiling, sharding, and replication across device hierarchies.

Enables a single kernel design with collective primitives for diverse hardware.

Abstract

Scaling modern deep learning workloads demands coordinated placement of data and compute across device meshes, memory hierarchies, and heterogeneous accelerators. We present Axe Layout, a hardware-aware abstraction that maps logical tensor coordinates to a multi-axis physical space via named axes. Axe unifies tiling, sharding, replication, and offsets across inter-device distribution and on-device layouts, enabling collective primitives to be expressed consistently from device meshes to threads. Building on Axe, we design a multi-granularity, distribution-aware DSL and compiler that composes thread-local control with collective operators in a single kernel. Experiments show that our unified approach can bring performance close to hand-tuned kernels on across latest GPU devices and multi-device environments and accelerator backends.