linalg indexing maps

Affine indexing maps map each loop iteration to operand coordinates and define the op's memory access pattern.

One map per operand

The generic form of matmul carries three maps over the loop dimensions (d0, d1, d2), which are (m, n, k):

#a = affine_map<(d0, d1, d2) -> (d0, d2)>   // A[m][k]
#b = affine_map<(d0, d1, d2) -> (d2, d1)>   // B[k][n]
#c = affine_map<(d0, d1, d2) -> (d0, d1)>   // C[m][n]

func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>,
                  %C: memref<?x?xf32>) {
  linalg.generic {
      indexing_maps = [#a, #b, #c],
      iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : memref<?x?xf32>, memref<?x?xf32>)
    outs(%C : memref<?x?xf32>) {
  ^bb0(%a: f32, %b: f32, %c: f32):
    %p = arith.mulf %a, %b : f32
    %s = arith.addf %c, %p : f32
    linalg.yield %s : f32
  }
  return
}

linalg.matmul is this contract with the maps and the payload fixed.

Reading a map

A map’s inputs are loop coordinates and its results are operand subscripts. (d2, d1) says the iteration (m, n, k) touches B at [k][n], so the map is the address function written in loop coordinates. Results count equals the operand rank; inputs count equals the iteration space rank. Getting either wrong fails verification.

Which operand sets the bound

The maps also decide where each loop bound comes from: a dimension is bounded by the shape of an operand whose map indexes it, so m is pinned by A and C, n by B and C, and k by A and B. The shapes have to agree. With static shapes a disagreement is a verifier error; with dynamic shapes the extents are SSA values, and nothing compares them until your reference check does.

Broadcast and transpose are map properties

Dropping a dimension means the operand does not vary with that loop: (d0, d1) -> (d0) is a rank-1 input reused along n, which is a broadcast with no extra op. Swapping results, (d1, d0), transposes the access instead of moving data. That is why a layout question is often answered by editing a map rather than inserting a transpose.

Things that break

A wrong result count fails verification, which is the easy case. The quiet case is a map that is structurally valid and names the wrong dimension: (d1, d0) where (d0, d1) was meant verifies and computes a different product. Check it at a non-square shape, because square operand shapes hide the swap entirely.

Sources