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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

PEAK: Generating High-Performance Schedules in MLIR

  • Amir Mohammad Tavakkoli,
  • Sameeran Joshi,
  • Shreya Singh,
  • Yufan Xu,
  • P. Sadayappan,
  • Mary Hall

摘要

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.