kit

o1_remat.sh (4742B)

---

 #!/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"