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