42. Lowering the Loops Last to First: Testing the Idea Chapter 39 Left Open¶
What you will understand: how a hypothesis that comes out of a profile gets tested. Chapter 39 found why MLIR's SCFToControlFlow is quadratic in the number of loops in one function (each loop splits its block and moves everything behind it) and said, in its limits section, that lowering the loops from the last to the first "would make each split move a short tail; that is an idea suggested by the mechanism, not something tested here". This chapter tests it. It writes a small pass, --mg-scf-to-cf-reverse, that lowers the top-level statements of a function one at a time, last first, and measures it: against MLIR's own pass on a synthetic function (312 seconds against 1.2 seconds at 64,000 loops), on the real 200-step training program (the lowering pass goes from 2,106 seconds to 3 seconds with no outlining at all), and on its output, which turns out to be byte-for-byte the same IR. It also measures what the pass does not fix.
What you need to know first: Chapter 38 (the quadratic, outlining) and Chapter 39 (the profile and the mechanism: Block::splitBlock, transferNodesFromList). The new pass is about 40 lines of C++ on top of MLIR's own conversion patterns; you do not need to know C++ to follow the measurements.
Compile and run
cd docs/part42/code
./build.sh # mg-opt of this chapter: Chapter 38's plus --mg-scf-to-cf-reverse (a few minutes; needs mlir-18-tools, libmlir-18-dev, cmake)
./scaling.py > scaling_out.txt # synthetic loops, three lowerings, 2,000 to 64,000 loops (about 7 minutes)
./nested_limit.py > nested_limit_out.txt # what the pass does not fix (about 1 minute)
./check_reverse.py > check_reverse_out.txt # the checks of this page (about a minute; needs valgrind)
./mutation.py > mutation_out.txt # 6 WRONG passes plus one wrong setting that the checker must catch (about 10 minutes: each rebuilds mg-opt)
./full200.sh > full200_out.txt # the 200-step training program, stage by stage (about 4 minutes)
./pass_times_200.sh > pass_times_200_out.txt # every pass of the lowering stage on that program; the ordinary pass takes about 35 minutes: run it in the background
cd ../../part15/code && ./run_lit.sh # the whole test suite
mgc now takes --reverse-loops (it substitutes the new pass for --convert-scf-to-cf); nothing else about it changes, and the option is off by default (Chapter 43 returns to what the defaults should be).
Why this page shows a mutation run, and why 'caught' is the wanted result
mutation.py writes broken versions of the pass and expects the checker to fail on each. "Caught" is good; "NOT CAUGHT" would be a gap. Nothing on this page is a failing test.
The idea, and why it should work¶
Chapter 39's profile says: lowering the loop at position i of a block splits the block there and moves every operation behind it into a new block, and the loops are lowered first to last, so loop i still has n − i unlowered nests behind it. About n²/2 nests are moved in all.
Now lower them in the other order. The last loop is split first: nothing is behind it. The next-to-last loop is split second, and what is behind it is the already lowered last loop, which is no longer a pile of operations in this block: the split left a branch behind and the loop's blocks live elsewhere. So every split moves only the operations between two neighbouring loops, a constant number, and the total is linear.
MLIR's pass cannot be told to do this (the conversion driver visits operations in program order), so the new pass drives MLIR's own lowering patterns itself:
// --mg-scf-to-cf-reverse: lower scf.for / scf.if / scf.while to control flow like MLIR's --convert-scf-to-cf, but one top-level statement at a time, LAST TO FIRST.
//
// Why (Chapter 39): lowering a loop splits the block at the loop and moves every operation after it into a new block, so n loops in one block move about n^2/2 operations in all
// when they are lowered first to last. Lowered last to first, the tail after loop k has already been turned into a few branches, so each split moves only the few operations between
// two neighbouring loops. Chapter 42 measures whether that is enough. The result is the same control flow, in a different order of creation.
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/DialectConversion.h"
using namespace mlir;
namespace {
struct ScfToCfReversePass : public PassWrapper<ScfToCfReversePass, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScfToCfReversePass)
StringRef getArgument() const final { return "mg-scf-to-cf-reverse"; }
StringRef getDescription() const final { return "Lower scf to cf one top-level statement at a time, from the last to the first"; }
void getDependentDialects(DialectRegistry ®istry) const override { registry.insert<cf::ControlFlowDialect, arith::ArithDialect>(); }
void runOnOperation() override {
func::FuncOp func = getOperation();
if (func.getBody().empty()) return;
MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx);
populateSCFToControlFlowConversionPatterns(patterns);
FrozenRewritePatternSet frozen(std::move(patterns));
ConversionTarget target(*ctx);
target.addIllegalOp<scf::ForOp, scf::IfOp, scf::WhileOp, scf::ExecuteRegionOp, scf::IndexSwitchOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
// Only the top-level statements of the entry block are handled one by one; they hold everything nested inside them. Collect them first, then go backwards.
SmallVector<Operation *> tops;
for (Operation &op : func.getBody().front())
if (isa<scf::ForOp, scf::IfOp, scf::WhileOp, scf::ExecuteRegionOp, scf::IndexSwitchOp>(&op)) tops.push_back(&op);
for (auto it = tops.rbegin(); it != tops.rend(); ++it) {
Operation *one[] = {*it};
if (failed(applyPartialConversion(ArrayRef<Operation *>(one), target, frozen))) { signalPassFailure(); return; }
}
}
};
} // namespace
namespace mg {
void registerScfToCfReversePass() { PassRegistration<ScfToCfReversePass>(); }
}
It takes the statements of the entry block that are structured control flow (scf.for, scf.if, scf.while, scf.execute_region, scf.index_switch), then, from the last to the first, runs MLIR's own populateSCFToControlFlowConversionPatterns on that one statement (applyPartialConversion accepts a list of operations to convert; the statement's nested operations are converted with it). The patterns are MLIR's, so the lowering of each loop is by construction the one MLIR would have made; only the order differs. That is a claim, and the checks below test it.
mg-opt did not register the scf dialect (nothing ever parsed scf text, since the loops arrive as affine.for and --lower-affine creates the scf ones). This chapter adds it to mg-opt.cpp, so that the edge-case program below can be written directly in scf; no other behaviour changes.
Does it fix the quadratic?¶
The synthetic function of Chapter 38 (n tiny loops, each copying four numbers, all in one function), lowered by MLIR's pass, by the new one, and, as the control that was always fast, with the loops spread over functions of 500. Times are the pass's own, from --mlir-timing; one run each:
time of the lowering pass in seconds, by number of loop nests n (each copies 4 numbers); one run each
n convert-scf-to-cf last-to-first 500 per function growth for 2x the loops: standard, last-to-first
2000 0.067 0.018 0.021
4000 0.514 0.044 0.042 7.6x, 2.4x
8000 4.029 0.083 0.085 7.8x, 1.9x
16000 19.202 0.185 0.159 4.8x, 2.2x
32000 78.359 0.457 0.311 4.1x, 2.5x
64000 311.681 1.201 0.797 4.0x, 2.6x
At 64,000 loops, MLIR's pass takes 311.7 seconds; lowering last to first takes 1.2 seconds, 260 times less, and is within 1.5× of the spread-over-functions control (0.8 seconds). The standard pass grows by about 4× per doubling at the top of the table (4.0× at the last row: quadratic), the new pass by 2.4× to 2.6×.
That 2.4× to 2.6× is a little more than the 2× of a linear algorithm, and the control column, which is linear by construction (the same functions of 500, only more of them), grows by 2.5× to 2.6× over the same last rows too. I did not find out why: memory effects at these sizes are a candidate that I did not test. What the data support is that the new pass behaves like the control, not that it is exactly linear.
The count of instructions¶
Timings are noisy; instruction counts under valgrind --tool=callgrind are exact. The checker (below) counts the instructions spent in ilist_traits<Operation>::transferNodesFromList, the function Chapter 39 identified:
ok convert-scf-to-cf: transferNodesFromList grows x3.97 and x3.98 per doubling (4 = quadratic)
ok last to first: transferNodesFromList grows x2.00 and x2.00 per doubling (2 = linear)
ok last to first: block splitting is 0.0% of the instructions at 2000 loops (convert-scf-to-cf: 5%)
(Lines copied from the checker's output further down.) The mechanism is the one Chapter 39 named, and it is gone: ×2.00 per doubling is exactly linear.
What it costs when the quadratic is not there¶
At 2,000 loops the quadratic part is only 5% of the instructions, and the new pass runs one conversion per statement, which has a fixed price: 706 million instructions in all against 697 million for MLIR's pass, about 1% more. So by instructions the new pass is slightly more expensive for small functions (the checker allows up to 5% more). By time it is already ahead at 2,000 loops (0.018 seconds against 0.067 in the table above), because the instructions of the quadratic part are the expensive kind (Chapter 39: pointer-chasing that misses the cache). Both numbers are tiny at that size; the gap becomes large from 4,000 loops up.
The real program¶
Chapter 36's 200-step training program has 190,636 loop nests in one function (main). Chapter 38 got its compile time from about 40 minutes to 94 seconds by outlining the nests into 75 shared functions. Here is the same program without outlining, lowered last to first, with the time of each stage (full200.sh; the program is the 200-step one compiled from source):
Chapter 36's recorded output (addresses and blank lines removed): md5 cab50027c661
flags: --reverse-loops
front end 0.7 s | mg->affine(+outline) 4.7 s | lower (lower-affine, scf->cf, LLVM dialect) 62.0 s | mlir-translate 38.4 s | clang -O0 22.8 s | run 3.6 s | total compile 128.6 s
output: 76 lines, md5 of the output with addresses and blank lines removed: cab50027c661
-> output IDENTICAL to Chapter 36's recorded output (compiled without outlining and with the ordinary lowering)
flags: --reverse-loops --outline
front end 0.7 s | mg->affine(+outline) 6.1 s | lower (lower-affine, scf->cf, LLVM dialect) 14.0 s | mlir-translate 13.3 s | clang -O0 48.0 s | run 3.3 s | total compile 82.2 s
output: 76 lines, md5 of the output with addresses and blank lines removed: cab50027c661
-> output IDENTICAL to Chapter 36's recorded output (compiled without outlining and with the ordinary lowering)
Without outlining, the whole compile takes 129 seconds, and the output is identical to the output Chapter 36 recorded when it compiled the same program with the ordinary lowering and no outlining (the checker's md5 of the 76 lines, with the addresses and blank lines removed, is the same: cab50027c661). With both the new pass and outlining it takes 82 seconds.
Where did the time go? MLIR's own --mlir-timing report for the lowering stage, for the new pass and for the ordinary pass (pass_times_200.sh):
################ --mg-scf-to-cf-reverse
Total Execution Time: 62.8096 seconds
3.4562 ( 5.2%) 3.4562 ( 5.5%) Parser
3.0524 ( 4.6%) 3.0524 ( 4.9%) ConvertAffineToStandard
3.0916 ( 4.7%) 3.0784 ( 4.9%) 'func.func' Pipeline
3.0916 ( 4.7%) 3.0784 ( 4.9%) {anonymous}::ScfToCfReversePass
6.5084 ( 9.9%) 6.5084 ( 10.4%) {anonymous}::MgAssertToStderrPass
2.4105 ( 3.7%) 2.4105 ( 3.8%) ConvertMathToLLVMPass
5.8842 ( 8.9%) 5.8842 ( 9.4%) ArithToLLVMConversionPass
8.6890 ( 13.2%) 8.6890 ( 13.8%) FinalizeMemRefToLLVMConversionPass
0.2466 ( 0.4%) 0.2466 ( 0.4%) (A) DataLayoutAnalysis
9.0285 ( 13.7%) 9.0285 ( 14.4%) ConvertFuncToLLVMPass
0.4463 ( 0.7%) 0.4463 ( 0.7%) (A) DataLayoutAnalysis
7.8416 ( 11.9%) 7.8416 ( 12.5%) ReconcileUnrealizedCasts
10.9995 ( 16.7%) 10.9995 ( 17.5%) Output
1.6416 ( 2.5%) 1.6416 ( 2.6%) Rest
65.9013 (100.0%) 62.8096 (100.0%) Total
################ --convert-scf-to-cf
Total Execution Time: 2173.0708 seconds
3.2893 ( 0.2%) Parser
3.2489 ( 0.1%) ConvertAffineToStandard
2106.2219 ( 96.9%) SCFToControlFlow
8.2455 ( 0.4%) {anonymous}::MgAssertToStderrPass
3.3522 ( 0.2%) ConvertMathToLLVMPass
6.2020 ( 0.3%) ArithToLLVMConversionPass
9.2687 ( 0.4%) FinalizeMemRefToLLVMConversionPass
0.2598 ( 0.0%) (A) DataLayoutAnalysis
10.1083 ( 0.5%) ConvertFuncToLLVMPass
0.4519 ( 0.0%) (A) DataLayoutAnalysis
8.8744 ( 0.4%) ReconcileUnrealizedCasts
12.3533 ( 0.6%) Output
1.9064 ( 0.1%) Rest
2173.0708 (100.0%) Total
- Ordinary
SCFToControlFlow: 2,106 seconds of the stage's 2,173 (97%). The new pass: 3.1 seconds of 62.8. Everything else in the stage is the same in both: a list of passes taking 3 to 11 seconds each, none of which dominates. - The stage at 62 seconds is then 48% of the compile; the rest is
mlir-translate(38 s) andclang -O0(23 s), also linear in the size of the program. The quadratic really was in one pass. - The 2,106-second figure is one run, made while other work (builds and checks) was running on the same four-core machine; treat it as "about 35 minutes", which is consistent with Chapter 36's "about 40 minutes" for compile and run together, and not as a measurement to the second.
One more observation, not explained: with outlining added, clang -O0 took 48 seconds, against 23 without. So outlining shortens the MLIR stages but lengthens the last one on this machine (at -O0). The total is still better (82 s against 129 s). I did not look into why clang is slower on the outlined file.
Is it the same program?¶
The new pass only helps if it changes nothing else. The checker tests this four ways.
1. The book's examples. On every Mountain Goat example of the book that the front end accepts and that has under 400 lines, --lower-affine followed by each lowering gives output that is byte-for-byte identical: 107 programs, none differ (13 are skipped: 400 lines or more, or deliberately rejected by the front end).
2. Control flow that the examples do not have. The front end only ever produces plain loops, so the examples never contain an scf.if or an scf.while, a loop with a result or a loop inside a loop with iter_args. examples/edge_cases.mlir is hand-written scf with all of them, each later statement using the results of earlier ones:
// Chapter 42: structured control flow that is NOT a plain run of loops, to see that lowering the top-level statements last to first still gives the right program:
// a loop with a result (iter_args) used by a later loop, an scf.if with results, an scf.while, loops nested inside loops, and a loop whose bounds come from an earlier result.
module {
func.func private @printF64(f64)
func.func private @printNewline()
func.func @main() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c5 = arith.constant 5 : index
%c10 = arith.constant 10 : index
%z = arith.constant 0.0 : f64
%one = arith.constant 1.0 : f64
// 1: a loop with a result: the sum 0+1+2+3+4 = 10
%s = scf.for %i = %c0 to %c5 step %c1 iter_args(%acc = %z) -> (f64) {
%fi = arith.index_cast %i : index to i64
%f = arith.sitofp %fi : i64 to f64
%n = arith.addf %acc, %f : f64
scf.yield %n : f64
}
call @printF64(%s) : (f64) -> ()
call @printNewline() : () -> ()
// 2: an scf.if with a result that uses the result of loop 1: 10 > 5, so 1000
%lim = arith.constant 5.0 : f64
%cond = arith.cmpf ogt, %s, %lim : f64
%pick = scf.if %cond -> (f64) {
%big = arith.constant 1000.0 : f64
scf.yield %big : f64
} else {
%small = arith.constant -1.0 : f64
scf.yield %small : f64
}
call @printF64(%pick) : (f64) -> ()
call @printNewline() : () -> ()
// 3: an scf.while that doubles until the value passes 100, starting from 1: 128
%w = scf.while (%v = %one) : (f64) -> f64 {
%hundred = arith.constant 100.0 : f64
%go = arith.cmpf olt, %v, %hundred : f64
scf.condition(%go) %v : f64
} do {
^bb0(%u: f64):
%two = arith.constant 2.0 : f64
%d = arith.mulf %u, %two : f64
scf.yield %d : f64
}
call @printF64(%w) : (f64) -> ()
call @printNewline() : () -> ()
// 4: nested loops: 5 x 5 times adding 1, starting from the result of 3: 128 + 25 = 153
%nest = scf.for %i = %c0 to %c5 step %c1 iter_args(%a0 = %w) -> (f64) {
%inner = scf.for %j = %c0 to %c5 step %c1 iter_args(%a1 = %a0) -> (f64) {
%n = arith.addf %a1, %one : f64
scf.yield %n : f64
}
scf.yield %inner : f64
}
call @printF64(%nest) : (f64) -> ()
call @printNewline() : () -> ()
// 5: a loop whose upper bound is computed (from 4 steps of loop 1's count): 10 iterations adding the loop-2 pick, 153 + 10 * 1000 = 10153
%last = scf.for %i = %c0 to %c10 step %c1 iter_args(%a2 = %nest) -> (f64) {
%n = arith.addf %a2, %pick : f64
scf.yield %n : f64
}
call @printF64(%last) : (f64) -> ()
call @printNewline() : () -> ()
return
}
}
Both lowerings give the same text for it, and the lowered program, run, prints 10, 1000, 128, 153, 10153, the values worked out in the comments by hand.
3. Real programs run the same. An attention example, the tiny transformer and the bigram training run print identical output with and without --reverse-loops.
4. A big function. One function of 3,000 loops lowers to byte-identical IR (1.6 megabytes) from both.
Tests¶
One new lit file test/reverse42/check.mlir (the suite is now 154 tests). The full output of check_reverse.py:
#!/usr/bin/env python3
"""Checks the Chapter 42 claims about --mg-scf-to-cf-reverse (lowering scf to cf last to first):
1. SAME PROGRAM: on every Mountain Goat example of the book that the front end accepts and that has under 400 lines, the output of --lower-affine --mg-scf-to-cf-reverse is byte-for-byte the output of
--lower-affine --convert-scf-to-cf (so nothing downstream can differ);
2. NOT ONLY PLAIN LOOPS: a hand-written function with a loop that has a result, an scf.if with results, an scf.while, nested loops, and a loop that uses earlier results lowers to the same text
and RUNS to the values worked out by hand (10, 1000, 128, 153, 10153);
3. SAME OUTPUT WHEN RUN: three real programs (an attention example, the tiny transformer, the bigram training run) print the same thing with and without --reverse-loops;
4. THE COST IS LINEAR (and one function of 3000 loops lowers to the same text), MEASURED IN INSTRUCTIONS (callgrind, exact and repeatable): in one function of n loops, the cost of moving operations between blocks (ilist_traits::transferNodesFromList)
grows by 4x per doubling under --convert-scf-to-cf and by at most 2.3x under --mg-scf-to-cf-reverse, where it is under 5% of the program's instructions at 2000 loops.
Environment variables (CHK_*) change an input; mutation.py sets them to check that each claim can fail. Usage: check_reverse.py Exit status 0 only if all pass."""
import glob, os, re, subprocess, sys, tempfile
here = os.path.dirname(os.path.abspath(__file__)); root = os.path.join(here, "..", ".."); E = os.environ.get
OPT = E("MG_OPT") or os.path.join(here, "build", "mg-opt"); W = tempfile.mkdtemp(prefix="ch42c_")
REV = E("CHK_REVERSE", "--mg-scf-to-cf-reverse"); STD = "--convert-scf-to-cf"
failures = 0
def report(ok, name, detail=""):
global failures; failures += not ok
print(f" {'ok ' if ok else 'FAIL'} {name}" + ("" if ok else f"\n {detail}"))
def sh(*a, **k): return subprocess.run(list(a), capture_output=True, text=True, **k)
# ---- 1
print("-- 1. same program on the book's examples")
same = diff = skipped = 0; bad = []
for f in sorted(glob.glob(root + "/part*/code/examples/*.mg") + glob.glob(root + "/tour/code/*.mg")):
if sum(1 for _ in open(f)) >= 400: skipped += 1; continue
m = sh(sys.executable, os.path.join(here, "mgfront.py"), f)
if m.returncode: skipped += 1; continue
open(W + "/x.mlir", "w").write(m.stdout); a = sh(OPT, W + "/x.mlir", "--convert-mg-to-affine")
if a.returncode: skipped += 1; continue
open(W + "/y.mlir", "w").write(a.stdout); r = [sh(OPT, W + "/y.mlir", "--lower-affine", p) for p in (STD, REV)]
if r[0].returncode == 0 and r[1].returncode == 0 and r[0].stdout == r[1].stdout: same += 1
else: diff += 1; bad.append(os.path.relpath(f, root))
report(same >= 100 and diff == 0, f"{same} programs lower to byte-identical IR with --mg-scf-to-cf-reverse and --convert-scf-to-cf, {diff} differ ({skipped} skipped: 400 lines or more, or rejected by the front end)", ", ".join(bad[:5]))
# ---- 2
print("-- 2. structured control flow that is not a plain run of loops")
edge = os.path.join(here, "examples", "edge_cases.mlir"); outs = []
for p in (STD, REV):
r = sh(OPT, edge, p, "--convert-arith-to-llvm", "--convert-func-to-llvm", "--reconcile-unrealized-casts", "-o", f"{W}/e{p}.mlir"); outs.append(r.returncode)
report(outs == [0, 0] and open(f"{W}/e{STD}.mlir").read() == open(f"{W}/e{REV}.mlir").read(), "loop with result, scf.if, scf.while, nested loops: the same text from both lowerings", f"exit codes {outs}")
sh("mlir-translate-18", "--mlir-to-llvmir", f"{W}/e{REV}.mlir", "-o", W + "/e.ll")
sh("clang-18", W + "/e.ll", "-o", W + "/e.exe", "-L/usr/lib/llvm-18/lib", "-lmlir_c_runner_utils", "-lm", "-Wl,-rpath,/usr/lib/llvm-18/lib")
got = sh(W + "/e.exe").stdout.split()
report(got == ["10", "1000", "128", "153", "10153"], "the lowered program prints 10, 1000, 128, 153, 10153 (worked out by hand)", f"got {got}")
# ---- 3
print("-- 3. real programs print the same with and without --reverse-loops")
MGC = os.path.join(here, "mgc"); env = dict(os.environ, MG_OPT=OPT); progs = [os.path.join(root, p) for p in ("part29/code/examples/08_self_attention_block.mg", "part30/code/examples/03_tiny_transformer.mg", "part31/code/examples/04_training.mg")]
progs = [p for p in progs if os.path.exists(p)]
res = []
for p in progs:
a = sh(MGC, "run", p, env=env); b = sh(MGC, "run", p, "--reverse-loops", env=env)
norm = lambda s: re.sub(r"base@ = 0x[0-9a-f]+", "base@ = 0x", s)
res.append((os.path.basename(p), a.returncode == 0 and b.returncode == 0 and norm(a.stdout) == norm(b.stdout) and len(a.stdout) > 20))
report(len(res) == 3 and all(ok for _, ok in res), "three programs print identical output: " + ", ".join(n for n, _ in res), str(res))
# ---- 4
print("-- 4. cost in instructions (callgrind): the splitting cost is quadratic under convert-scf-to-cf, linear under the reverse pass")
def gen(n):
out = ["module {", " func.func @f(%a: memref<4xf64>, %b: memref<4xf64>) {"]
for i in range(n): out.append(f" affine.for %i{i} = 0 to 4 {{\n %x{i} = affine.load %a[%i{i}] : memref<4xf64>\n affine.store %x{i}, %b[%i{i}] : memref<4xf64>\n }}")
return "\n".join(out + [" return", " }", "}"]) + "\n"
def profile(n, lowering):
src = f"{W}/g.mlir"; open(src, "w").write(gen(n)); cg = f"{W}/cg.out"
subprocess.run(["valgrind", "--tool=callgrind", f"--callgrind-out-file={cg}", OPT, src, "--lower-affine", lowering, "-o", "/dev/null"], capture_output=True, check=True)
text = sh("callgrind_annotate", cg).stdout; tot = tr = 0
for l in text.split("\n"):
m = re.match(r"\s*([\d,]+) \(\s*[\d.]+%\)\s+(.*)", l)
if not m: continue
ir = int(m.group(1).replace(",", ""))
if "PROGRAM TOTALS" in m.group(2): tot = ir
elif "ilist_traits<mlir::Operation>::transferNodesFromList" in m.group(2): tr += ir
return tot, tr
open(W + "/big.mlir", "w").write(gen(3000)); big = [sh(OPT, W + "/big.mlir", "--lower-affine", p).stdout for p in (STD, REV)]
report(len(big[0]) > 10**6 and big[0] == big[1], "one function of 3000 loops: byte-identical IR from both lowerings (%d bytes)" % len(big[0]))
sizes = (500, 1000, 2000)
std = {n: profile(n, STD) for n in sizes}; rev = {n: profile(n, REV) for n in sizes}
gs = [std[b][1] / max(std[a][1], 1) for a, b in ((500, 1000), (1000, 2000))]; gr = [rev[b][1] / max(rev[a][1], 1) for a, b in ((500, 1000), (1000, 2000))]
report(all(3.9 <= g <= 4.1 for g in gs), "convert-scf-to-cf: transferNodesFromList grows %s per doubling (4 = quadratic)" % " and ".join("x%.2f" % g for g in gs))
report(all(g <= 2.3 for g in gr), "last to first: transferNodesFromList grows %s per doubling (2 = linear)" % " and ".join("x%.2f" % g for g in gr), "")
share = rev[2000][1] / rev[2000][0]
report(share < 0.05, "last to first: block splitting is %.1f%% of the instructions at 2000 loops (convert-scf-to-cf: %.0f%%)" % (100 * share, 100 * std[2000][1] / std[2000][0]))
report(rev[2000][0] <= 1.05 * std[2000][0], "last to first costs at most 5%% more in all at 2000 loops (one conversion per statement): %d million instructions against %d million" % (rev[2000][0] // 10**6, std[2000][0] // 10**6))
print("all checks pass" if not failures else f"{failures} CHECK(S) FAILED"); sys.exit(1 if failures else 0)
(The listing is the checker itself.) Its output:
-- 1. same program on the book's examples
ok 107 programs lower to byte-identical IR with --mg-scf-to-cf-reverse and --convert-scf-to-cf, 0 differ (13 skipped: 400 lines or more, or rejected by the front end)
-- 2. structured control flow that is not a plain run of loops
ok loop with result, scf.if, scf.while, nested loops: the same text from both lowerings
ok the lowered program prints 10, 1000, 128, 153, 10153 (worked out by hand)
-- 3. real programs print the same with and without --reverse-loops
ok three programs print identical output: 08_self_attention_block.mg, 03_tiny_transformer.mg, 04_training.mg
-- 4. cost in instructions (callgrind): the splitting cost is quadratic under convert-scf-to-cf, linear under the reverse pass
ok one function of 3000 loops: byte-identical IR from both lowerings (1634222 bytes)
ok convert-scf-to-cf: transferNodesFromList grows x3.97 and x3.98 per doubling (4 = quadratic)
ok last to first: transferNodesFromList grows x2.00 and x2.00 per doubling (2 = linear)
ok last to first: block splitting is 0.0% of the instructions at 2000 loops (convert-scf-to-cf: 5%)
ok last to first costs at most 5% more in all at 2000 loops (one conversion per statement): 706 million instructions against 697 million
all checks pass
CI now builds this chapter's mg-opt (the newest build, which contains everything before it); run_lit.sh, ci.sh, the workflow and its checker, and Chapter 37's scripts that look for the newest build were updated to find it.
Are the checks good enough? Seven wrong passes.¶
NOTE: this script deliberately breaks the pass. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
the pass is replaced by MLIR's own --convert-scf-to-cf (CHK_REVERSE): caught by 1 failing check(s):
last to first: transferNodesFromList grows x3.97 and x3.98 per doubling (2 = linear)
the statements are lowered first to last, not last to first: caught by 2 failing check(s):
last to first: transferNodesFromList grows x3.97 and x3.98 per doubling (2 = linear)
last to first costs at most 5% more in all at 2000 loops (one conversion per statement): 739 million instructions against 699 million
the whole function is converted in one call (what --convert-scf-to-cf does): caught by 1 failing check(s):
last to first: transferNodesFromList grows x3.97 and x3.98 per doubling (2 = linear)
a top-level scf.if is never lowered: caught by 1 failing check(s):
loop with result, scf.if, scf.while, nested loops: the same text from both lowerings
a top-level scf.while is never lowered: caught by 1 failing check(s):
loop with result, scf.if, scf.while, nested loops: the same text from both lowerings
only the first 2500 top-level statements are lowered: caught by 1 failing check(s):
one function of 3000 loops: byte-identical IR from both lowerings (1634222 bytes)
nested loops are never lowered (only statements whose body has no loop: a loop with a loop inside is skipped): caught by 2 failing check(s):
2 programs lower to byte-identical IR with --mg-scf-to-cf-reverse and --convert-scf-to-cf, 105 differ (13 skipped: 400 lines or more, or rej
loop with result, scf.if, scf.while, nested loops: the same text from both lowerings
restored: ScfToCfReverse.cpp is back to its original and mg-opt is rebuilt
All seven are caught. Look at which check catches each:
- The two that make the order wrong (first to last, or the whole function in one call) are caught only by the instruction count: the output is still the right program, so no output comparison could notice them. A pass that gave the right answer but quietly lost the speed-up would be invisible without that check.
- A pass that skips top-level
scf.iforscf.whileis caught only by the hand-written file, because the book's examples contain neither. If I had trusted the 107 programs alone, this would have gone unnoticed. - Lowering only the first 2,500 statements is caught only by the 3,000-loop comparison (the instruction counts stop at 2,000 loops and look at no output of a big function, so I expected this mutant to slip through and added the 3,000-loop comparison before the run; I did not run the mutants without it).
What the pass does not fix¶
time of the lowering pass in seconds: ONE outer loop whose body contains n loops
n convert-scf-to-cf last-to-first
2000 0.066 0.066
4000 0.484 0.650
8000 5.186 5.052
16000 23.892 24.444
The pass reorders only the top-level statements of a function. One outer loop whose body holds n inner loops is lowered as one statement, and inside it MLIR's driver still goes in program order, so the quadratic is back: 23.9 seconds against 24.4 seconds at 16,000 inner loops, no better at all (the small differences at 2,000 to 8,000 are noise from one run each). The training programs do not look like this (each matrix operation is its own top-level nest), which is why the pass works on them, but a program that put all its operations inside one big loop body would not benefit. Lowering the nested loops last to first too is not a matter of recursing: a lowered inner loop leaves its parent's body with several blocks, and an scf.for must have a single-block body, so the order of conversion has to change more than this pass does. I did not do that.
Limits and what is not established¶
- One machine, one run per cell in the timing tables (the instruction counts are exact; the times are not). The 2,106-second figure was measured with other work running.
- The new pass is not in MLIR. It drives MLIR's patterns from outside and handles only the entry block's top-level statements. A real fix would go in the conversion driver or in
SCFToControlFlowitself, and I have not read that code to see whether such a change would be accepted or would break something else. - Byte-identical output was shown on 107 programs, one synthetic function and one hand-written one, not for every possible
scfprogram. For any case where MLIR's patterns depend on the order of conversion the outputs could differ; none was found. - The nested case is not fixed (above), and functions that are not
func.func(agpu.module, for example) are not looked at by this pass. - Still unexplained: why the new pass and the control both grow by about 2.5× per doubling at the top of the table, and why
clang -O0is slower on outlined input. - The default is unchanged. Whether
mgcshould lower last to first by default, and whether outlining should be on by default, is the next chapter's question.
Reproducing¶
cd docs/part42/code && ./build.sh
./scaling.py > scaling_out.txt && ./nested_limit.py > nested_limit_out.txt
./check_reverse.py > check_reverse_out.txt && ./mutation.py > mutation_out.txt
./full200.sh > full200_out.txt # and ./pass_times_200.sh in the background (about 35 minutes)
cd ../../part15/code && ./run_lit.sh # 154 tests
Chapter summary¶
- Chapter 39's open idea is confirmed: lowering
scfto control flow from the last statement to the first removes the quadratic. On 64,000 synthetic loops the pass goes from 311.7 to 1.2 seconds; on the 200-step training program (190,636 nests in one function), from 2,106 to 3.1 seconds, with no outlining, and the whole compile takes 129 seconds with identical output. - The instruction count of block splitting goes from ×4 per doubling to ×2.00: the mechanism Chapter 39 found is the thing removed.
- The output is byte-identical to MLIR's on 107 book programs, a 3,000-loop function and a hand-written function with
if,whileand nested loops. - The pass costs about 1% more at small sizes and helps only with many top-level statements; it does not help a function whose loops are nested inside one outer loop.
- Seven broken passes are caught, and which check catches each is part of the result.
Self-check questions¶
Each answer is collapsed; try the question first.
-
Why does lowering the loops from last to first avoid the quadratic?
Answer
Lowering a loop splits its block there and moves everything behind it. First to last, loop i has all the later unlowered nests behind it, so about n²/2 are moved in all. Last to first, what is behind a loop is already lowered (it left a branch and its blocks live elsewhere), so each split moves only the few operations between two loops.
-
The two wrong orders are caught only by the instruction count, not by any output check. Why?
Answer
MLIR's patterns give the same lowering of each loop whatever the order; only how long it takes differs. The output of a pass with the wrong order is still byte-identical. Only a measurement of cost can see that the speed-up has gone.
-
Why does the hand-written
edge_cases.mlirmatter when 107 book programs already agree?Answer
The front end only produces plain loops, so none of the book's programs has a top-level
scf.iforscf.while, or loops with results used later. The mutants that skip those statements are caught only by the hand-written file. -
The new pass is about 1% more expensive at 2,000 loops. Why, and what does that say about when to use it?
Answer
It runs one conversion per top-level statement, and each conversion has a fixed price, while the quadratic part it removes is only 5% of the instructions at that size. It pays off for functions with many top-level loops (the gap opens from 4,000 loops) and costs a little for small ones, so it should be chosen by size or accepted as a small tax.
-
What does the nested-loops experiment show, and what would be needed to fix it?
Answer
One outer loop with n inner loops takes the same time under both passes (about 24 seconds at 16,000), because the pass reorders only top-level statements. Fixing it means lowering inner loops last to first as well, but a lowered inner loop leaves the outer loop's body with several blocks, which an
scf.forcannot have, so the order of conversion has to change more than this pass does. -
The 2,106-second figure should be read with care. Why?
Answer
It is one run, made while other work was running on the same four cores. It supports "about 35 minutes", consistent with the roughly 40 minutes Chapter 36 measured for compile and run together, not a figure to the second.
-
Name two things this chapter does not establish.
Answer
Any of: why the new pass and the control both grow by about 2.5× per doubling at the top of the table; why
clang -O0is slower on outlined input; that the outputs match for every possiblescfprogram; whether the same change would be accepted into MLIR; any benefit for loops nested in one body.