OPTIMIZATION LAB #03 — "It's fast with the pipeline, slow without — which pass is X?"

TRACK: optimization | DIFFICULTY: 2 | NEEDS: labs 01–02 of this track TOOL: tools/ablate.py (reusable on any kernel)

What this teaches

Labs 01 and 02 ablated by hand. In real work you inherit a pipeline of 50–120 passes (a production backend’s Linalg→LLVM pass adds 126) and a report like "kernel K got 5× slower after the rebase" or "K is fast on one target, slow on our CPU pipeline". You need a mechanical way to find the one pass that matters.

tools/ablate.py does leave-one-out ablation:

FULL: every listed pass.
-P: every pass except P, one row per pass.
NONE: none of them.

Each variant is built, checked for correctness by the harness, timed (pinned, interleaved rounds, min kept), and compared against FULL.

Why leave-one-out instead of add-one-in? Because passes interact. Lab 02 showed fusion alone buys 1.6× over nothing, but removing it from the full pipeline costs 5.5×: eliminate-empty-tensors only pays off once fusion has removed the intermediates. Ablating from the optimised end measures each pass in the context it actually runs in.

STEP 1 — Ablate the fusion pipeline, with some neutral passes mixed in

python3 tools/ablate.py --payload 04_optimization/fusion.mlir \
    --harness 04_optimization/bench_fusion.c \
    --passes "canonicalize,linalg-fuse-elementwise-ops,cse,linalg-fold-unit-extent-dims,eliminate-empty-tensors" \
    --post "one-shot-bufferize{bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map},canonicalize" \
    --rounds 3

Real output:

variant                                     min ms   vs FULL  result
FULL                                          8.05     1.00x  OK
-canonicalize                                 8.24     1.02x  OK
-linalg-fuse-elementwise-ops                 44.49     5.53x  OK
-cse                                          8.30     1.03x  OK
-linalg-fold-unit-extent-dims                 8.30     1.03x  OK
-eliminate-empty-tensors                     37.06     4.60x  OK
NONE                                         46.55     5.78x  OK

How to read it:

Two passes matter, three are noise (≤ 1.03× is within this machine's noise, lab 01 STEP 4). That's typical: on any one kernel, most of a pipeline is neutral. The job is finding the few that aren't.
The two that matter aren't additive: 5.5× + 4.6× doesn't "explain" 5.8×. They depend on each other. Report that as one finding: "fusion AND empty-tensor elimination are both required to avoid the temp buffer".
result is OK everywhere. When a removal makes a row WRONG, you've also found a correctness dependency (a later pass assumes an earlier one ran). That's a bug worth filing by itself.

STEP 2 — After you find X, explain X (never stop at the table)

The table gives you a suspect; the IR gives you the explanation. For each row that moved:

# what does the kernel look like WITHOUT X?  (drop the pass from the pipeline and diff the IR)
mlir-opt 04_optimization/fusion.mlir --pass-pipeline='builtin.module(linalg-fuse-elementwise-ops,one-shot-bufferize{bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map},canonicalize)' > /tmp/without.mlir
mlir-opt 04_optimization/fusion.mlir --pass-pipeline='builtin.module(linalg-fuse-elementwise-ops,eliminate-empty-tensors,one-shot-bufferize{bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map},canonicalize)' > /tmp/with.mlir
diff <(grep -E 'alloc|copy|generic' /tmp/without.mlir) <(grep -E 'alloc|copy|generic' /tmp/with.mlir)

Then back the explanation with a counter (lab 02 STEP 3: page faults) and write it down as: pass → IR change → hardware effect → time.

The same method on the matmul (by hand, lab 01)

run_matmul.sh is the add-one-in version for transform-script stages (T, V, H). To do leave-one-out there too: ./run_matmul.sh one noT "V H", ./run_matmul.sh one noV "T H", ./run_matmul.sh one noH "T V", and compare with T V H. (Try it: vectorizing without tiling is a compile-time disaster at 512³. Some "Xs" are prerequisites, not optimisations.)

When the pipeline is in C++ (a vendor MLIR backend / a production MLIR backend)

You can't pass --passes to a pipeline hardcoded in LinalgToLLVMPass.cpp. Three options, from least to most work:

Use its switches. Real pipelines guard passes with options (if (options.enableX) pm.addPass(...)). Those are your ablation knobs; a vendor MLIR backend sets them from Triton kernel kwargs in its backend options module.
Dump and replay. Run the compiler with --mlir-print-ir-after-all / MLIR_ENABLE_DUMP=1 (Triton), save the IR before the pipeline, and replay the pipeline as a --pass-pipeline string with the tool's -opt binary (the backend’s own *-opt binary). Then this lab applies directly.
Add an option for the pass you suspect. That's also the patch you'll send upstream once you've proved it matters.

Try this yourself

Add func.func(linalg-generalize-named-ops) to the list. Does it matter here? Why not?
Put eliminate-empty-tensors before linalg-fuse-elementwise-ops in FULL. Does order matter? (ablate.py keeps the order you give.)
Point ablate.py at a Triton-derived kernel from track 07 once you've built one.