commit a69929714a36eaaad30511e1582b14c81e63cab1
parent d9902a7ea9f62a06393ee58db352fe8f635e88de
Author: Ryan Sepassi <rsepassi@gmail.com>
Date: Mon, 15 Jun 2026 23:34:10 -0700
opt: rematerialization instead of spilling at O1 (O1.md W2)
Diffstat:
3 files changed, 281 insertions(+), 13 deletions(-)
diff --git a/mk/test.mk b/mk/test.mk
@@ -859,7 +859,7 @@ test-macho: lib $(TEST_RT_DEP) $(ROUNDTRIP_BIN_MACHO) $(LINK_EXE_RUNNER) $(JIT_R
OPT_TEST_BIN = build/test/cg_ir_lower_test
TINY_INLINE_TEST_BIN = build/test/tiny_inline_test
-test-opt: bin $(OPT_TEST_BIN) test-opt-tiny-inline test-opt-inline test-opt-zero-arg test-opt-static-prune-aa64 test-opt-aa64-tail test-opt-x64-win-tail-sret test-opt-prologue-tier test-opt-whole-program-inline test-opt-lto-phase1 test-opt-redundant-copy-ext test-opt-redundant-frame-sub test-opt-o1-branch-cleanup test-opt-hot-slot-order
+test-opt: bin $(OPT_TEST_BIN) test-opt-tiny-inline test-opt-inline test-opt-zero-arg test-opt-static-prune-aa64 test-opt-aa64-tail test-opt-x64-win-tail-sret test-opt-prologue-tier test-opt-whole-program-inline test-opt-lto-phase1 test-opt-redundant-copy-ext test-opt-redundant-frame-sub test-opt-o1-branch-cleanup test-opt-hot-slot-order test-opt-o1-remat
$(OPT_TEST_BIN)
@@ -887,6 +887,12 @@ test-opt-o1-branch-cleanup: bin
test-opt-hot-slot-order: bin
@KIT=$(abspath $(BIN)) bash test/opt/hot_slot_order.sh
+# Structural disasm check: O1 rematerialization recomputes cheap spilled values
+# (small load_imm, addr_of[local]) at the use instead of reloading from a slot.
+.PHONY: test-opt-o1-remat
+test-opt-o1-remat: bin
+ @KIT=$(abspath $(BIN)) bash test/opt/o1_remat.sh
+
test-opt-tiny-inline: bin $(TINY_INLINE_TEST_BIN)
$(TINY_INLINE_TEST_BIN)
diff --git a/src/opt/pass_lower.c b/src/opt/pass_lower.c
@@ -1398,12 +1398,149 @@ static Reg scratch_for(Func* f, u8 cls, u32* next) {
return r;
}
+/* ---------------------------------------------------------------------------
+ * Rematerialization (doc/plan/O1.md W2).
+ *
+ * A spilled value whose single def is *cheaper to recompute than to reload* is
+ * recomputed at each use instead of reloaded from its slot; the spill store is
+ * dropped entirely (the original def then writes a dead scratch, removed by
+ * mir_dce). The v1 set is restricted to input-less producers whose recompute is
+ * <= reload cost, so the recompute is valid at any use site:
+ * - IR_LOAD_IMM whose immediate materializes in <= 2 instructions
+ * (movz, or movz;movk -- i.e. <= 2 nonzero 16-bit halfwords);
+ * - IR_ADDR_OF of an OPK_LOCAL (one add off the stable frame base).
+ * IR_ADDR_OF(global) (hoisted by addr_of_global_cse) and large IR_LOAD_CONST
+ * (4-insn materialize, worse than a 1-insn reload) are deliberately excluded.
+ *
+ * The classification is a single linear pre-pass (opt_mark_remat) over all
+ * instructions; the choice at each spill site is local. No global analysis.
+ * ------------------------------------------------------------------------- */
+
+/* Cost gate: how many machine instructions an immediate of `type` takes to
+ * materialize, using the movz/movk model (count of nonzero 16-bit halfwords).
+ * This matches aa64's aa_load_imm_words exactly and is conservative for the
+ * other backends (a 2-halfword imm is at most a 2-3 insn rebuild everywhere),
+ * so a constant that costs more than its slot reload never rematerializes. */
+static u32 remat_imm_insn_count(const Func* f, KitCgTypeId type, i64 imm) {
+ u32 width = type_size_fallback(f, type);
+ u32 halfwords = (width > 4u) ? 4u : 2u;
+ u64 v = (u64)imm;
+ u32 n = 0;
+ for (u32 i = 0; i < halfwords; ++i)
+ if (((v >> (i * 16u)) & 0xffffu) != 0u) ++n;
+ return n ? n : 1u; /* zero is one movz */
+}
+
+/* True iff `in` is in the conservative v1 rematerialization set. */
+static int remat_inst_is_candidate(const Func* f, const Inst* in) {
+ switch ((IROp)in->op) {
+ case IR_LOAD_IMM:
+ return remat_imm_insn_count(f, in->type, in->extra.imm) <= 2u;
+ case IR_ADDR_OF:
+ return in->nopnds >= 2 && in->opnds[1].kind == OPK_LOCAL;
+ default:
+ return 0;
+ }
+}
+
+/* Per-function rematerialization table, computed once by opt_mark_remat and
+ * read at every spill site. remat_def[v] is an arena snapshot of v's single
+ * defining inst when v is a remat candidate, else NULL. */
+typedef struct RematInfo {
+ Inst** remat_def; /* indexed by PReg; NULL entry = not a remat candidate */
+ u32 nregs;
+} RematInfo;
+
+typedef struct RematCountCtx {
+ u8* def_count; /* saturating at 2 per PReg */
+} RematCountCtx;
+
+static void remat_count_def(Func* f, Inst* in, Operand* op, int is_def,
+ void* arg) {
+ (void)in;
+ if (!is_def || op->kind != OPK_REG) return;
+ PReg v = (PReg)op->v.reg;
+ if (v == PREG_NONE || v == 0 || v >= opt_reg_count(f)) return;
+ RematCountCtx* c = (RematCountCtx*)arg;
+ if (c->def_count[v] < 2u) ++c->def_count[v];
+}
+
+/* Linear pre-pass: record remat_def[v] iff v has exactly one def and that def
+ * is in the v1 set. One walk to count defs + snapshot candidate defs, one PReg
+ * sweep to drop multiply-defined entries. */
+static void opt_mark_remat(Func* f, RematInfo* ri) {
+ u32 nregs = opt_reg_count(f);
+ memset(ri, 0, sizeof *ri);
+ ri->nregs = nregs;
+ if (nregs == 0) return;
+ ri->remat_def = arena_zarray(f->arena, Inst*, nregs);
+ u8* def_count = arena_zarray(f->arena, u8, nregs);
+ RematCountCtx cc = {def_count};
+ for (u32 b = 0; b < f->nblocks; ++b) {
+ Block* bl = &f->blocks[b];
+ for (u32 i = 0; i < bl->ninsts; ++i) {
+ Inst* in = &bl->insts[i];
+ opt_walk_inst_operands(f, in, remat_count_def, &cc);
+ if (!remat_inst_is_candidate(f, in)) continue;
+ PReg dv = (PReg)in->def;
+ if (dv == PREG_NONE || dv == 0 || dv >= nregs) continue;
+ /* Only the single-def case is sound; the count sweep below nullifies any
+ * PReg that ends up with more than one def. Snapshot the inst (and its
+ * operand array) so later block mutation cannot disturb the recipe. */
+ Inst* snap = arena_znew(f->arena, Inst);
+ *snap = *in;
+ if (in->nopnds) {
+ snap->opnds = arena_array(f->arena, Operand, in->nopnds);
+ memcpy(snap->opnds, in->opnds, sizeof(Operand) * in->nopnds);
+ }
+ ri->remat_def[dv] = snap;
+ }
+ }
+ for (u32 v = 0; v < nregs; ++v)
+ if (def_count[v] != 1u) ri->remat_def[v] = NULL;
+}
+
+static Inst* remat_def_for(const RematInfo* ri, PReg v) {
+ if (!ri || !ri->remat_def || v == PREG_NONE || v == 0 || v >= ri->nregs)
+ return NULL;
+ return ri->remat_def[v];
+}
+
typedef struct RewriteCtx {
RewriteList* before;
RewriteList* after;
+ const RematInfo* remat;
u32 next_scratch[OPT_REG_CLASSES];
} RewriteCtx;
+/* Shared "materialize this spilled PReg at this use" helper. Every spilled-use
+ * rewrite path routes through this. When v rematerializes, clone its recorded
+ * def into `before` with dst = scratch; otherwise append an IR_LOAD from the
+ * slot. The caller has already pointed the operand at `scratch`. */
+static void remat_or_reload_use(Func* f, RewriteCtx* c, Operand* op, PReg v,
+ Reg scratch) {
+ Inst* def = remat_def_for(c->remat, v);
+ if (def) {
+ Inst* re = list_push(f, c->before, (IROp)def->op);
+ InstId id = re->id;
+ *re = *def;
+ re->id = id;
+ re->opnds = arena_array(f->arena, Operand, def->nopnds ? def->nopnds : 1u);
+ if (def->nopnds)
+ memcpy(re->opnds, def->opnds, sizeof(Operand) * def->nopnds);
+ re->opnds[0] = *op; /* dst already rewritten to scratch */
+ re->def = (Val)scratch;
+ re->type = op->type;
+ return;
+ }
+ Inst* ld = list_push(f, c->before, IR_LOAD);
+ ld->opnds = arena_array(f->arena, Operand, 2);
+ ld->opnds[0] = *op;
+ ld->opnds[1] = spill_addr(f, v);
+ ld->nopnds = 2;
+ ld->extra.mem = spill_mem(f, v);
+}
+
static void rewrite_one_operand(Func* f, Inst* owner, Operand* op, int is_def,
void* arg) {
RewriteCtx* c = (RewriteCtx*)arg;
@@ -1426,12 +1563,12 @@ static void rewrite_one_operand(Func* f, Inst* owner, Operand* op, int is_def,
}
op->v.reg = scratch;
if (!is_def) {
- Inst* ld = list_push(f, c->before, IR_LOAD);
- ld->opnds = arena_array(f->arena, Operand, 2);
- ld->opnds[0] = *op;
- ld->opnds[1] = spill_addr(f, v);
- ld->nopnds = 2;
- ld->extra.mem = spill_mem(f, v);
+ remat_or_reload_use(f, c, op, v, scratch);
+ } else if (remat_def_for(c->remat, v)) {
+ /* Every use of this PReg rematerializes (the producer is input-less and all
+ * three spilled-use paths route through remat_or_reload_use), so no frame
+ * use of the slot remains -- drop the spill store. The original def now
+ * writes a dead scratch that mir_dce removes. */
} else {
Inst* st = list_push(f, c->after, IR_STORE);
st->opnds = arena_array(f->arena, Operand, 2);
@@ -1442,16 +1579,31 @@ static void rewrite_one_operand(Func* f, Inst* owner, Operand* op, int is_def,
}
}
-static void rewrite_call_arg_operand(Func* f, Operand* op) {
+/* Spilled call-arg use. Without remat the slot becomes a direct OPK_LOCAL frame
+ * operand; with remat we recompute into a scratch so the dropped spill store
+ * leaves no dangling frame read. */
+static void rewrite_call_arg_operand(Func* f, Operand* op, RewriteCtx* ctx) {
if (!op || op->kind != OPK_REG) return;
PReg v = (PReg)op->v.reg;
if (v == PREG_NONE || v == 0 || v >= opt_reg_count(f)) return;
u8 alloc_kind = opt_preg_alloc_kind(f, v);
if (alloc_kind == OPT_ALLOC_HARD) {
op->v.reg = opt_preg_hard_reg(f, v);
- } else if (alloc_kind == OPT_ALLOC_SPILL) {
- *op = spill_addr(f, v);
+ return;
}
+ if (alloc_kind != OPT_ALLOC_SPILL) return;
+ if (remat_def_for(ctx->remat, v)) {
+ u8 cls = opt_preg_loc_cls(f, v);
+ Reg scratch = scratch_for(f, cls, &ctx->next_scratch[cls]);
+ if (scratch == (Reg)REG_NONE)
+ compiler_panic(f->c, (SrcLoc){0, 0, 0},
+ "opt rewrite: no scratch register for spilled class %u",
+ (unsigned)cls);
+ op->v.reg = scratch;
+ remat_or_reload_use(f, ctx, op, v, scratch);
+ return;
+ }
+ *op = spill_addr(f, v);
}
static void rewrite_store_value_operand(Func* f, Inst* owner, Operand* op,
@@ -1467,6 +1619,12 @@ static void rewrite_store_value_operand(Func* f, Inst* owner, Operand* op,
return;
}
if (alloc_kind == OPT_ALLOC_SPILL) {
+ /* A remat candidate must recompute into a scratch (the slot may no longer
+ * be written); a non-remat spill may fold into a direct frame operand. */
+ if (remat_def_for(ctx->remat, v)) {
+ rewrite_one_operand(f, owner, op, 0, ctx);
+ return;
+ }
*op = spill_addr(f, v);
return;
}
@@ -1488,11 +1646,11 @@ static void rewrite_call_arg_value(Func* f, Inst* owner, CGABIValue* v,
RewriteCtx* ctx) {
if (!v) return;
rewrite_call_arg_indirect_base(f, owner, &v->storage, ctx);
- rewrite_call_arg_operand(f, &v->storage);
+ rewrite_call_arg_operand(f, &v->storage, ctx);
for (u32 i = 0; i < v->nparts; ++i) {
Operand* op = (Operand*)&v->parts[i].op;
rewrite_call_arg_indirect_base(f, owner, op, ctx);
- rewrite_call_arg_operand(f, op);
+ rewrite_call_arg_operand(f, op, ctx);
}
}
@@ -1566,6 +1724,8 @@ static void rewrite_func(Func* f, const OptLiveInfo* live_info) {
u64* live = arena_zarray(f->arena, u64, words ? words : 1u);
u32 ncall_save_pregs = 0;
PReg* call_save_pregs = rewrite_collect_call_save_pregs(f, &ncall_save_pregs);
+ RematInfo remat;
+ opt_mark_remat(f, &remat);
InstRefs refs;
memset(&refs, 0, sizeof refs);
u32 live_active_words = 0;
@@ -1599,6 +1759,7 @@ static void rewrite_func(Func* f, const OptLiveInfo* live_info) {
memset(&ctx, 0, sizeof ctx);
ctx.before = &before;
ctx.after = &after;
+ ctx.remat = &remat;
if ((IROp)in.op == IR_CALL) {
IRCallAux* aux = (IRCallAux*)in.extra.aux;
if (aux) {
@@ -1607,7 +1768,7 @@ static void rewrite_func(Func* f, const OptLiveInfo* live_info) {
for (u32 k = 0; k < aux->plan.nargs; ++k) {
rewrite_call_arg_indirect_base(f, &in, &aux->plan.args[k].src,
&ctx);
- rewrite_call_arg_operand(f, &aux->plan.args[k].src);
+ rewrite_call_arg_operand(f, &aux->plan.args[k].src, &ctx);
}
for (u32 k = 0; k < aux->plan.nrets; ++k)
rewrite_one_operand(f, &in, &aux->plan.rets[k].dst, 1, &ctx);
diff --git a/test/opt/o1_remat.sh b/test/opt/o1_remat.sh
@@ -0,0 +1,101 @@
+#!/usr/bin/env bash
+# Structural check for -O1 rematerialization (doc/plan/O1.md W2).
+#
+# At a use of a spilled value whose single def is cheaper to recompute than to
+# reload, kit recomputes at the use instead of reloading from the slot, and
+# drops the spill store entirely when every use rematerializes. The conservative
+# v1 set is input-less producers whose recompute is <= reload cost:
+# - IR_LOAD_IMM whose immediate is <= 2 instructions (movz / movz;movk);
+# - IR_ADDR_OF of a local (one `add` off the stable frame base).
+#
+# The probe drives a small marker constant (KMARK, one movz) and the address of
+# a local (one `add xN, x29, #off`) to spill: both are low-frequency but live
+# across two trailing calls, while sixteen high-frequency array values win every
+# callee-save register. After W2 each spilled use of the constant / address must
+# RECOMPUTE (a `movz #KMARK` / a repeated `add xN, x29, #off`) rather than reload
+# from a slot, and the spill store must be gone (no `str` of the recomputed
+# value into a frame slot it is never read from).
+#
+# Pinned on aarch64 (the reference backend) where the mnemonics are stable.
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
+KIT="${KIT:-$ROOT/build/kit}"
+WORK="$ROOT/build/test/opt/o1_remat"
+mkdir -p "$WORK"
+
+SRC="$WORK/case.c"
+cat > "$SRC" <<'EOF'
+extern void use_ptr(int *);
+extern void sink(long);
+extern void barrier(void);
+
+/* KMARK = 0x4321 is a single-movz constant; &loc is one add off the frame base.
+ * Both are low-frequency but live across the two trailing calls, so the
+ * sixteen high-frequency array values displace them from the callee-saves and
+ * they spill. W2 should recompute both at their uses instead of reloading. */
+long remat_probe(long *arr) {
+ int loc;
+ const long KMARK = 0x4321;
+ long a=arr[0],b=arr[1],c=arr[2],d=arr[3],e=arr[4],f=arr[5],g=arr[6],h=arr[7];
+ long i=arr[8],j=arr[9],k=arr[10],l=arr[11],m=arr[12],o=arr[13],p=arr[14],q=arr[15];
+ barrier();
+ long s = a*a+b*b+c*c+d*d+e*e+f*f+g*g+h*h+i*i+j*j+k*k+l*l+m*m+o*o+p*p+q*q;
+ s += a*b+b*c+c*d+d*e+e*f+f*g+g*h+h*i+i*j+j*k+k*l+l*m+m*o+o*p+p*q+q*a;
+ s += a+b+c+d+e+f+g+h+i+j+k+l+m+o+p+q;
+ use_ptr(&loc); /* use 1: address of a local as a call arg */
+ sink(KMARK); /* use 2: the constant as a call arg */
+ use_ptr(&loc); /* use 3: the address again */
+ return s + KMARK; /* use 4: the constant again */
+}
+EOF
+
+OBJ="$WORK/case.o"
+"$KIT" cc -target aarch64-linux-gnu -O1 -std=c11 -c "$SRC" \
+ -o "$OBJ" > "$WORK/cc.out" 2>&1
+"$KIT" objdump -d "$OBJ" > "$WORK/dis.out" 2>&1
+
+BODY="$(awk '
+ /<remat_probe>:/ { in_fn = 1; next }
+ in_fn && /^[0-9a-f]+ </ { exit }
+ in_fn { print }
+' "$WORK/dis.out")"
+
+fail() {
+ printf 'o1_remat FAILED: %s\n' "$1" >&2
+ printf ' --- remat_probe disassembly ---\n' >&2
+ printf '%s\n' "$BODY" | sed 's/^/ | /' >&2
+ exit 1
+}
+
+[ -n "$BODY" ] || fail "remat_probe not found in disassembly"
+
+# (1) Constant rematerialization: KMARK (0x4321) is recomputed at each of its
+# two uses. A reload-based baseline would materialize it once and reload via
+# ldr; remat emits the `movz #0x4321` again at every use -> count >= 2.
+movz_n="$(printf '%s\n' "$BODY" | grep -cE 'movz[ \t]+x[0-9]+, 0x4321' || true)"
+[ "$movz_n" -ge 2 ] || fail "constant KMARK not rematerialized (movz 0x4321 x$movz_n, expected >=2)"
+
+# (2) Address rematerialization: the same `add xN, x29, #off` recompute appears
+# at both address uses. Find the most common add-off-frame-base offset and
+# require it to recur (a single def would appear once).
+addr_off="$(printf '%s\n' "$BODY" \
+ | grep -oE 'add x[0-9]+, x29, #[0-9]+$' \
+ | grep -oE '#[0-9]+$' | sort | uniq -c | sort -rn | head -1 | awk '{print $2}')"
+[ -n "$addr_off" ] || fail "no 'add xN, x29, #off' (addr_of local) recompute found"
+addr_n="$(printf '%s\n' "$BODY" | grep -cE "add x[0-9]+, x29, ${addr_off}\$" || true)"
+[ "$addr_n" -ge 2 ] || fail "local address not rematerialized (add x29,${addr_off} x$addr_n, expected >=2)"
+
+# (3) Store-drop: the rematerialized constant and address must leave no spill
+# store behind. With every use rematerialized there must be no `str` of the
+# KMARK value or the recomputed address into a slot. We approximate by
+# requiring that the marker constant is never stored: a kept spill store
+# would be `movz xN,#0x4321 ; str xN,[x29,#..]` -- i.e. a movz immediately
+# followed by a str of the same register. None should remain.
+if printf '%s\n' "$BODY" | grep -A1 -E 'movz[ \t]+x[0-9]+, 0x4321' \
+ | grep -qE 'str[ \t]+x[0-9]+, \[x29'; then
+ fail "spill store for rematerialized constant survived (movz;str)"
+fi
+
+printf 'o1_remat: OK (const movz x%s, addr add-off-x29 x%s, spill stores dropped)\n' \
+ "$movz_n" "$addr_n"
