MLIR: A Compiler Infrastructure for the End of Moore’s Law

Title: MLIR: A Compiler Infrastructure for the End of Moore’s Law

Authors: Chris Lattner (Google), Mehdi Amini (Google), Uday Bondhugula (IISc), Albert Cohen (Google), Andy Davis (Google), Jacques Pienaar (Google), River Riddle (Google), Tatiana Shpeisman (Google), Nicolas Vasilache (Google), Oleksandr Zinenko (Google)

Abstract Summary:

MLIR presents a novel compiler infrastructure designed to address software fragmentation, improve compilation for heterogeneous hardware, and reduce the cost of building domain-specific compilers. It aims to facilitate the design and implementation of code generators, translators, and optimizers across various levels of abstraction and application domains.

Key Concepts:

1. Reusable and extensible compiler infrastructure
2. Heterogeneous hardware support
3. Domain-specific compilers
4. Intermediate Representation (IR)
5. Static Single Assignment (SSA)
6. Declarative system for IR dialects
7. Multithreaded compilation support
8. Progressive lowering

Problem Statement:

The main problem addressed by the paper is the high cost and complexity of building and maintaining domain-specific compilers and the lack of a unified infrastructure that can efficiently handle heterogeneous hardware and various levels of abstraction in compiler design.

Methods and Techniques:

• Static Single Assignment (SSA): Standardizing SSA-based IR data structures to simplify dataflow analysis and transformations.
• IR Dialects: Providing a declarative system for defining various IR dialects to support different abstraction levels and domains.
• Common Infrastructure: Offering a wide range of common compiler infrastructure components, such as parsing and printing logic, location tracking, and pass management, to reduce the engineering effort.
• Progressive Lowering: Allowing for gradual lowering from high-level representations to low-level code, facilitating better optimization and support for heterogeneous targets.

Key Results:

• Reduction in Cost: MLIR significantly reduces the cost of developing compilers by providing reusable components and infrastructure.
• Support for Heterogeneity: It efficiently supports heterogeneous compilation by enabling the creation of IRs that span different abstraction levels and hardware targets.
• Improved Compilation: The framework shows improved compilation performance and capability through diverse use cases, such as TensorFlow graph optimizations and polyhedral code generation.

Contributions and Innovations:

• Extensible Compiler Infrastructure: MLIR introduces a novel, extensible framework that supports the creation and evolution of domain-specific compilers.
• Dialect Support: It allows for the easy definition of new IR dialects, which can be mixed and matched within the same compilation pipeline.
• Reusability: The infrastructure supports reusable passes and transformations, enhancing efficiency and consistency across different compiler projects.

Future Work:

The authors suggest further research into areas such as:

• Enhanced support for more complex data structures and programming models.
• Integration with higher-level languages and more advanced IR designs.
• Further improvements in parallel and heterogeneous compilation.

Applications:

• Machine Learning Frameworks: Optimizing TensorFlow graphs for various hardware targets.
• High-Performance Computing: Applying polyhedral code generation techniques.
• Domain-Specific Compilers: Facilitating the creation of compilers for specific tasks, such as lattice regression.

Relevant Links:

• MLIR website 
• Mojo API for MLIR