Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers
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.
Code (0)
등록된 구현이 없습니다.
Similar Papers 제목 키워드 기반
Hexcute: A Tile-based Programming Language with Automatic Layout and Task-Mapping Synthesis
Deep learning (DL) workloads mainly run on accelerators like GPUs. Recent DL quantization techniques demand a new matrix multiplication operator with mixed input data types, further complicating GPU optimization. Prior h…
GPUQuantizationCMLCompiler: A Unified Compiler for Classical Machine Learning
Classical machine learning (CML) occupies nearly half of machine learning pipelines in production applications. Unfortunately, it fails to utilize the state-of-the-practice devices fully and performs poorly. Without a un…
CPUGPUALT: Boosting Deep Learning Performance by Breaking the Wall between Graph and Operator Level Optimizations
Deep learning models rely on highly optimized tensor libraries for efficient inference on heterogeneous hardware. Current deep compilers typically predetermine layouts of tensors and then optimize loops of operators. How…
A Tensor Compiler for Unified Machine Learning Prediction Serving
Machine Learning (ML) adoption in the enterprise requires simpler and more efficient software infrastructure---the bespoke solutions typical in large web companies are simply untenable. Model scoring, the process of obta…
BIG-bench Machine LearningCPUGPUUniSparse: An Intermediate Language for General Sparse Format Customization
The ongoing trend of hardware specialization has led to a growing use of custom data formats when processing sparse workloads, which are typically memory-bound. These formats facilitate optimized software/hardware implem…
AttributeCode GenerationCPUGPU