Skip to content

MLIR

  • Multi-Level Intermediate Representation
  • Not Machine Learning IR
  • Reusable and extensible compiler infrastructure
  • "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."

Dialects

  • MLIR's builtins are minimal
  • Dialects define operations, types and attributes
  • Multiple dialects co-exist in one program
  • The key to progressive lowering

Operations

Define

  • Uses Operation Definition Syntax (ODS) based on TableGen
tablegen
def LeakyReluOp: Op<"leaky_relu",
    // a list of traits used for verification and optimization.
    [NoSideEffect, SameOperandsAndResultType]> {
  let summary = "Leaky Relu operator";
  let description = [{ Element-wise Leaky ReLU operator x -> x >= 0 ? x : (alpha * x) }];
  let arguments = (ins AnyTensor:$input, F32Attr:$alpha);
  let results = (outs AnyTensor:$output);
}

Invoke

  • In MLIR language
  • With source location information
mlir
%result = "dialect.operation"(%arg0, %arg1) [%succ_block]
({
  // Region
}) {attribute="value" : type}
 : (arg1_type, arg2_type) -> (res_type)
 loc("example/file/path":12:1)

Regions and Blocks

  • Each region contains one or more blocks, which then contain operations
  • Each operation may have attached one or more regions
mlir
// Ops may have regions attached.
"affine.for"(%arg0) ({
  // Regions consist of a CFG of blocks with arguments.
  ^bb0(%arg4: index):
    // Blocks are lists of operations. (SSA style)
    // ...

    // Blocks end with a terminator Op.
    "affine.terminator"() : () -> ()
}) {lower_bound = () -> (0), step = 1 : index}
  : (index) -> ()
  • Blocks are connected via terminators, instead of ϕ nodes
mlir
// ... true branch
"cf.br" ^merge_block(%val_from_true)

// ... false branch
"cf.br" ^merge_block(%val_from_false)

^merge_block(%result: i32):
  // %result is either %val_from_true or %val_from_false
  ...

Attributes

  • Extensible via dialects
  • Values can be numbers, strings, or external data structures
mlir
// Attribute aliases can be forward-declared.
#map1 = ()[s0] -> (s0)

"affine.for"(%arg0) ({
  // ...
}) {
  lower_bound = () -> (0),
  step = 1 : index,
  upper_bound = #map1
} : (index) -> ()

Parallel Compilation

  • SSA use-def chains cannot cross the region boundaries - No global use-def chain
  • Global objects are referenced through symbol table entries
  • Constants are implemented as operations with associated attributes

meow

  • String (semantic) operation/attribute names are necessary
  • Dialects may implement custom printing and parsing formats for Operations
  • The MLIR syntax of operations has so many components

References