Beginner’s Guide to MLIR’s linalg dialect

A beginner’s guide to MLIR’s linalg dialect: how to read linalg.generic’s indexing maps and iterator types, and lower it to scf.for loops by hand.

If you work on MLIR-based compilers chances are you heard of linalg dialect.
And if you work on a compiler which uses it (e.g. IREE), then understanding ins and outs of linalg.generic (pun intended) is a prerequisite for the job.

I have read the docs, the rationale, and watched a couple of relevant talks, but still was struggling to understand it.
Until I realized that linalg.generic is just a fancy notation for good old for-loops.

I set myself a goal to convert linalg.generic to a bunch of scf.fors by hand, and this article is a summary of my pen-and-paper exercise.

memref refresher

A side note on tensors vs memrefs: depending on which part of the compiler stack you are dealing with, you are more likely to encounter linalg on tensors as opposed to memrefs.
For the purposes of this article I chose to use memrefs as IMO they are a bit easier to work with and we don’t need to be concerned about intricacies of tensor and scf.for APIs.

The topics discussed here apply to both types interchangeably.

Roughly speaking, a memref is a multidimensional memory buffer, similar to C’s void *, but with a bit more metadata attached both at compile and run time:

• rank: the number of dimensions this memref “holds”. The rank is known at compile time (in the context of linalg, that is).

• dimension size or an extent: either static (non-negative integer) or dynamic (?). Static extent is known at compile time, the dynamic is only known at run time.

• element_type: this is the underlying type a memref holds. memref<i32> holds integers, memref<f32> holds floats, memref<vector<8xf64>> holds vectors of doubles.

There are several operations working with memrefs, for this article we only need three of them:

• memref.dim(%mem, index): returns the size of a specific dimension

• memref.load(%mem, [indexes]) : retrieves a specific value

• memref.store(%mem, %new_val, [indexes]) : “overrides” a specific value

Some examples:

Anatomy of a linalg.generic operation

A linalg.generic operation is a compact representation of K nested loops with N inputs ins and M outputs outs.

It consists of metadata such as indexing_maps with a total number of N+M affine_maps (one for each operand ins and outs), total K iterator_types (one for each loop) and a loop nest (a body) of the inner-most loop. The loop nest accepts N+M operands each corresponding to a specific ins / outs value. The loop nest does some work (possibly nothing) and yields M values “back”.
A linalg.generic operates on memref(or tensor) values with potentially distinct ranks and potentially distinct underlying types (Tin_0, ..., Tin_N, Tout_0,..., Tout_M).

Below is a generic scheme of a linalg.generic:

There are a lot of details and now we are going to look into each component in isolation. We do so in reverse order: loop nest, iterator types, and indexing maps.

Loop nest

The loop nest describes precisely what happens to inputs/outputs at each iteration of the loop.

The loop nest is a single (basic) block with N+M block arguments. Each block argument corresponds to a specific ins or outs operand and has the element type of the corresponding linalg.generic operand.
For example:

• ins(%A: memref<f32>) outs(%B: memref<f64>) → ^bb0(%a: f32, %b: f64): memrefs with different element types

• ins(%A: memref<4xf32>) outs(%B: memref<12x14xf64>) → ^bb0(%a: f32, %b: f64): ranks do not matter, block arguments still have element types

• ins(%A: memref<4xvector<8xf32>>) outs(%B: memref<4xvector<8xf64>>) → ^bb0(%a: vector<8xf32>, %b: vector<8xf64>): element types may be non-scalars

The loop nest always ends with a linalg.yield which itself has exactly M operands, with the same element type as each corresponding outs operand.

Iterator types

The iterator types (parallel or reduction) are defined for each nested loop: they depend on K and not on N or M .
Roughly speaking, they matter most for optimizations, e.g. parallel means that the correctness does not depend on the order in which this specific loop is running, and thus it can be reordered or parallelized.

For the purposes of this article, the exact iterator types can be ignored.

Indexing maps

This is perhaps the most information-dense and thus confusing part, at least its purpose isn’t that obvious by looking at the operation definition.

The indexing_maps provides the following info:

1. The number of nested loops. All the maps in the array share the same domain (the left-hand side (d0, d1, …, dK) part), i.e. a specific linalg.generic has K nested loops. In a way, each dimension dk can be seen as an induction variable (an index) of each nested loop. The actual names do not matter: affine_map<(i, j, k) → (…)> is a valid map. Thus, a linalg.generic with the i, j, k map represents the following loop:

2. The address/index computation for each data access. indexing_map array has an entry for each linalg.generic input and output. In the “anatomy” picture above, each %IN_n has a corresponding #in_n affine map, %OUT_m has a map #out_m respectively. A map acts as a “classical” function e.g. f(x) → y. To compute an address of a value accessed by an iteration of a loop, one needs to apply the affine map/function to the induction variables.

   In the first example below, an identity map <(d0) → (d0)> produces the same value: #id(%i) → %i, #id(%j) → %j, #id(42) → 42.
   The second example is a matmul using a naïve implementation: there are three nested loops with indexes i, j, k . At each iteration, the loop nest accesses values from A, B, C as you’d see in a text book:

   C[i, j] = A[i, k] * B[k, j] + C[i, j]

3. Upper bounds for the loops. All the scf.for so far conveniently omitted the upper bounds in the examples. And the reason is that determining the exact limits for each loop is a bit involved.
   Let’s take a look at a matmul example. In this case, there are three operands all of which are 2d matrices of floats: A shape is 4x8, B shape is 8x6 and C has a dynamic shape.
   Each indexing map corresponds to a linalg.generic operand, and the number of results of an affine map (the right-hand side after →) equals to the rank of a corresponding memref. In the example below m, k corresponds to 4x8, k, n corresponds to 8x6, and m, n corresponds to ?x?.

   

In other words, each “symbol” from an affine map corresponds to the dimensions of an operand, and thus the upper bound of each nested loop is determined by those dimensions:

The mapping must be consistent, i.e. the dims of each results must be equal:

• k equals 8 at both entries

• i equals to 4 or ? , or simply 4

• and j equals 6 or ?, or simply 6

Thus, even though %C has dynamic shape only known at run time, at compile time we can deduce that its only valid shape is 4x6 .

Converting linalg.generic to loops

Given the information above we can now convert any linalg.generic into loops by hand. To convert a linalg.generic to an scf.for loop we need to take several steps:

• determine the number of nested loops

• determine the upper bounds for each nested loop

• create the nested loops with a corresponding induction variable and a corresponding upper bound

• “remap” the loop nest arguments into memref.loads based on the induction variables (and indexing maps)

• copy the original loop body

• replace the terminating linalg.yield with memref.stores

Let’s consider a matmul example again, now with all the shapes being dynamic and only known at run time.

This operation has three nested loops with induction variables %i, %j, %k , each loop has a corresponding upper bound:

• %i_ub = memref.dim %A, 0

• %j_ub = memref.dim %B, 1

• %k_ub = memref.dim %A, 1

Loop nest arguments %a, %b, %c are loaded from the corresponding ins and outs %A, %B, %C based on the indexing maps:

• %a = memref.load %A[%i, %k]

• %b = memref.load %B[%k, %j]

• %c = memref.load %C[%i, %j]

The loop nest stays the same, except the final linalg.yield is replaced by a store:
memref.store %c, %C[%i, %j] .

Stitching it all together we have the following implementation:

You can also see in action on Compiler Explorer.

Missing bits

There are a few details that I omitted for brevity, from the top of my head:

• “syntactic sugar” aka named/elementwise operations

• important non-generic operations such as linalg.pack/unpack, linalg.index, and some more

• the indexing maps view is a simplified one: the maps are a bit more involved and not always simple index mappings, e.g. there could be constants or arithmetic operations

• linalg.generic operands are not always memrefs or tensors

• unranked tensors/memrefs (though they are not relevant for linalg dialect)

Perhaps there are some more topics I missed, but I hope with the information above you can go and fill in the gaps on your own.

Do not hesitate reaching out if you find mistakes or plain wrong claims in the article.