Source-linked AI summary

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

Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, Oleksandr Zinenko

arXiv:2002.11054v2cs.PLcs.LG

TL;DR

Compiler systems often duplicate high-level IR infrastructure because existing platforms do not flexibly span abstraction levels and domains. MLIR addresses this with extensible, SSA-based IR structures and reusable mechanisms, and its applications and adoption demonstrate broad applicability while abstraction best practices remain unsettled.

  • Problem

    Compiler infrastructure often lacks flexible support for high-level abstractions, leaving languages to build domain-specific IRs and duplicate compiler technology.

  • Method

    MLIR provides customizable SSA-based structures, dialects, progressive lowering, and reusable compiler infrastructure for diverse abstractions and applications.

  • Results

    MLIR’s applications span polyhedral code generation, Fortran IR design, and domain-specific compilation, demonstrating applicability across domains.

  • Takeaways & Limitations

    MLIR adoption included more than 26 dialects in development and 7 projects replacing custom infrastructure, indicating broad participation and use.

  • Takeaways & Limitations

    MLIR provides little guidance on which abstractions work best, and several years of research may be needed to establish best practices.

Abstract

from arXiv · show

This work presents MLIR, a novel approach to building reusable and extensible compiler infrastructure. MLIR aims to address software fragmentation, improve compilation for heterogeneous hardware, significantly reduce the cost of building domain specific compilers, and aid in connecting existing compilers together. MLIR facilitates the design and implementation of code generators, translators and optimizers at different levels of abstraction and also across application domains, hardware targets and execution environments. The contribution of this work includes (1) discussion of MLIR as a research artifact, built for extension and evolution, and identifying the challenges and opportunities posed by this novel design point in design, semantics, optimization specification, system, and engineering. (2) evaluation of MLIR as a generalized infrastructure that reduces the cost of building compilers-describing diverse use-cases to show research and educational opportunities for future programming languages, compilers, execution environments, and computer architecture. The paper also presents the rationale for MLIR, its original design principles, structures and semantics.

1 Introduction

MLIR addresses fragmented compiler infrastructure by making new abstraction levels cheap to define while supplying reusable infrastructure for diverse domains and hardware targets. The paper presents its design and reports applications and adoption indicating broad usability.

  • Existing compiler platforms commonly expose one abstraction level, while languages and machine-learning systems build separate domain-specific IRs for higher-level needs.
  • High engineering costs for domain-specific IRs can produce slow compilation, buggy implementations, weak diagnostics, and poor debugging experiences.
  • MLIR makes new abstraction levels cheap by standardizing SSA-based structures, enabling declarative dialect definitions, and providing common compiler infrastructure.
  • The paper explores MLIR’s design points, applications across problems, and implications for language design and education.
  • The paper contributes a novel infrastructure, scalable and modular construction approaches, diverse applications, and shared development experience.
  • MLIR responds to fragmented compiler stacks by enabling shared infrastructure, progressive upgrades, heterogeneous compilation, and research into new abstractions.

2 Design Principles

MLIR’s design principles prioritize extensibility, retained high-level structure, progressive lowering, and declarative validation and transformation mechanisms. These choices support heterogeneous compilation but leave abstraction design and ecosystem compatibility as open challenges.

  • MLIR minimizes built-in concepts and keeps abstractions customizable, allowing machine-learning graphs, ASTs, polyhedral representations, CFGs, and instruction-level IRs.
  • Customizability risks internal fragmentation when independently designed abstractions are poorly compatible, motivating reusable abstractions beyond their initial scope.
  • SSA and nested regions represent dataflow and higher-level structures while supporting structured control flow, concurrency, closures, and heterogeneous compilation.
  • MLIR lets compilation algorithms choose nested regions or linearized control flow, departing from LLVM’s normalization-only orientation.
  • Progressive lowering moves programs through multiple abstraction levels in small steps to support varied platforms and programming models.
  • The design retains higher-level semantics and permits mixed abstraction levels so lower-level accelerator operations can coexist with reusable higher-level structure.
  • Extensible IR requires declarative validation, traceable source locations and transformations, and rewrite rules that can expose properties such as complexity and completion.
  • Translation validation and modern compiler testing remain open problems for an extensible compiler ecosystem.

3 IR Design Details

MLIR models compiler structure through extensible operations, dialects, types, attributes, regions, blocks, symbols, and locations. Its recursive, SSA-based IR supports nested control flow, customizable semantics, strict typing, and compiler traceability.

  • Operations are MLIR’s semantic unit, covering instructions, functions, and modules while carrying operands, results, attributes, regions, block arguments, and location information.
  • Attributes: Attributes provide typed compile-time static information in extensible key-value dictionaries, including constants and dialect-defined meanings.
  • Regions and blocks: An operation may contain regions, regions contain blocks, and blocks contain operations, enabling recursive nested structures.
  • Regions and blocks: Regions use terminators and successor blocks to define control-flow graphs, while standard SSA-based control flow operates within each region.
  • SSA: MLIR uses functional SSA: terminators pass values to typed successor block arguments instead of using φ nodes.
  • Visibility: Values must obey SSA dominance, nesting, and enclosing-operation restrictions, with isolated-from-above operations acting as scope barriers.
  • Dialects: Dialects group operations, attributes, and types under namespaces without imposing new semantics, enabling extensible operation support.
  • Type system: Types are user-extensible, may reference foreign type systems, and use strict equality checking without built-in conversion rules.

4 IR Infrastructure

MLIR combines extensible IR abstractions with declarative definitions, rewriting, verification, textual round-tripping, and flexible pass management. Its infrastructure is designed to make new abstractions and compiler transformations easier to define, inspect, test, and execute.

  • Infrastructure: MLIR provides infrastructure for defining dialects, Ops, pattern rewrites, verification, and reusable passes.These facilities support extensibility and ease of use when defining new abstractions and using MLIR as an optimization toolkit.
  • Operation description: ODS declaratively specifies an Op’s structure and verifier components, then generates interoperable C++ code.Definitions can include names, traits, arguments, results, type constraints, documentation, and custom textual forms.
  • Declarative rewrites: DRR expresses constrained source-to-target DAG equivalences declaratively and can combine generated C++ with handwritten rewrite patterns.The framework keeps common rewrites simple without restricting more complex transformations.
  • Pass management: The pass manager operates on arbitrary operations and supports concurrent IR traversal and modification through region isolation invariants.MLIR avoids specializing pass management to a fixed set of module, function, or loop operations.
  • Textual representation: MLIR’s textual IR fully reflects its in-memory representation, enabling debugging, comprehension, and separately testable compiler passes.Round-trippable input and output make transformations easy to trace, while the absence of hidden state preserves individual-pass behavior within the full pipeline.
  • Documentation and verification: Generated documentation reuses ODS descriptions and constraints that also generate verification code, helping documentation remain synchronized with runtime behavior.The generated material includes argument and result type constraints.

5 Evaluation: Applications of MLIR

MLIR is evaluated through diverse applications spanning machine-learning graphs, affine code generation, Fortran compilation, and specialized machine-learning compilers. These cases illustrate its extensibility, abstraction mixing, progressive lowering, dialect reuse, and practical performance gains.

  • Evaluation approach: MLIR’s primary evaluation metric is adoption across diverse projects, supported by community activity and detailed use cases.The evaluation emphasizes generality, extensibility, and implementation of the customizability design principle.
  • Community adoption: MLIR community activity included participants from 16 universities, 4 national laboratories, 14 multinational companies, and more than 100 industry developers.These figures describe workshop participation, company endorsement, and an LLVM Developer Meeting roundtable, respectively.
  • TensorFlow graphs: In TensorFlow, MLIR models high-level dataflow graphs and supports algebraic optimization, accelerator retargeting, mobile lowering, and native-code generation.TensorFlow graph nodes represent computations deployable across devices, including specialized hardware accelerators.
  • Affine dialect: The affine dialect combines static-control-flow loops and conditionals with structured multidimensional memory references for progressive lowering.Affine maps and integer sets are modeled as attributes, while operations apply affine restrictions to code.
  • Affine dialect differences: MLIR’s affine representation mixes typed SSA operations with polyhedral transformations, preserving loop structure and reducing the representation gap.This allows traditional compiler analyses and transformations to interleave with polyhedral transformations without raising into a drastically different representation.
  • Affine dialect differences: MLIR explicitly avoids polyhedron scanning because preserving loops supports compilation speed, unlike approaches relying on exponential-complexity algorithms.Existing polyhedral approaches are described as relying heavily on integer linear programming and polyhedron scanning.
  • Affine dialect results: Experience with the affine dialect shows usefulness across a wide range of code-generation projects.Its development provided practical exploration of MLIR’s design space.
  • Fortran IR: FIR uses MLIR to represent Fortran-specific semantics, including virtual dispatch tables, enabling robust devirtualization and reuse of language-independent dialects.Shared OpenMP and GPU-oriented dialects and passes support reuse across Fortran, C, and heterogeneous-platform workflows.

6 Consequences of the MLIR Design

MLIR shifts compiler engineering toward reusable abstractions, extensible operations, and transformations that can span dialects and abstraction levels. This design improves modularity but leaves abstraction design and best practices as active challenges.

  • Extensible design: MLIR reuses generic compiler infrastructure while allowing new operations, types, and dialects to model new abstractions.This design makes adding domain-specific constructs the primary extension mechanism.
  • Reusable compiler passes: Operation traits and privileged hooks let generic passes use operation properties, constant folding, and canonicalization across abstraction domains.Canonicalization patterns support extensible algebraic simplifications and subsume several specialized compiler passes.
  • Reusable compiler passes: Optimization interfaces keep dialect-specific transformation logic inside dialects while allowing generic passes such as inlining to operate across domains.Operations or dialects can register interfaces; otherwise the pass behaves conservatively.
  • Reusable compiler passes: Dialect-specific passes remain useful when transformations require full dialect semantics or specialized constraints.Custom machine-instruction scheduling is given as an example where generalization is unnecessary.
  • Mixing dialects together: Mixing operations from different dialects in one program enables reuse such as applying generic polyhedral transformations across problem domains.The affine dialect can combine with arithmetic and target-specific accelerator dialects, while OpenMP abstractions can be reused across source-language IRs.
  • Unopinionated design provides new challenges: MLIR’s flexible abstraction design provides little guidance about which IR abstractions work best, and established best practices remain under development.The authors describe continued learning and several years of research needed to fully understand the design points.

7 Related Work

MLIR occupies a broad position among compiler infrastructures, heterogeneous programming systems, metaprogramming frameworks, parser generators, and domain-specific code generators. Its distinguishing aim is to provide shared infrastructure across abstraction levels and domains rather than a single specialized flow.

  • Compiler infrastructures: MLIR resembles LLVM but targets richer first-class data structures and algorithms, including tensor algebra, graphs, and heterogeneous compilation.Its pattern-rewriting infrastructure supports composable local transformations and flexible lowering.
  • Heterogeneous compilation: Unlike specialized heterogeneous compilation systems, MLIR is positioned as reusable infrastructure that can support multiple current code-generation strategies.The related systems discussed include XLA, Glow, TVM, Halide, PolyMage, and other accelerator or polyhedral compilers.
  • Language and frontend systems: MLIR complements parser-generation and metaprogramming systems, while lacking general parser generation, AST construction, and AST modeling functionality itself.The paper identifies combining MLIR with systems such as ANTLR as a possible route from user input through code generation.

8 Conclusion and Future Work

The paper concludes that MLIR is a flexible, extensible compiler-construction infrastructure with applications across domains and implications for research and education. Future work extends its abstractions, frontends, transformations, and teaching uses while leaving several language-support challenges open.

  • Conclusion: MLIR is presented as a flexible and extensible infrastructure whose design, applicability, and research implications were demonstrated across important domains.The conclusion characterizes the work as both a concrete system description and an exploration of original engineering implications.
  • Future work: Future directions include richer data structures, symbolic-shape and rank-polymorphic operations, and transformations combining symbolic reasoning with dataflow and control flow optimization.The stated application areas include machine learning and high-performance computing.
  • Future work: MLIR’s open challenges include a C++ mid-level frontend and support for garbage-collected, higher-order, and polymorphic type systems with type inference.The paper identifies these as missing or unresolved capabilities for general-purpose languages.
  • Future work: First-class parallel constructs in MLIR can support higher-level transformations before lowering to LLVM, where regular transformations operate on lowered code.This is contrasted with the invasive and difficult layering required for parallelism and concurrency constructs in LLVM.
  • Education: Textual IR and additional optimization-visualization tooling could support compiler education, an area not commonly covered in undergraduate curricula.The paper specifically connects textual IR with demystifying high-performance compilation for students.
Loading 2002.11054v2…