LEGO: A Layout Expression Language for Code Generation of Hierarchical Mapping

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Tavakkoli, Amir Mohammad, Oancea, Cosmin, Hall, Mary
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866909961457500160
author Tavakkoli, Amir Mohammad
Oancea, Cosmin
Hall, Mary
author_facet Tavakkoli, Amir Mohammad
Oancea, Cosmin
Hall, Mary
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_08091
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle LEGO: A Layout Expression Language for Code Generation of Hierarchical Mapping
Tavakkoli, Amir Mohammad
Oancea, Cosmin
Hall, Mary
Programming Languages
Distributed, Parallel, and Cluster Computing
Performance
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.
title LEGO: A Layout Expression Language for Code Generation of Hierarchical Mapping
topic Programming Languages
Distributed, Parallel, and Cluster Computing
Performance
url https://arxiv.org/abs/2505.08091