LEGO: A Layout Expression Language for Code Generation of Hierarchical Mapping
Amir Mohammad Tavakkoli, Cosmin Oancea, Mary Hall

TL;DR
LEGO is a novel layout expression language that optimizes data movement and hierarchical code generation for GPUs, integrating with existing compiler frameworks to improve performance and flexibility.
Contribution
LEGO introduces a layout-independent code expression approach that enhances data and computation layout optimization in GPU programming.
Findings
LEGO achieves performance comparable to Triton.
LEGO integrates effectively with MLIR and CUDA.
LEGO broadens data and thread layout optimization capabilities.
Abstract
We describe LEGO, a new approach to optimizing data movement whereby code is expressed as a layout-independent computation and composed with layouts for data and computation. This code generator organization derives complex indexing expressions associated with hierarchical parallel code and data movement for GPUs. LEGO maps from layout specification to indexing expressions, and can be integrated into existing compilers and code templates. It facilitates the exploration of data layouts in combination with other optimizations. We demonstrate LEGO's integration with the Triton and MLIR compilers, and with CUDA templates. We show that LEGO is capable of deriving performance competitive with Triton, and shows broad applicability for data and thread layout mapping optimizations in its integration with CUDA and MLIR.
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsParallel Computing and Optimization Techniques · Computer Graphics and Visualization Techniques · Logic, programming, and type systems
