llvm-project/mlir/g3doc/Rationale.md

1129 lines
52 KiB
Markdown
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# MLIR Rationale
This document is intended to capture some of the alternatives considered and
open debates in the design of MLIR, along with the rationale for certain
decisions we made. This is not intended to be a "finely groomed" document - we
prefer the ability to dump in interesting tidbits without worrying too much
about their consistency or readability.
[TOC]
## Abstract
MLIR is a compiler intermediate representation with similarities to traditional
three-address SSA representations (like
[LLVM IR](http://llvm.org/docs/LangRef.html) or
[SIL](https://github.com/apple/swift/blob/master/docs/SIL.rst)), but which
introduces notions from the polyhedral loop optimization works as first class
concepts. This hybrid design is optimized to represent, analyze, and transform
high level dataflow graphs as well as target-specific code generated for high
performance data parallel systems. Beyond its representational capabilities, its
single continuous design provides a framework to lower from dataflow graphs to
high performance target specific code.
MLIR stands for one of "Multi-Level IR" or "Multi-dimensional Loop IR" or
"Machine Learning IR" or "Mid Level IR", we prefer the first. This document only
provides the rationale behind MLIR -- its actual
[specification document](LangRef.md) and other content is hosted elsewhere.
## Introduction and Motivation
The Multi-Level Intermediate Representation (MLIR) is intended for easy
expression and optimization of computations involving deep loop nests and dense
matrices of high dimensionality. It is thus well-suited to deep learning
computations in particular. Yet it is general enough to also represent arbitrary
sequential computation. The representation allows high-level optimization and
parallelization for a wide range of parallel architectures including those with
deep memory hierarchies --- general-purpose multicores, GPUs, and specialized
neural network accelerators.
MLIR uses ideas drawn from IRs of LLVM and Swift for lower level constructs
while combining them with ideas from the polyhedral abstraction to represent
loop nests, multi-dimensional data (tensors), and transformations on these
entities as first class concepts in the IR.
MLIR is a multi-level IR, i.e., it represents code at a domain-specific
representation such as HLO or TensorFlow graphs, all the way down to the machine
level. MLIR is able to represent arbitrary control flow and arbitrary data
accesses, and is general enough to represent nearly all sequential computation.
This is a key distinction from existing polyhedral representation
implementations (such LLVM [Polly](https://polly.llvm.org/)) that are able to
use the polyhedral abstraction in a way isolated from the LLVM IR and only for
affine loop nests, i.e., portions of the code where array accesses, loop bounds,
and conditionals are regular (involve linear functions of loop iterators and
constant symbols). The presence of statically unpredictable data accesses or
control flow does not preclude representation in MLIR, but only limits to a
certain extent the ability to reason about and apply transformations using the
polyhedral abstraction.
Maps, sets, and relations with affine constraints are the core structures
underlying a polyhedral representation of high-dimensional loop nests and
multi-dimensional arrays. These structures are represented as textual
expressions in a form close to their mathematical form. These structures are
used to capture loop nests, tensor data structures, and how they are reordered
and mapped for a target architecture. All structured or "conforming" loops are
captured as part of the polyhedral information, and so are tensor variables,
their layouts, and subscripted accesses to these tensors in memory.
The information captured in the IR allows a compact expression of all loop
transformations, data remappings, explicit copying necessary for explicitly
addressed memory in accelerators, mapping to pre-tuned expert written
primitives, and mapping to specialized vector instructions. Loop transformations
that can be easily implemented include the body of affine transformations: these
subsume all traditional loop transformations (unimodular and non-unimodular)
such as loop tiling, interchange, permutation, skewing, scaling, relative
shifting, reversal, fusion, and distribution/fission. Transformations on data
layout such as padding and transforming to blocked layouts are also captured.
The design of the IR allows a progressive lowering to target-specific forms.
Besides high-level transformations for loop nests and data layout that a typical
mid-level optimizer is expected to deal with, MLIR is also designed to perform
certain low-level scheduling and mapping decisions that a typical backend IR is
entrusted with: these include mapping to specialized vector instructions,
auto-vectorization, and software pipelining. The need to support these
transformations stems from the fact that neural network accelerators have
specialized units that deal with large chunks of data whose computation maps
back to chunks of more than one loop of the loop nests as viewed by a program at
a level closer to the original specification. Such specialized units or
instructions operate on multidimensional data chunks from a programmer's
viewpoint. It thus makes it hard or infeasible for a backend operating on a very
low-level IR close to assembly to lift and reconstruct loops and perform such a
mapping. This is in contrast to classic instruction selection and scheduling in
today's compilers that primarily only deals with the body of the innermost loop.
MLIR also facilitates automatic mapping to expert pre-tuned primitives or vendor
libraries operating on data at higher levels (or at the highest level) of the
memory hierarchy.
**Strengths**
* MLIR is closed under the kind of transformations needed to lower to TPUs;
MLIR can be used to represent both the input and output of emitters
* MLIR allows us to build modular and reusable target independent and target
dependent passes - since each pass/emitter can read in another's output.
## Design Decisions
This section sheds light on some of the design decisions -- some of these are
indirectly implied by the specification document.
### Loads and stores
The 'load' and 'store' instructions are specifically crafted to fully resolve to
an element of a memref. These instructions take as arguments n+1 indices for an
n-ranked tensor. This disallows the equivalent of pointer arithmetic or the
ability to index into the same memref in other ways (something which C arrays
allow for example). Furthermore, in an affine construct, the compiler can follow
use-def chains (e.g. through
[affine.apply instructions](Dialects/Affine.md#affineapply-operation)) to
precisely analyze references at compile-time using polyhedral techniques. This
is possible because of the
[restrictions on dimensions and symbols](Dialects/Affine.md#restrictions-on-dimensions-and-symbols).
A scalar of element-type (a primitive type or a vector type) that is stored in
memory is modeled as a 0-d memref. This is also necessary for scalars that are
live out of for loops and if conditionals in a function, for which we don't yet
have an SSA representation --
[an extension](#mlfunction-extensions-for-"escaping-scalars") to allow that is
described later in this doc.
### Symbols and types
The current MLIR disallows use of symbols in types. For example, when a tensor
or memref dimension is statically unknown, it is denoted in the type as '?'. An
SSA symbol is then bound to it when a memref is created. The actual value of the
unknown dimension can be queried using the "dim" builtin as shown below.
Example:
```mlir {.mlir}
func foo(...) {
%A = alloc <8x?xf32, #lmap> (%N)
...
call bar(%A) : (memref<8x?xf32, #lmap>)
}
func bar(%A : memref<8x?xf32, #lmap>) {
// Type of %A indicates that %A has dynamic shape with 8 rows
// and unknown number of columns. The number of columns is queried
// dynamically using dim instruction.
%N = dim %A, 1 : memref<8x?xf32, #lmap>
affine.for %i = 0 to 8 {
affine.for %j = 0 to %N {
// A[i,j] += 1
%s1 = load %A [%i, %j] : memref<8x?xf32, #lmap>
%s2 = add %s1, 1
store %s2 to %A [%i, %j] : memref<8x?xf32, #lmap>
}
}
return
}
```
An alternative design is to embed the reference to symbols directly in the
type - memref<8x%Nxf32>. We went for the current approach in MLIR because it
simplifies the design --- types remain immutable when the values of symbols
change.
### Block Arguments vs PHI nodes
MLIR Regions represent SSA using "[block arguments](LangRef.md#blocks)" rather
than [PHI instructions](http://llvm.org/docs/LangRef.html#i-phi) used in LLVM.
This choice is representationally identical (the same constructs can be
represented in either form) but block arguments have several advantages:
1. LLVM PHI nodes always have to be kept at the top of a block, and
transformations frequently have to manually skip over them. This is defined
away with BB arguments.
1. LLVM has a separate function Argument node. This is defined away with BB
arguments, because the arguments to the entry block serve this purpose.
1. Blocks of PHI nodes in LLVM execute atomically, which is surprising and
super confusing to compiler engineers and it is easy to introduce bugs with
this (very related to the
"[lost copy](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.524.5461&rep=rep1&type=pdf)"
problem in SSA lowering literature.) With the BB argument representation,
this confusion is defined away.
1. The entry list of PHI nodes in LLVM are unordered, and some blocks have
thousands of predecessors (e.g. unwind blocks). This can cause long compile
time problems because transformations have to linearly scan this list. This
is defined away with BB argument representation.
1. LLVM has no way to represent values that are available only in one successor
but not the other, e.g. its invoke instruction cannot produce the exception
value JUST on the exception edge. Instead, the
[landingpad instruction](http://llvm.org/docs/LangRef.html#landingpad-instruction)
is a hack used to represent this. MLIR doesn't make use of this capability,
but SIL uses it extensively, e.g. in the
[switch_enum instruction](https://github.com/apple/swift/blob/master/docs/SIL.rst#switch-enum).
For more context, block arguments were previously used in the Swift
[SIL Intermediate Representation](https://github.com/apple/swift/blob/master/docs/SIL.rst),
and described in
[a talk on YouTube](https://www.youtube.com/watch?v=Ntj8ab-5cvE). The section of
interest
[starts here](https://www.google.com/url?q=https://youtu.be/Ntj8ab-5cvE?t%3D596&sa=D&ust=1529450150971000&usg=AFQjCNFQHEWL7m8q3eO-1DiKw9zqC2v24Q).
### Index type disallowed in vector/tensor/memref types
Index types are not allowed as elements of `vector`, `tensor` or `memref` type.
Index types are intended to be used for platform-specific "size" values and may
appear in subscripts, sizes of aggregate types and affine expressions. They are
also tightly coupled with `affine.apply` and load/store operations; having
`index` type is a necessary precondition of a value to be acceptable by these
operations. While it may be useful to have `memref<?xindex>` to express indirect
accesses in MLFunctions, e.g. sparse matrix manipulations or lookup tables, it
creates problems MLIR is not ready to address yet. MLIR needs to internally
store constants of aggregate types and emit code operating on values of those
types, which are subject to target-specific size and alignment constraints.
Since MLIR does not have a target description mechanism at the moment, it cannot
reliably emit such code. Moreover, some platforms may not support vectors of
type equivalent to `index`.
Indirect access use cases can be alternatively supported by providing and
`index_cast` instruction that allows for conversion between `index` and
fixed-width integer types, at the SSA value level. It has an additional benefit
of supporting smaller integer types, e.g. `i8` or `i16`, for small indices
instead of (presumably larger) `index` type.
### Bit width of a non-primitive types and `index` is undefined
The bit width of a compound type is not defined by MLIR, it may be defined by a
specific lowering pass. In MLIR, bit width is a property of certain primitive
_type_, in particular integers and floats. It is equal to the number that
appears in the type definition, e.g. the bit width of `i32` is `32`, so is the
bit width of `f32`. The bit width is not _necessarily_ related to the amount of
memory (in bytes) or the size of register (in bits) that is necessary to store
the value of the given type. These quantities are target and ABI-specific and
should be defined during the lowering process rather than imposed from above.
For example, `vector<3xi57>` is likely to be lowered to a vector of four 64-bit
integers, so that its storage requirement is `4 x 64 / 8 = 32` bytes, rather
than `(3 x 57) ceildiv 8 = 22` bytes as can be naively computed from the
bitwidth. Individual components of MLIR that allocate space for storing values
may use the bit size as the baseline and query the target description when it is
introduced.
The bit width is not defined for dialect-specific types at MLIR level. Dialects
are free to define their own quantities for type sizes.
### Signless types
Integers in the builtin MLIR type system have a bitwidth (note that the `index`
type has a symbolic width equal to the machine word size), but they do not have
an intrinsic sign. This means that the "standard ops" operation set has things
like `addi` and `muli` which do two's complement arithmetic, but some other
operations get a sign, e.g. `divis` vs `diviu`.
LLVM uses the [same design](http://llvm.org/docs/LangRef.html#integer-type),
which was introduced in a revamp rolled out
[in the LLVM 2.0 integer type](http://releases.llvm.org/2.0/docs/LangRef.html#t_derived).
Prior to that, from
[LLVM 1.0](http://releases.llvm.org/1.0/docs/LangRef.html#t_classifications) to
[1.9](http://releases.llvm.org/1.9/docs/LangRef.html#t_classifications), LLVM
uses signed types like "sbyte" and "ubyte". This shift was important and has
served LLVM well over the years. The reason this is important is that it is a
good thing for an intermediate representation to represent the same computation
with the same instruction. Signed types got in the way, because (e.g.) an "add
of an sbyte" does the same computation as an "add of a ubyte", but the type
system made them look artificially different. This split also required casts
like "cast from sbyte to ubyte" which do nothing at the machine level. Removing
signs from the type system eliminated these problems, making the compiler
simpler.
More information about this split is available in an old
[talk on youtube](https://www.youtube.com/watch?v=VeRaLPupGks) talking about
LLVM 2.0.
Note that this rationale only applies to the "standard ops" dialect in which we
can express an opinion about its design. Other dialects generally try to model
an external system, and should aim to reflect its design as closely as possible.
### Splitting floating point vs integer operations
The MLIR "standard" operation set splits many integer and floating point
operations into different categories, for example `addf` vs `addi` and `cmpf` vs
`cmpi`
([following the design of LLVM](http://llvm.org/docs/LangRef.html#binary-operations)).
These instructions _are_ polymorphic on the number of elements in the type
though, for example `addf` is used with scalar floats, vectors of floats, and
tensors of floats (LLVM does the same thing with its scalar/vector types).
This split is important because floating point and integer operations are quite
different in practice: for example, floating point values include NaN's, so
[integer comparisons](http://llvm.org/docs/LangRef.html#icmp-instruction) and
[floating point comparisons](http://llvm.org/docs/LangRef.html#fcmp-instruction)
should use different comparison opcodes. On the arithmetic side of things,
floating point operations support rounding modes, floating point contractions,
["fast math"](http://llvm.org/docs/LangRef.html#fadd-instruction), and integers
may want to have two's complement overflow behavior or be undefined on
[various forms of wrapping](http://llvm.org/docs/LangRef.html#add-instruction)
for performance.
We are a long way from this sort of thing being a priority to care about in
MLIR, but since we have experience and know the right way to do this, we'd
rather design it in from the beginning.
Note that this rationale only applies to the "standard ops" dialect in which we
can express an opinion about its design. Other dialects generally try to model
an external system, and should aim to reflect its design as closely as possible.
### Specifying sign in integer comparison operations
Since integers are [signless](#signless-types), it is necessary to define the
sign for integer comparison operations. This sign indicates how to treat the
foremost bit of the integer: as sign bit or as most significant bit. For
example, comparing two `i4` values `0b1000` and `0b0010` yields different
results for unsigned (`8 > 3`) and signed (`-8 < 3`) interpretations. This
difference is only significant for _order_ comparisons, but not for _equality_
comparisons. Indeed, for the latter all bits must have the same value
independently of the sign. Since both arguments have exactly the same bit width
and cannot be padded by this operation, it is impossible to compare two values
whose bit representations would differ while the values are interpreted as
equal.
### Specifying comparison kind as attribute
Unlike arithmetic, comparison operators share several common properties, e.g.
they cannot be considered associative. In practice, comparisons are sometimes
implemented by the same instruction or its variants so it makes sense to group
them together at the IR level.
An alternative would be introducing ten distinct operators for all currently
supported kinds of integer comparisons. These operators would have increased the
number of "reserved" names used by standard operations as well as the size of
the C++ API while their implementations would have been mostly identical.
The comparison kind is internally an integer attribute. However, for the sake of
readability by humans, custom assembly form accepts string literals that are
mapped to the underlying integer values: `cmpi "eq", %lhs, %rhs` better implies
integer equality comparison than `cmpi 0, %lhs, %rhs` where it is unclear what
gets compared to what else. This syntactic sugar is possible thanks to parser
logic redefinitions for custom assembly form of non-builtin operations.
Supporting it in the full notation would have required changing how the main
parsing algorithm works and may have unexpected repercussions. While it had been
possible to store the predicate as string attribute, it would have rendered
impossible to implement switching logic based on the comparison kind and made
attribute validity checks (one out of ten possible kinds) more complex.
### 'select' operation to implement min/max
Although `min` and `max` operations are likely to occur as a result of
transforming affine loops in ML functions, we did not make them first-class
operations. Instead, we provide the `select` operation that can be combined with
`cmpi` to implement the minimum and maximum computation. Although they now
require two operations, they are likely to be emitted automatically during the
transformation inside MLIR. On the other hand, there are multiple benefits of
introducing `select`: standalone min/max would concern themselves with the
signedness of the comparison, already taken into account by `cmpi`; `select` can
support floats transparently if used after a float-comparison operation; the
lower-level targets provide `select`-like instructions making the translation
trivial.
This operation could have been implemented with additional control flow: `%r =
select %cond, %t, %f` is equivalent to
```mlir
^bb0:
br_cond %cond, ^bb1(%t), ^bb1(%f)
^bb1(%r):
```
However, this control flow granularity is not available in the ML functions
where min/max, and thus `select`, are likely to appear. In addition, simpler
control flow may be beneficial for optimization in general.
### Regions
#### Attributes of type 'Block'
We considered representing regions through `ArrayAttr`s containing a list of a
special type `IRBlockAttr`, which in turn would contain a list of operations.
All attributes in MLIR are uniqued within the context, which would make the IR
inside the regions immortal for no good reason.
#### Use "inlined" functions as regions
We considered attaching a "force-inline" attribute on a function and/or a
function `call` operation. Even the minimal region support (use cases in
affine.for and affine.if existing before the regions) requires access to the
values defined in the dominating block, which is not supported by functions.
Conceptually, function bodies are instances of regions rather than the inverse;
regions can also be device kernels, alternative sections, etc.
#### Dedicated `region` operation
This would mean we have a special kind of operation that is allowed to have
regions while other operations are not. Such distinction is similar to the
Stmt/Op difference we have had and chose to remove to make the IR simpler and
more flexible. It would also require analyses and passes to consider the
interplay between operations (e.g., an `affine.for` operation must be followed
by a region operation). Finally, a region operation can be introduced using the
current implementation, among other operations and without being special in any
sense.
#### Explicit capture of the values used in a region
Being able to use values defined outside the region implies that use-def chains
may contain uses from different nested regions. Consequently, IR transformations
and analyses can pull the instruction defining the value across region
boundaries, for example in case of TableGen-defined canonicalization patterns.
This would not be the case if all used values had been passed as region
arguments. One of the motivations for introducing regions in the IR is precisely
to enable cross-region analyses and transformations that are simpler than
inter-procedural transformations. Having uses from different regions appear in
the same use-def chain, contrary to an additional data structure maintaining
correspondence between function call arguments as uses of the original
definitions and formal arguments as new definitions, enables such
simplification. Since individual operations now belong to blocks, which belong
to regions, it is always possible to check if the definition of the value
belongs to the same region as its particular use. The risk is that any IR
traversal will need to handle explicitly this situation and it is easy to forget
a check (or conversely it isnt easy to design the right check in a tablegen
pattern for example): traversing use-def chains potentially crosses implicitly
semantic barriers, making it possible to unknowingly break region semantics.
This is expected to be caught in the verifier after the transformation.
At the same time, one may choose to pass certain or all values as region
arguments to explicitly break the use-def chains in the current proposal. This
can be combined with an attribute-imposed semantic requirement disallowing the
body of the region to refer to any value from outside it.
### Quantized integer operations
We haven't designed integer quantized operations in MLIR, but experience from
TensorFlow suggests that it is better to put information about the quantization
range/scale into the type itself, rather than have a single type like "qint8"
and put these on attributes of the operation.
There are a few ways to do this with MLIR, including at least:
* We could do the same thing TensorFlow does - and we will _have_ to support
that model to some extent for compatibility.
* We can encode the fp range of quantized integers directly into the types
when they are constants. The best practice on this seems to be to encode the
zero point as well as a scale factor. This ensures that 0.0 is always
exactly representable, e.g. `qi8<-1.42, 31.23x>`.
* We could theoretically encode dynamically determined ranges into the types
using something like `qi8<?,?>` with the bounds being determined through the
SSA dataflow graph dynamically - similar to how dynamic shapes are handled.
We will definitely need to do #1 for compatibility, we probably want to do #2,
and we should investigate #3 over time. That said, our short term plan is to get
more implementation experience with the rest of the system first, then come back
to re-examine the representation for quantized arithmetic when we have that
experience. When we do, we should chat with benoitjacob@ and
[read the paper](https://arxiv.org/abs/1712.05877).
### Dialect type extensions
This section describes the design decisions that shaped the dialect extensible
type system present in MLIR.
#### Reserving dialect type kinds
Dialects that wish to define type extensions must reserve a range of type kinds
within a '.def' file within the core IR library. This means that every dialect
wishing to define custom types must modify this file, but it guarantees that all
type casting checkings are performed in O(1) time.
#### Interactions between dialects
There are two different interactions between dialects that are important to
understand. When types of a dialect are:
* In operations of other dialects
- For standard/builtin operations, only standard/builtin types are
allowed. This restriction allows for operations to clearly understand
the invariants that they are working under.
- Outside of standard/builtin operations, dialects are expected to verify
the allowable operation types per operation.
* In types of other dialects
- For standard/builtin types, these types are allowed to contain types
from other dialects. This simplifies the type system and removes the
need for dialects to redefine all of the standard aggregate types, e.g.
tensor, as well as the memref type. Dialects are expected to verify that
a specific type is valid within a standard type, e.g. if a type can be
an element of a tensor.
- For dialect types, the dialect is expected to verify any type
invariants, e.g. if the standard tensor type can contain a specific type
of that dialect.
#### Separating builtin and standard types
Following the separation between the built-in and standard dialect, it makes
sense to separate built-in types and standard dialect types. Built-in types are
required for the validity of the IR itself, e.g. the function type (which
appears in function signatures and generic assembly forms of operations).
Integer, float, vector, memref and tensor types, while important, are not
necessary for IR validity.
#### Unregistered types
MLIR supports unregistered operations in generic assembly form. MLIR also
supports a similar concept for types. When parsing, if the dialect for dialect
type has not been registered the type is modeled as an 'OpaqueType'. This allows
for types to be round-tripped without needing to link in the dialect library
that defined them. No additional information about opaque types, outside of
parsing/printing, will be available.
#### Dialect type syntax
Dialect extended types are represented as string literals wrapped inside of the
dialect namespace. This means that the parser delegates to the dialect for
parsing specific type instances. This differs from the representation of dialect
defined operations, of which have a identifier name that the parser uses to
identify and parse them.
This representation was chosen for several reasons:
##### Dialects must provide custom type parsers
Dialect type parsing cannot plug into the existing parser infrastructure as
operations do with the OpAsmParser/Printer. Operations have a defined syntax
structure that is the same across all dialects. Types, on the other hand, may
have many different, and sometimes conflicting, parsing constraints that would
be difficult/unmaintainable to provide within a single interface.
This also has the added benefit of encouraging dialects to reuse existing
external type parsers. For example, an LLVM dialect may provide an MLIR LLVM
type that is simply a wrapper around LLVM types. The LLVM dialect would then use
the existing LLVM type parsing infrastructure.
Example:
```mlir {.mlir}
%s = "foo"() : () -> !llvm<"i32*">
```
##### Types do not always have canonical names
Unlike operations, types generally do not have a formal canonical name. For
example, function types have no defined keyword and integer types are defined by
a regular expression to support arbitrary bitwidth. Dialects with existing type
systems, e.g. LLVM, are likely to provide wrappers around their existing type
systems. For these wrapper types there is no simple canonical name, it's logical
to think of these types as existing within the namespace of the dialect. If a
dialect wishes to assign a canonical name to a type, it can be done via
[type aliases](LangRef.md#type-aliases).
### Tuple types
The MLIR type system provides first class support for defining
[tuple types](LangRef.md#tuple-type). This is due to the fact that `Tuple`
represents a universal concept that is likely to, and has already begun to,
present itself in many different dialects. Though this type is first class in
the type system, it merely serves to provide a common mechanism in which to
represent this concept in MLIR. As such, MLIR provides no standard operations
for interfacing with `tuple` types. It is up to dialect authors to provide
operations, e.g. extract_tuple_element, to interpret and manipulate them. When
possible, operations should prefer to use multiple results instead. These
provide a myriad of benefits, such as alleviating any need for tuple-extract
operations that merely get in the way of analysis and transformation.
### Assembly forms
MLIR decides to support both generic and custom assembly forms under the
following considerations:
MLIR is an open system; it is designed to support modular and pluggable
dialects. Depending on whether there exists a corresponding dialect and whether
the dialect is plugged in, operations may or may not be registered into MLIR
system. Yet we still need a way to investigate these operations. So the generic
assembly form is mandated by this aspect of MLIR system. It provides a default
textual form for operations.
On the other hand, an assembly form is for assisting developers to investigate
the IR. The generic form serves as a safe fallback but it can be too verbose for
certain ops. Therefore, MLIR gives each dialect the choice to define a custom
assembly form for each operation according to the operation's semantics and
specific needs. The custom assembly form can de-duplicate information from the
operation to derive a more concise form, thus better facilitating the
comprehension of the IR.
## Examples
This section describes a few very simple examples that help understand how MLIR
represents computation.
### Non-affine control flow
```mlir {.mlir}
// A simple linear search in every row of a matrix
for (i=0; i<N; i++) {
for (j=0; j<N; j++) {
// dynamic control flow
if (a[i][j] == key) {
s[i] = j;
break;
}
}
}
```
The presence of dynamic control flow leads to an inner non-affine function
nested in an outer function that using affine loops.
```mlir {.mlir}
func @search(memref<?x?xi32 %A, <?xi32> %S, i32 %key) {
%ni = dim %A, 0 : memref<?x?xi32>
// This loop can be parallelized
affine.for %i = 0 to %ni {
call @search_body (%A, %S, %i) : (memref<?x?xi32>, memref<?xi32>, i32)
}
return
}
func @search_body(%A: memref<?x?xi32>, %S: memref<?xi32>, %key: i32) {
%nj = dim %A, 1 : memref<?x?xi32>
br ^bb1(0)
^bb1(%j: i32)
%p1 = cmpi "lt", %j, %nj : i32
br_cond %p1, ^bb2, ^bb5
^bb2:
%v = load %A[%i, %j] : memref<?x?xi32>
%p2 = cmpi "eq", %v, %key : i32
br_cond %p2, ^bb3(%j), ^bb4
^bb3(%j: i32)
store %j, %S[%i] : memref<?xi32>
br ^bb5
^bb4:
%jinc = addi %j, 1 : i32
br ^bb1(%jinc)
^bb5:
return
}
```
As per the [MLIR spec](LangRef.md), the restrictions on dimensions and symbol
identifiers to be used with the affine.apply operation only apply to accesses
inside `affine.for` and `affine.if` operations. However, an analysis of accesses
inside the called function (`@search_body`) is necessary to determine if the
`%i` loop could be parallelized: such function access analysis is calling
context sensitive.
### Non-affine loop bounds
Loop bounds that are not affine lead to a nesting of functions as shown below.
```c
for (i=0; i <N; i++)
  for (j=0; j<N; j++)
// non-affine loop bound for k loop
    for (k=0; k<pow(2,j); k++)
       for (l=0; l<N; l++) {
        // block loop body
        ...
       }
```
```mlir {.mlir}
func @outer_nest(%n) : (i32) {
affine.for %i = 0 to %n {
affine.for %j = 0 to %n {
call @inner_nest(%i, %j, %n)
}
}
return
}
func @inner_nest(%i: i32, %j: i32, %n: i32) {
%pow = call @pow(2, %j) : (f32, f32) -> f32
// TODO(missing cast from f32 to i32)
call @inner_nest2(%pow, %n)
return
}
func @inner_nest2(%m, %n) -> i32 {
affine.for %k = 0 to %m {
affine.for %l = 0 to %n {
...
}
}
return
}
```
### Reference 2D Convolution
The following example illustrates a reference implementation of a 2D
convolution, which uses an integer set `#domain` to represent valid input data
in a dilated convolution.
```mlir {.mlir}
// Dilation factors S0 and S1 can be constant folded if constant at compile time.
#domain = (d0, d1)[S0,S1,S2,S3]: (d0 % S0 == 0, d1 % S1 == 0, d0 >= 0, d1 >= 0,
S3 - d0 - 1 >= 0, S4 - d1 - 1 >= 0)
// Identity map (shown here for illustration).
#map0 = (d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)
// Affine map from output to input coordinate space.
// d0 = output_h, d1 = output_w, d2 = kernel_h, d3 = kernel_w
// S0 = h_stride, S1 = w_stride, S2 = h_kernel_dilation, S3 = w_kernel_dilation
// S4 = h_pad_low, S5 = w_pad_low
// %out0 = %0#1 * %h_stride + %0#4 * %h_kernel_dilation - %h_pad_low
// %out1= %0#2 * %w_stride + %0#5 * %w_kernel_dilation - %w_pad_low
#map1_0 = (d0, d1, d2, d3) [S0, S1, S2, S3, S4, S5] -> (d0 * S0 + d2 * S2 - %S4)
#map1_1 = (d0, d1, d2, d3) [S0, S1, S2, S3, S4, S5] -> (d1 * S1 + d3 * S3 - %S5)
// Semi-affine map to undilated input coordinate space.
// d0 = input_h, d1 = input_w, S0 = h_base_dilation, S1 = w_base_dilation.
#map2_0 = (d0, d1) [S0, S1] -> (d0 / S0)
#map2_1 = (d0, d1) [S0, S1] -> (d1 / S1)
// Conv2D shapes:
// input: [batch, input_height, input_width, input_feature]
// kernel: [kernel_height, kernel_width, input_feature, output_feature]
// output: [batch, output_height, output_width, output_feature]
func @conv2d(memref<16x1024x1024x3xf32, #lm0, vmem> %input,
memref<5x5x3x32xf32, #lm0, vmem> %kernel,
memref<16x512x512x32xf32, #lm0, vmem> %output) {
affine.for %b = 0 to %batch {
affine.for %oh = 0 to %output_height {
affine.for %ow = 0 to %output_width {
affine.for %of = 0 to %output_feature {
affine.for %kh = 0 to %kernel_height {
affine.for %kw = 0 to %kernel_width {
affine.for %if = 0 to %input_feature {
// Calculate input indices.
%1_0 = affine.apply #map1_0 (%0#1, %0#2, %0#4, %0#5)
[%h_stride, %w_stride, %h_kernel_dilation, %w_kernel_dilation,
%h_pad_low, %w_pad_low]
%1_1 = affine.apply #map1_1 (%0#1, %0#2, %0#4, %0#5)
[%h_stride, %w_stride, %h_kernel_dilation, %w_kernel_dilation,
%h_pad_low, %w_pad_low]
// Check if access is not in padding.
affine.if #domain(%1_0, %1_1)
[%h_base_dilation, %w_kernel_dilation, %h_bound, %w_bound] {
%2_0 = affine.apply #map2 (%1_0, %1_1)
%2_1 = affine.apply #map2 (%1_0, %1_1)
// Compute: output[output_indices] += input[input_indices] * kernel[kernel_indices]
call @multiply_accumulate(%input, %kernel, %output, %b, %oh, %ow, %of, %kh, %kw, %if, %2_0, %2_1)
}
}
}
}
}
}
}
}
return
}
```
TODO (Add more examples showing the IR for a variety of interesting cases)
## Design alternatives and extensions
This is a list of some design alternatives and extensions that we discussed in
detail but did not include in the spec or postponed them for future
consideration on demand. We will revisit these discussions when we have more
implementation experience and learn more about the challenges and limitations of
our current design in practice.
### Polyhedral code representation alternatives: schedule lists vs schedules trees vs affine loop/if forms
The current MLIR uses a representation of polyhedral schedules using a tree of
if/for loops. We extensively debated the tradeoffs involved in the typical
unordered polyhedral instruction representation (where each instruction has
multi-dimensional schedule information), discussed the benefits of schedule tree
forms, and eventually decided to go with a syntactic tree of affine if/else
conditionals and affine for loops. Discussion of the tradeoff was captured in
this document:
[ MLIR: The case for a simplified polyhedral form](RationaleSimplifiedPolyhedralForm.md).
At a high level, we have two alternatives here:
1. Schedule tree representation instead of an affine loop AST form: The current
proposal uses an affine loop and conditional tree form, which is syntactic
and with no separation of domains as sets and schedules as multidimensional
affine functions. A schedule tree form however makes polyhedral domains and
schedules a first class concept in the IR allowing compact expression of
transformations through the schedule tree without changing the domains of
instructions. Such a representation also hides prologues, epilogues, partial
tiles, complex loop bounds and conditionals making loop nests free of
"syntax". Cost models instead look at domains and schedules. In addition, if
necessary such a domain schedule representation can be normalized to
explicitly propagate the schedule into domains and model all the cleanup
code. An example and more detail on the schedule tree form is in the next
section.
1. Having two different forms of MLFunctions: an affine loop tree form
(AffineLoopTreeFunction) and a polyhedral schedule tree form as two
different forms of MLFunctions. Or in effect, having four different forms
for functions in MLIR instead of three: CFG Function,
AffineLoopTreeFunction, Polyhedral Schedule Tree function, and external
functions.
#### Schedule Tree Representation for MLFunctions
This representation is based on a simplified form of the domain/schedule
representation used by the polyhedral compiler community. Domains represent what
has to be executed while schedules represent the order in which domain elements
are interleaved. We model domains as non piece-wise convex integer sets, and
schedules as affine functions; however, the former can be disjunctive, and the
latter can be piece-wise affine relations. In the schedule tree representation,
domain and schedules for instructions are represented in a tree-like structure
which is called a schedule tree. Each non-leaf node of the tree is an abstract
polyhedral dimension corresponding to an abstract fused loop for each ML
instruction that appears in that branch. Each leaf node is an ML Instruction.
```mlir {.mlir}
// A tiled matmul code (128x128x128) represented in schedule tree form
// #map0 = (d0, d1, d2, d3, d4, d5) -> (128*d0 + d3, 128*d1 + d4, 128*d2 + d5)
#intset_ij = (i, j) [M, N, K] : i >= 0, -i + N - 1 >= 0, j >= 0, -j + N-1 >= 0
#intset_ijk = (i, j, k) [M, N, K] : i >= 0, -i + N - 1 >= 0, j >= 0,
-j + M-1 >= 0, k >= 0, -k + N - 1 >= 0)
func @matmul(%A, %B, %C, %M, %N, %K) : (...) { // %M, N, K are symbols
// t1, t2, t3, t4, t5, t6 are abstract polyhedral loops
mldim %t1 : {S1,S2,S3,S4,S5} floordiv (i, 128) {
mldim %t2 : {S1,S2,S3,S4,S5} floordiv (j, 128) {
// (%i, %j) = affine.apply (d0, d1) -> (128*d0, 128*d1) (%t1, %t2)
call dma_hbm_to_vmem(%C, %i, %j, %M, %N, %K)
with @intset_ij(%i, %j) [%M, %N, %K]
mldim %t3 : {S2,S3,S4,S5} floordiv (k, 128) {
// (%i, %j, %k) = affine.apply (d0, d1, d2)
// -> (128*d0, 128*d1, 128*d2) (%t1, %t2, %t3)
call dma_hbm_to_vmem(%A, ...) with #inset_ijk (%i, %j, %k) [%M, %N, %K]
// (%i, %j, %k) = affine.apply (d0, d1, d2)
// -> (128*d0, 128*d1, 128*d2) (%t1, %t2, %t3)
call dma_hbm_to_vmem(%B, ...) with #inset_ijk (%i, %j, %k) [%M, %N, %K]
mldim %t4 : {S4} i mod 128 {
mldim %t5 : {S4} j mod 128 {
mldim %t6 : {S4} k mod 128 {
// (%i, %j, %k) = affine.apply #map0 (%t1, %t2, %t3, %t4, %t5, %t6)
call matmul_body(A, B, C, %i, %j, %k, %M, %N, %K)
with #inset_ijk(%i, %j, %k) [%M, %N, %K]
} // end mld4im t6
} // end mldim t5
} // end mldim t4
} // end mldim t3
// (%i, %j) = affine.apply (d0, d1) -> (128*d0, 128*d1) (%t1, %t2)
call $dma_vmem_to_hbm_C ... with #intset(%i, %j) [%M, %N, %K]
} // end mldim t2
} // end mldim t1
return
}
```
### Affine Relations
The current MLIR spec includes affine maps and integer sets, but not affine
relations. Affine relations are a natural way to model read and write access
information, which can be very useful to capture the behavior of opaque external
library calls, high-performance vendor libraries, or user-provided / user-tuned
routines.
An affine relation is a relation between input and output dimension identifiers
while being symbolic on a list of symbolic identifiers and with affine
constraints on the identifiers.
Syntax:
``` {.ebnf}
// Affine relation definition at the top of file
affine-rel-def ::= affine-rel-id `=` affine-relation-inline
affine-rel-id ::= `##` prefixed-id
affine-relation-inline ::=
`(` input-dims `)` (`[` symbols `]`)? `->`
`(` output-dims `)` : affine-constraint-conjunction
input-dims ::= bare-id-list
output-dims ::= bare-id-list
symbols ::= bare-id-list
affine-rel ::= affine-rel-id | affine-relation-inline
// Usage
affine-rel-spec ::= affine-rel dim-and-symbol-use-list
```
All identifiers appearing in input-dims, output-dims, and symbol-dims are
pairwise distinct. All affine-constraint non-terminals in the above syntax are
allowed to contain identifiers only from input-dims, output-dims, and
symbol-dims.
Affine relations are used to model read, write, may_read, and may_write sets of
functions in the IR. The output dimension identifiers correspond to the data
dimensions.
Example:
```mlir {.mlir}
// read relation: two elements ( d0 <= r0 <= d0+1 )
##aff_rel9 = (d0) -> (r0) : r0 - d0 >= 0, d0 - r0 + 1 >= 0
func @count (memref<128xf32, (d0) -> (d0)> %A, i32 %pos) -> f32
reads: {%A ##aff_rel9 (%pos)}
writes: /* empty */
may_reads: /* empty */
may_writes: /* empty */ {
bb0 (%0, %1: memref<128xf32>, i64):
%val = load %A [(d0) -> (d0) (%pos)]
%val = load %A [(d0) -> (d0 + 1) (%pos)]
%p = mulf %val, %val : f32
return %p
}
```
### Regions
#### Making function definition an operation
MLIR supports values of a Function type. Instead of having first-class IR
concept for functions, one could define an operation with a body region that
defines a function value. The particularity of functions is that their names are
globally visible and can be referred to before being defined, unlike SSA values
that must be defined first. Implementing a "function definition" operation would
require to relax some of the SSA constraints in a region, and also make the IR
Module a region as well. It would also affect the core infrastructure (e.g.,
function passes) only for the sake of concept unification.
#### Having types on a region
Instead of inspecting the types of arguments of the first block, one could give
the region itself a type. This type would be redundant with block argument
types, which must have values and create room for type mismatches. While
functions do have types that are partly redundant with the arguments of the
first block in the function, this is necessary to support function declarations
that do not have a body which we can refer to in order to obtain the argument
types. A region is always contained in an operation or a function that can be
queried to obtain the “type” of the region if necessary.
A type on a region can be justified if Regions were to be considered separately
from the enclosing entity (operation or function) and had their own semantics
that should be checked.
#### Attaching attributes to regions
Regions could be annotated with dialect attributes to use attribute verification
hooks. An operation could take multiple regions as arguments, and each of them
may require different attributes. However, there are currently very few
practical cases where this would be necessary. Instead, one could simulate
per-region attributes with array attributes attached to the entity containing
the region (operation or function). This decreases the overall complexity of the
IR and enables more concise and op-specific forms, e.g., when all regions of an
op have the same attribute that can be only mentioned once. Since the semantics
of the region is entirely defined by the enclosing entity, it also makes sense
to have attributes attached to that entity rather than to the region itself.
This can be reconsidered in the future if we see a non-neglectable amount of use
cases.
### Read/Write/May_Read/May_Write sets for External Functions
Having read, write, may_read, and may_write sets for external functions which
include opaque ones, high-performance vendor libraries such as CuDNN, CuB, MKL,
FFT libraries, user-provided/optimized functions, or data movement runtimes such
as DMA ones is a powerful feature. It allows the compiler to perform analysis,
composition/transformation in the presence of such calls and with loops around
such calls on sub-tensors. For user-provided or custom hand-tuned functions, the
read/write/may_read/may_write sets could be provided a-priori by a user as part
of the external function signature or they could be part of a database.
TODO: Design this, and update to use function attribute syntax.
Example:
```mlir {.mlir}
##rel9 ( ) [s0] -> (r0, r1) : 0 <= r0 <= 1023, 0 <= r1 <= s0 - 1
func @cblas_reduce_ffi(memref<1024 x ? x f32, #layout_map0, hbm> %M) -> f32 [
reads: {%M, ##rel9() }
writes: /* empty */
may_reads: /* empty */
may_writes: /* empty */
]
func @dma_hbm_to_vmem(memref<1024 x f32, #layout_map0, hbm> %a,
offset, memref<1024 x f32, #layout_map0, vmem> %b,
memref<1024 x f32, #layout_map0> %c
) [
reads: {%M, ##rel9() }
writes: /* empty */
may_reads: /* empty */
may_writes: /* empty */
]
```
### Memref Extensions
1. Arbitrary polyhedral shapes for tensors: e.g., triangular shapes in tensor
dimensions where there is symmetry: use integer set (affine constraints) to
model tensor data space (instead of just extents). Requires some changes to
the IR and the in-memory form.
1. Layout maps
1. Allow piece-wise affine maps for layouts: allows clean modeling of
boundary cases for images/tensors through padding, wrapping, mirroring,
padding where padded values are the results of computation as opposed to
data, padding in the interior as opposed to just boundaries.
1. Allow many-to-one layout maps: Index and layout maps in the current
proposal are bijective. Extending them to many-to-one layout maps allows
cleaner(?) modeling of broadcast/reduce style computations while reusing
memory.
Proposal 2(a) requires non-trivial changes to the IR and the in-memory
representation. 2(b) requires no change, but impacts how cost models look at
index and layout maps.
### `affine.if` and `affine.for` Extensions for "Escaping Scalars"
We considered providing a representation for SSA values that are live out of
`if/else` conditional bodies and loop carried in `affine.for` loops. We
ultimately abandoned this approach due to its complexity. In the current design
of MLIR, scalar variables cannot escape for loops or if instructions. In
situations, where escaping is necessary, we use zero-dimensional tensors and
memrefs instead of scalars.
**TODO**: This whole section is obsolete and should be updated to use block
arguments and a yield like terminator in for/if instructions.
The abandoned design of supporting escaping scalars is as follows:
#### affine.for Instruction
Syntax:
``` {.ebnf}
[<out-var-list> =]
for %<index-variable-name> = <lower-bound> ... <upper-bound> step <step>
[with <in-var-list>] { <loop-instruction-list> }
```
out-var-list is a comma separated list of SSA values defined in the loop body
and used outside the loop body. in-var-list is a comma separated list of SSA
values used inside the loop body and their initializers. loop-instruction-list
is a list of instructions that may also include a yield instruction.
Example:
```mlir {.mlir}
// Return sum of elements in 1-dimensional mref A
func int32 @sum(%A : memref<?xi32>, %N : i32) -> (i32) {
%init = 0
%result = affine.for %i = 0 to N with %tmp(%init) {
%value = load %A[%i]
%sum = %value + %tmp
yield %sum
}
return %result
}
```
#### affine.if/else Instruction
Syntax:
``` {.ebnf}
<out-var-list> = affine.if (<cond-list>) {...} [else {...}]
```
Out-var-list is a list of SSA values defined by the if-instruction. The values
are arguments to the yield-instruction that occurs in both then and else clauses
when else clause is present. When if instruction contains only if clause, the
escaping value defined in the then clause should be merged with the value the
variable had before the if instruction. The design captured here does not handle
this situation.
Example:
```mlir {.mlir}
// Compute sum of half of the array
func int32 @sum_half(%A, %N) {
%s0 = 0
%s1 = affine.for %i = 1 ... N step 1 with %s2 (%s0) {
%s3 = if (%i >= %N / 2) {
%v0 = load %A[%i]
%s4 = %s2 + %v0
yield %s4
}
yield %s3
}
return %s1
}
```
### Multithreading the compiler
People want compilers to go fast, and one simple way to do that is to
multi-thread them. There are multiple strategies for this, but a simple one is
to optimize and compile separate functions in parallel. LLVM's original pass
manager anticipated this demand, and the CallGraphSCCPass manager is even
designed to support this as well, but unfortunately, a few early design
decisions in LLVM prevent this from ever happening. Instead, things like ThinLTO
are forced to split programs into separate LLVM modules/context and optimize
those chunks independently.
The problem is that LLVM has several objects in its IR that are globally uniqued
and also mutable: notably constants like `i32 0`. In LLVM, these constants are
`Value*r`'s, which allow them to be used as operands to instructions, and that
they also have SSA use lists. Because these things are uniqued, every `i32 0` in
any function share a use list. This means that optimizing multiple functions in
parallel won't work (at least without some sort of synchronization on the use
lists, which would be unbearably inefficient).
MLIR now supports a multithreaded pass manager. We do this through several
design choices:
1) MLIR makes use of extensive uniqued immutable data structures (affine
expressions, types, etc are all immutable, uniqued, and immortal). 2) constants
are defined in per-function pools, instead of being globally uniqued. 3)
functions themselves are not SSA values either, so they don't have the same
problem as constants. 4) FunctionPasses are copied (through their copy ctor)
into one instances per thread, avoiding sharing of local state across threads.
This allows MLIR function passes to support efficient multithreaded compilation
and code generation.