COMPASS: An Agent for MLIR Compilation Pass Pipeline Generation
摘要
MLIR is the state-of-the-art multi-level compilation infrastructure. It provides a reusable and extensible compiler intermediate representation (IR) framework that supports the collaboration of IRs across different levels. MLIR has been widely adopted in various compilers, forming a domain-specific compilation ecosystem. However, the features of MLIR also bring challenges to compilation pass pipeline arrangement. Managing mixed-level abstractions and diverse IR extensions requires significant expertise, which increases the complexity of developing MLIR compilation path. COMPASS is proposed as an agent to generate MLIR pass pipelines, aiming to simplify the construction process and reduce reliance on manual expertise. Confronted with uncertain compilation passes and fragmented domain-specific knowledge, we address the problem by decomposing it into a unified compilation methodology and an extensible knowledge base. To achieve this, COMPASS combines an MLIR knowledge base with a compilation path generator to automate pass pipeline arrangement. The knowledge base is constructed using a large language model (LLM), while the generator performs backtracking over a pass-selection tree. Evaluation on a range of MLIR-based compilers shows that COMPASS achieves a success rate exceeding 90% across 1,256 test cases, successfully generating pass pipelines and compiling source files to their corresponding target abstraction levels. These results demonstrate that COMPASS enables the automatic generation of pass pipelines for MLIR-based compilers, reducing the human effort involved in the MLIR ecosystem.