WRITE-A-PASS LAB — your own MLIR pass in C++, out of tree, tested and measured

TRACK: write_a_pass | DIFFICULTY: 4 | NEEDS: cmake, ninja, clang++, the build-full LLVM FILES: CMakeLists.txt, include/Lab/Passes.{td,h}, lib/DivToMul.cpp, tools/lab-opt.cpp, test/div_to_mul.mlir, bench/ (scale.mlir, bench_scale.c, run.sh), sync_headers.sh

What this teaches

Everything so far used passes other people wrote. On a production MLIR backend you'll write them. This lab builds the smallest complete one: TableGen definition, rewrite pattern, driver tool, FileCheck tests, and a benchmark. It's modelled on a rewrite that production backends really do carry (x / broadcast(c) → x * broadcast(1/c)).

include/Lab/Passes.td   -- the pass as DATA: name, anchor op, options, statistics, docs
          | mlir-tblgen -gen-pass-decls        (CMake: mlir_tablegen + add_public_tablegen_target)
build/include/Lab/Passes.h.inc  -- generated base class LabDivToMulBase<>, options struct, registration
include/Lab/Passes.h    -- GEN_PASS_DECL + GEN_PASS_REGISTRATION: the public header
lib/DivToMul.cpp        -- GEN_PASS_DEF + a RewritePattern + runOnOperation()
tools/lab-opt.cpp       -- MlirOptMain + dialect registry + registerLabPasses()
test/div_to_mul.mlir    -- RUN/CHECK tests (lit + FileCheck)

This is the exact layout of a vendor MLIR backend (include/<project>/Transforms/Passes.td → lib/Transforms/*.cpp → bin/<project>-opt), IREE and triton-shared.

STEP 0 — A build-environment trap you will hit (and did, writing this)

./sync_headers.sh

LLVM build commit : 89d02eebfb7d
source tree HEAD  : 6506e7380c7f
221 headers differ; overlaid llvm/include + mlir/include at 89d02eebfb7d (49M)

The first build linked with:

undefined reference to `llvm::cl::Option::addOccurrence(unsigned int, llvm::StringRef, llvm::StringRef)'

nm -DC libLLVM.so | grep addOccurrence showed the library exports a 4-argument version. The headers had a 3-argument one. Cause: The LLVM source tree was git pulled on Sep 23, but the build directory was built on Sep 16. MLIRConfig.cmake points at the source tree's include/, so we compiled against newer headers than the library we linked. sync_headers.sh extracts the include trees at the build's commit (read from clang --version) with git archive, without touching your checkout, and CMakeLists.txt puts them first. (First attempt: overlaying only the 221 changed files. That failed with 'PDLPatternMatch.h.inc' file not found, because headers include siblings by relative path. Overlay whole trees or nothing.)

Rule for a production MLIR backend: the LLVM you build against and the headers you include must come from the same commit. Pin the LLVM commit in the project (a production MLIR backend, a vendor MLIR backend and triton-shared all do), and never git pull an LLVM tree that has a build other projects link against.

STEP 1 — Build

source ../tools/env.sh
cmake -G Ninja -S . -B build -DMLIR_DIR=$LLVM_BUILD/lib/cmake/mlir -DCMAKE_BUILD_TYPE=Release \
      -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_C_COMPILER=clang
ninja -C build                        # ~15 s; produces build/lab-opt
build/lab-opt --help | grep -A1 lab-div-to-mul

      --lab-div-to-mul          -   Rewrite x / C into x * (1/C) when that is legal
        --allow-inexact         - Rewrite even when 1/C is not exact and the op has no arcp flag

The help text and the option came from Passes.td. You wrote no parsing code.

STEP 2 — Read the pass (lib/DivToMul.cpp, ~100 lines)

The three ideas that make up nearly every pass you'll write:

Match one op, bail out with a reason. matchAndRewrite(arith::DivFOp op, ...) returns notifyMatchFailure(op, "why not") for every case it declines. (With an assertions build, -debug-only=greedy-rewriter prints those reasons.)
Legality is the real work. The rewrite itself is 2 lines. Deciding when it's allowed is the rest:
- APFloat::divide computes 1/C and reports whether rounding happened (opOK = exact). 1/4 = 0.25 is exact, so it's always legal. 1/3 isn't, so x * 0.333… can differ from x / 3 in the last bit.
- Inexact is legal only if the op carries fastmath<arcp> ("allow reciprocal") or the user opts in (allow-inexact=true).
- Never touch 0, inf, NaN or denormal divisors.
Greedy driver. applyPatternsGreedily applies patterns until nothing changes. Most MLIR "optimisation passes" are a pattern set plus this call.

STEP 3 — Tests

llvm-lit -v test

PASS: lab-pass :: div_to_mul.mlir (1 of 1)

test/div_to_mul.mlir has one function per legality case (power of two, inexact refused, inexact allowed via option, arcp kept on the new op, splat vector, never-touch cases), and two RUN lines (default and allow-inexact). Proven to fail: changing the expected 2.500000e-01 to 2.600000e-01 makes it FAIL.

STEP 4 — Measure: what legality costs

bench/scale.mlir: y = x / 3.0 over an L1-resident 4096-float buffer, 20 000 times. bench/bench_scale.c also compares every result bit-for-bit with exact x/3. Run bench/run.sh. (It lives in bench/ because LLVM's CMake helpers reject a stray .c next to a target: llvm_check_source_file_list fails the configure with "Found erroneous configuration for source file bench_scale.c". Another thing learned the hard way.)

div           0.300 ns/element   results differing from exact x/3: 0 of 4096      (32 vdivps)
pass_default  0.299 ns/element   results differing from exact x/3: 0 of 4096      (32 vdivps)
pass_inexact  0.100 ns/element   results differing from exact x/3: 1371 of 4096   (32 vmulps)

The default pass correctly refuses: same code, same bits.
allow-inexact is 3× faster and changes a third of the results (by one ulp). That's the trade -ffast-math makes. Whether it's acceptable is a product decision (usually fine for ML inference, not fine for a reference test).

STEP 5 — Check whether LLVM already does it (it partly does)

The first version of bench/scale.mlir had arith.divf ... fastmath<arcp>. Both variants compiled to vmulps and both differed in 1371 elements: LLVM's instcombine already rewrites fdiv arcp X, C into a multiply by 1/C, and fdiv X, 4.0 into fmul X, 0.25 with no flags at all. So for scalar constants this pass is redundant after lowering. Where an MLIR-level version earns its keep:

the divisor isn't a visible constant by the time LLVM sees it, e.g. the x / linalg.fill(c) case: after bufferization it's a load from a buffer, and LLVM can't prove it holds one value;
you want the decision made before vectorization, tiling or fusion, or applied to a target LLVM handles poorly.

Before writing a pass, always compile the input through the existing pipeline and look at the asm. Half of "missing optimisations" are already there.

What a vendor MLIR backend's version does differently (worth discussing with the team)

Its pass rewrites x / fill(y) → x * fill(1/y) without checking exactness or fastmath flags, i.e. it always takes the inexact trade (sensible for an accelerator where division is very expensive and ML tolerance is loose). For a CPU backend that may run reference comparisons, gating it on arcp/an option, as here, is one reasonable design. It's a policy question to raise, not a bug.

Statistics (and why this build can't show them)

Passes.td declares three statistics (num-exact, num-inexact, num-refused), and --mlir-pass-statistics would print them. On build-full the report is empty: LLVM statistics are compiled out unless LLVM is built with assertions or LLVM_FORCE_ENABLE_STATS=ON. That's the same assertions build lab 06/01 recommends for -debug-only.

Try this yourself

Extend the pattern to x / fill(c) inside a linalg.generic body (the splat-divisor case). You'll need to look through a linalg.fill / block argument to the constant.
Add a second pattern: x * 1.0 → x (exact, always legal). Does canonicalize already do it? (Check first!)
Build the pass as a plugin instead of a new tool: add_mlir_library(... PLUGIN) + mlir-opt --load-pass-plugin=.... (Look up mlir/examples/standalone/standalone-plugin.)