FForge
← Labsmlir-lowering-01MLIR to LLVM IR●●○○

Progressive lowering, one pass at a time, down to AArch64 asm

mlir-lowering/01_progressive_lowering.mlirneeds: mlir-opt, mlir-translate, opt, llc7 steps
0/7 steps
what this lab teaches

"Lowering to LLVM" is not one step. It's a set of CONVERSION passes, each responsible for one dialect, and they have to meet in the middle. Watch a 12-line function (sum of relu) go through each one, and learn to see: * loops (scf.for) become a CFG of blocks — iter_args become BLOCK ARGUMENTS, which become LLVM PHI nodes; * during the transition, `builtin.unrealized_conversion_cast` glues converted and not-yet-converted values together; * a memref argument explodes into the 5-word descriptor (lab 01/02); * after mlir-translate you're in plain LLVM, and opt/llc take over. Then two "fast with X, slow without X" facts about the LLVM stage that you will hit in a production MLIR backend: FP REASSOCIATION and the TARGET TRIPLE.

evidence to leave with
One IR file walked to AArch64 asm one pass at a time, with the pass that first changed the vectorization decision named
walk it, one step at a time
run this
$ mlir-opt 01_progressive_lowering.mlir --convert-scf-to-cf | sed -n '/func.func/,/^  }/p'
what it printed on the machine that wrote this lab
  func.func @sum_relu(%arg0: memref<?xf32>) -> f32 {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %cst = arith.constant 0.000000e+00 : f32
    %dim = memref.dim %arg0, %c0 : memref<?xf32>
    cf.br ^bb1(%c0, %cst : index, f32)
  ^bb1(%0: index, %1: f32):  // 2 preds: ^bb0, ^bb2
    %2 = arith.cmpi slt, %0, %dim : index
    cf.cond_br %2, ^bb2, ^bb3
  ^bb2:  // pred: ^bb1
    %3 = memref.load %arg0[%0] : memref<?xf32>
    %4 = arith.maximumf %3, %cst : f32
    %5 = arith.addf %1, %4 : f32
    %6 = arith.addi %0, %c1 : index
    cf.br ^bb1(%6, %5 : index, f32)
  ^bb3:  // pred: ^bb1
    return %1 : f32
  }
^bb1(%0: index, %1: f32) is the loop header: its two block arguments are
the induction variable and the accumulator (the iter_arg). The two
`cf.br ^bb1(...)` edges pass the initial and updated values. That is
exactly SSA phi semantics, spelled differently.
run this
$ mlir-opt 01_progressive_lowering.mlir --convert-scf-to-cf --convert-arith-to-llvm | sed -n '/func.func/,/^  }/p'
what it printed on the machine that wrote this lab
  func.func @sum_relu(%arg0: memref<?xf32>) -> f32 {
    %0 = llvm.mlir.constant(0 : i64) : i64
    %1 = builtin.unrealized_conversion_cast %0 : i64 to index
    %2 = llvm.mlir.constant(1 : i64) : i64
    %3 = llvm.mlir.constant(0.000000e+00 : f32) : f32
    %dim = memref.dim %arg0, %1 : memref<?xf32>
    %4 = builtin.unrealized_conversion_cast %dim : index to i64
    cf.br ^bb1(%1, %3 : index, f32)
  ^bb1(%5: index, %6: f32):  // 2 preds: ^bb0, ^bb2
    %7 = builtin.unrealized_conversion_cast %5 : index to i64
    %8 = llvm.icmp "slt" %7, %4 : i64
    cf.cond_br %8, ^bb2, ^bb3
  ^bb2:  // pred: ^bb1
    %9 = memref.load %arg0[%5] : memref<?xf32>
    %10 = llvm.intr.maximum(%9, %3) : (f32, f32) -> f32
    %11 = llvm.fadd %6, %10 : f32
    %12 = llvm.add %7, %2 : i64
    %13 = builtin.unrealized_conversion_cast %12 : i64 to index
    cf.br ^bb1(%13, %11 : index, f32)
  ^bb3:  // pred: ^bb1
    return %6 : f32
  }
arith.* are gone (llvm.icmp, llvm.fadd, llvm.intr.maximum), but memref.dim,
memref.load and the block arguments still want `index`. The converter
inserted `unrealized_conversion_cast index <-> i64` at every boundary.
These casts are PROMISES: "a later pass will convert the other side".
`reconcile-unrealized-casts` at the end deletes pairs that cancel out.
If one survives, the pipeline is incomplete (lab 02/02 shows the error).
run this
$ mlir-opt 01_progressive_lowering.mlir --convert-scf-to-cf --convert-arith-to-llvm --finalize-memref-to-llvm --convert-func-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts > out/sum_relu.llvm.mlir && grep -c unrealized out/sum_relu.llvm.mlir
what it printed on the machine that wrote this lab
0
Zero casts left. (grep -c exits 1 on zero matches; that's fine.)
run this
$ mlir-translate --mlir-to-llvmir out/sum_relu.llvm.mlir -o out/sum_relu.ll && sed -n '/^define/,/^}/p' out/sum_relu.ll
what it printed on the machine that wrote this lab
define float @sum_relu(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4) {
  %6 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %0, 0
  %7 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %6, ptr %1, 1
  %8 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %7, i64 %2, 2
  %9 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, i64 %3, 3, 0
  %10 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %9, i64 %4, 4, 0
  %11 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %10, 3
  %12 = alloca [1 x i64], i64 1, align 8
  store [1 x i64] %11, ptr %12, align 4
  %13 = getelementptr [1 x i64], ptr %12, i32 0, i64 0
  %14 = load i64, ptr %13, align 4
  br label %15

15:                                               ; preds = %19, %5
  %16 = phi i64 [ %25, %19 ], [ 0, %5 ]
  %17 = phi float [ %24, %19 ], [ 0.000000e+00, %5 ]
  %18 = icmp slt i64 %16, %14
  br i1 %18, label %19, label %26

19:                                               ; preds = %15
  %20 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %10, 1
  %21 = getelementptr inbounds nuw float, ptr %20, i64 %16
  %22 = load float, ptr %21, align 4
  %23 = call float @llvm.maximum.f32(float %22, float 0.000000e+00)
  %24 = fadd float %17, %23
  %25 = add i64 %16, 1
  br label %15

26:                                               ; preds = %15
  ret float %17
}

The descriptor struct is rebuilt with insertvalue (and even spilled to an alloca) just to read one field. This looks terrible and costs nothing: opt's SROA/instcombine erase it. Do NOT judge MLIR output before -O2. Note the two `phi` nodes: they are the block arguments of STEP 1.

First make the reassoc variant (one flag on one op):

run this
$ sed 's/arith.addf %acc, %relu : f32/arith.addf %acc, %relu fastmath<reassoc> : f32/' 01_progressive_lowering.mlir | mlir-opt --convert-scf-to-cf --convert-arith-to-llvm --finalize-memref-to-llvm --convert-func-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts | mlir-translate --mlir-to-llvmir -o out/sum_relu_reassoc.ll && grep -o 'fadd reassoc float' out/sum_relu_reassoc.ll
what it printed on the machine that wrote this lab
fadd reassoc float
(a) no triple, no reassoc:
run this
$ opt -O3 out/sum_relu.ll -pass-remarks-missed=loop-vectorize -pass-remarks-analysis=loop-vectorize -o /dev/null 2>&1 | head -1
what it printed on the machine that wrote this lab
remark: <unknown>:0:0: loop not vectorized: cannot prove it is safe to reorder floating-point operations
(b) no triple, WITH reassoc:
run this
$ opt -O3 out/sum_relu_reassoc.ll -pass-remarks-missed=loop-vectorize -pass-remarks-analysis=loop-vectorize -o /dev/null 2>&1 | head -1
what it printed on the machine that wrote this lab
remark: <unknown>:0:0: the cost-model indicates that vectorization is not beneficial
(c) aarch64 triple, no reassoc:
run this
$ opt -mtriple=aarch64-linux-gnu -O3 out/sum_relu.ll -pass-remarks=loop-vectorize -o /dev/null 2>&1
what it printed on the machine that wrote this lab
remark: <unknown>:0:0: vectorized loop (vectorization width: 4, interleaved count: 2)
(d) aarch64 triple, WITH reassoc:
run this
$ opt -mtriple=aarch64-linux-gnu -O3 out/sum_relu_reassoc.ll -pass-remarks=loop-vectorize -o /dev/null 2>&1
what it printed on the machine that wrote this lab
remark: <unknown>:0:0: vectorized loop (vectorization width: 4, interleaved count: 4)

(a) is the classic: summing floats in a different order can change the result, so without permission LLVM won't. (b) is the sneaky one: the permission is there, but with no target triple opt models a machine with NO vector registers and the cost model says "not beneficial". mlir-translate emits no triple unless the module has `llvm.target_triple`, so a pipeline that forgets it silently loses vectorization. Every real compiler sets it (a vendor MLIR backend in its translate step, IREE via --iree-llvmcpu-target-triple). When numbers look scalar, `grep target out.ll` FIRST.

run this
$ opt -mtriple=aarch64-linux-gnu -O3 out/sum_relu.ll -o - | llc -mtriple=aarch64-linux-gnu -O3 -o - | sed -n '/%vector.body/,/b.ne/p'
what it printed on the machine that wrote this lab
.LBB0_5:                                // %vector.body
                                        // =>This Inner Loop Header: Depth=1
	ldp	q3, q4, [x9, #-16]
	subs	x10, x10, #8
	add	x9, x9, #32
	fmax	v3.4s, v3.4s, v2.4s
	fmax	v4.4s, v4.4s, v2.4s
	mov	s5, v3.s[1]
	fadd	s0, s0, s3
	mov	s6, v3.s[2]
	mov	s3, v3.s[3]
	fadd	s0, s0, s5
	mov	s5, v4.s[2]
	fadd	s0, s0, s6
	fadd	s0, s0, s3
	mov	s3, v4.s[1]
	fadd	s0, s0, s4
	fadd	s0, s0, s3
	mov	s3, v4.s[3]
	fadd	s0, s0, s5
	fadd	s0, s0, s3
	b.ne	.LBB0_5

The relu is vectorized (fmax v.4s) but then every lane is pulled out with `mov s5, v3.s[1]` and added ONE AT A TIME into s0: eight dependent scalar fadds per iteration. This is an ORDERED ("strict") reduction — AArch64 lets LLVM vectorize without reassoc by keeping the exact summation order. The dependency chain on s0 is the bottleneck: ~8 x fadd latency per 8 elements.

run this
$ opt -mtriple=aarch64-linux-gnu -O3 out/sum_relu_reassoc.ll -o - | llc -mtriple=aarch64-linux-gnu -O3 -o - | sed -n '/%vector.body/,/b.ne/p'
what it printed on the machine that wrote this lab
.LBB0_7:                                // %vector.body
                                        // =>This Inner Loop Header: Depth=1
	ldp	q6, q7, [x10, #-32]
	subs	x11, x11, #16
	ldp	q16, q17, [x10], #64
	fmax	v6.4s, v6.4s, v2.4s
	fmax	v7.4s, v7.4s, v2.4s
	fmax	v16.4s, v16.4s, v2.4s
	fmax	v17.4s, v17.4s, v2.4s
	fadd	v5.4s, v5.4s, v6.4s
	fadd	v0.4s, v0.4s, v7.4s
	fadd	v3.4s, v3.4s, v16.4s
	fadd	v4.4s, v4.4s, v17.4s
	b.ne	.LBB0_7
16 elements per iteration, four INDEPENDENT vector accumulators
(v5, v0, v3, v4), zero lane extraction. The horizontal sum (faddp) runs
once after the loop. Count it yourself:
run this
$ for f in sum_relu sum_relu_reassoc; do echo "== $f"; opt -mtriple=aarch64-linux-gnu -O3 out/$f.ll -o - | llc -mtriple=aarch64-linux-gnu -O3 -o - | awk '/^sum_relu:/,/Lfunc_end/' | grep -oE '^\s+(fadd|faddp|fmax)\s+[vs]' | sed 's/\s\+/ /g' | sort | uniq -c; done
what it printed on the machine that wrote this lab
== sum_relu
      9  fadd s
      1  fmax s
      2  fmax v
== sum_relu_reassoc
      2  faddp s
      2  faddp v
      1  fadd s
      8  fadd v
      1  fmax s
      5  fmax v

Triton kernels hit this constantly: `tl.sum`, softmax, layernorm and attention are all reductions. Whether the backend marks them reassoc, and whether it vectorizes them ITSELF in MLIR (lab 04_optimization/03), is often a 4-10x difference on a real kernel. You'll measure it there.

take it further

1. Use `fastmath<fast>` instead of `reassoc` and diff the .ll files. Which other flags appear (nnan, ninf, nsz, arcp, contract, afn)? 2. Replace `arith.maximumf` with `arith.maxnumf` (different NaN semantics) and compare the aarch64 instruction chosen (fmax vs fmaxnm). 3. Add `-mattr=+sve` to both opt and llc. Does the loop use z-registers? (Scalable vectorization needs -scalable-vectorization=on or SVE-preferring cpus like -mcpu=neoverse-v1; try both.) ============================================================ relu(x) summed: a reduction loop with a select inside

read the source · mlir-lowering/01_progressive_lowering.mlir
func.func @sum_relu(%x: memref<?xf32>) -> f32 { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %zero = arith.constant 0.0 : f32 %n = memref.dim %x, %c0 : memref<?xf32> %r = scf.for %i = %c0 to %n step %c1 iter_args(%acc = %zero) -> (f32) { %v = memref.load %x[%i] : memref<?xf32> %relu = arith.maximumf %v, %zero : f32 %acc2 = arith.addf %acc, %relu : f32 scf.yield %acc2 : f32 } return %r : f32 }
◇
Finish the lab: write the evidence note

A lab is done when you can say what it showed on your machine — not when you have read it.