PEAK: Generating High-Performance Schedules in MLIR
摘要
Machine learning frameworks rely on vendor libraries or autotuning frameworks for high-performance implementations of key operators like matrix multiplication and convolution. The transform dialect has recently been developed in the MLIR framework to facilitate the composition of transformations to implement optimized schedules for tensor computations. However, its interface to users is complex and requires an understanding of the underlying compiler transformations. In this paper, we describe a higher level scheduling language PEAK that is built on top of the MLIR transform dialect to ease the process of developing optimized schedules for deep learning operators on GPUs. PEAK expresses a simplified scheduling language in the MLIR/IREE compiler; it exploits domain-specific properties about data reuse in tensor contractions to determine thread mapping strategies and stage data through the GPU memory hierarchy. PEAK integrates an autotuner to explore implementations of high-performance code with schedules based on the MLIR Transform Dialect. PEAK enables a significant reduction of effort to construct high-performance GPU code using MLIR. A comparison with the state-of-the-art TVM/Ansor autotuning compiler framework shows higher performance for matrix-vector products, nearly comparable performance for matrix-matrix multiplication, but lower performance for convolutions. The paper presents insights into the limitations of the MLIR/IREE infrastructure that currently impact the performance achievable with PEAK.