linalg.generic and indexing maps — the one idea behind ML compilers

$ lower_std 01_indexing_maps.mlir > out/im.mlir && mrun out/im.mlir | grep -v Unranked

[[2,   4,   6], 
 [8,   10,   12]]
[[1,   4], 
 [2,   5], 
 [3,   6]]
[6,  15]
[[11,   22,   33], 
 [14,   25,   36]]
[[14,   32], 
 [32,   77]]
add(a,a); transpose(a); rowsum(a) = [1+2+3, 4+5+6]; bias(a,b); a @ transpose(a).

$ mlir-opt 01_indexing_maps.mlir --linalg-specialize-generic-ops | grep -oE 'linalg\.[a-z_]+' | sort | uniq -c

      2 linalg.elementwise
      2 linalg.fill
      1 linalg.generic
      1 linalg.matmul
      1 linalg.transpose
      1 linalg.yield
The pass pattern-matched maps+body: our hand-written matmul generic IS
linalg.matmul, the transpose IS linalg.transpose, add and bias became
linalg.elementwise. The rowsum stayed generic (there's linalg.reduce, but
the specializer doesn't target it). Named ops matter because some
transforms and ukernels (IREE's mmt4d, and vendor matmul-to-ukernel passes)
match on the NAME.

$ mlir-opt 01_indexing_maps.mlir --linalg-specialize-generic-ops --linalg-generalize-named-ops | grep -oE 'linalg\.[a-z_]+' | sort | uniq -c

      7 linalg.generic
      7 linalg.yield
a vendor MLIR backend runs `createLinalgGeneralizePass()` in its pipeline for
exactly this reason: one form to match against.

$ mlir-opt 01_indexing_maps.mlir --one-shot-bufferize=bufferize-function-boundaries --convert-linalg-to-loops | sed -n '/@matmul_generic/,/^  }/p' | sed -n '12,25p'

    }
    scf.for %arg2 = %c0 to %c2 step %c1 {
      scf.for %arg3 = %c0 to %c2 step %c1 {
        scf.for %arg4 = %c0 to %c3 step %c1 {
          %0 = memref.load %arg0[%arg2, %arg4] : memref<2x3xf32, strided<[?, ?], offset: ?>>
          %1 = memref.load %arg1[%arg4, %arg3] : memref<3x2xf32, strided<[?, ?], offset: ?>>
          %2 = memref.load %alloc[%arg2, %arg3] : memref<2x2xf32>
          %3 = arith.mulf %0, %1 : f32
          %4 = arith.addf %2, %3 : f32
          memref.store %4, %alloc[%arg2, %arg3] : memref<2x2xf32>
        }
      }
    }
    return %alloc : memref<2x2xf32>
Three loops in map order (m, n, k), loads at exactly the positions the
maps say ([m,k], [k,n], [m,n]). This is the NAIVE lowering: C[m][n] is
loaded and stored on every k iteration, B is walked down a column (stride
= row length, cache-hostile). Track 04 replaces this with tiled,
vectorized code and measures the difference.