commit 57e2c48627cc19d3dbce6cb00d2e5af18ad57ca6
parent 856b7e33bbb9fb1cdb96714f784a6836ca82a6ce
Author: Ryan Sepassi <rsepassi@gmail.com>
Date: Tue, 16 Jun 2026 09:19:13 -0700
opt: fold sxtw/uxtw index into load/store addressing mode at O1 (O1-PATTERNS L8)
Add an extend rider to the optimizer's OPK_INDIRECT index operand
(OptOperand.v.ind.index_ext) and a matching NativeAddr.index_ext, plus a
NativeTarget capability hook can_fold_extend_into_addr (true only for aa64).
When a load/store address's index register is produced by a single-use 32->64
widening convert (CV_SEXT -> sxtw, CV_ZEXT -> uxtw), mir_combine's
try_addr_synth folds the extend into the addressing mode and the aa64 backend
emits [Xbase, Wm, sxtw/uxtw #scale] in one instruction instead of
sxtw; add ...,lsl #scale; [reg].
Gating is by target capability: the recognition pass only stamps the rider
when the consuming backend advertises the hook, so x64/rv64 keep the unfolded
(correct) form and never see a rider they cannot emit. Operand-equality in the
W5 load/compute CSE also compares index_ext so a riderless and ridered indirect
never alias.
aa64 done + validated (signed/unsigned/short/char index, scaled store, negative
index sign-extend; multiply-used + value-reused cases keep the sxtw alive).
x64/rv64 gated off (capability hook absent). test-opt + test-toy 1392/0 green.
Diffstat:
5 files changed, 165 insertions(+), 11 deletions(-)
diff --git a/src/arch/aa64/native.c b/src/arch/aa64/native.c
@@ -599,13 +599,26 @@ static u32 aa_str_uimm(u32 size, u32 rt, u32 rn, u32 byte_off) {
/* Register-offset load/store with an explicit load opcode (ld_opc is consulted
* only when load != 0; AA64_LDST_OPC_LDRS_X gives a sign-extending ldrsb/ldrsh
- * into the X register). */
+ * into the X register) and an explicit index-extend `option`
+ * (AA64_LDST_OPTION_*): LSL/UXTX (full X index, the default), or SXTW/UXTW for
+ * a 32-bit W index widened by the addressing mode (O1-PATTERNS L8). */
+static u32 aa_ldst_regoff_opt_v(u32 size, u32 v, u32 load, u32 ld_opc, u32 rt,
+ u32 rn, u32 rm, u32 option, u32 scaled) {
+ return aa64_ldst_regoff_pack((AA64LdStRegOff){
+ .size = size & 3u,
+ .V = v & 1u,
+ .opc = (load ? ld_opc : AA64_LDST_OPC_STR),
+ .Rm = rm & 0x1fu,
+ .option = option & 7u,
+ .S = scaled & 1u,
+ .Rn = rn & 0x1fu,
+ .Rt = rt & 0x1fu});
+}
+
static u32 aa_ldst_regoff_op_v(u32 size, u32 v, u32 load, u32 ld_opc, u32 rt,
u32 rn, u32 rm, u32 scaled) {
- return ((size & 3u) << 30) | 0x38200800u | ((v & 1u) << 26) |
- ((load ? ld_opc : AA64_LDST_OPC_STR) << 22) | ((rm & 0x1fu) << 16) |
- (3u << 13) | ((scaled & 1u) << 12) | ((rn & 0x1fu) << 5) |
- (rt & 0x1fu);
+ return aa_ldst_regoff_opt_v(size, v, load, ld_opc, rt, rn, rm,
+ AA64_LDST_OPTION_LSL, scaled);
}
static u32 aa_ldst_regoff_v(u32 size, u32 v, u32 load, u32 rt, u32 rn, u32 rm,
@@ -1008,8 +1021,16 @@ static void aa_emit_mem(AANativeTarget* a, int load, NativeLoc reg,
if (addr.index_kind != NATIVE_ADDR_INDEX_NONE) {
u32 use_base = base;
u32 scaled = 0;
+ /* L8 index extend: a 32-bit W index widened by the addressing mode. The
+ * Rm field still names the same register number; only the option bits
+ * change (010=UXTW, 110=SXTW). NONE keeps the full-width LSL/UXTX form. */
+ u32 option = AA64_LDST_OPTION_LSL;
if (addr.index_kind != NATIVE_ADDR_INDEX_REG)
aa_panic(a, "unsupported address index");
+ if (addr.index_ext == NATIVE_ADDR_IDX_EXT_SXTW)
+ option = AA64_LDST_OPTION_SXTW;
+ else if (addr.index_ext == NATIVE_ADDR_IDX_EXT_UXTW)
+ option = AA64_LDST_OPTION_UXTW;
if (off) {
use_base = AA_TMP1;
aa_emit_add_imm(a, use_base, base, off);
@@ -1021,8 +1042,9 @@ static void aa_emit_mem(AANativeTarget* a, int load, NativeLoc reg,
} else {
aa_panic(a, "unsupported memory address scale");
}
- aa_emit32(mc, aa_ldst_regoff_op_v(sz, native_loc_is_fp(reg), load, ld_opc,
- rt, use_base, addr.index.reg, scaled));
+ aa_emit32(mc, aa_ldst_regoff_opt_v(sz, native_loc_is_fp(reg), load, ld_opc,
+ rt, use_base, addr.index.reg, option,
+ scaled));
return;
}
if (off >= 0 && (((u32)off & ((1u << sz) - 1u)) == 0) &&
@@ -1058,11 +1080,27 @@ static int aa_addr_legal(NativeTarget* t, const NativeAddr* addr,
if (!addr) return 0;
if (addr->index_kind == NATIVE_ADDR_INDEX_NONE) return 1;
if (addr->index_kind != NATIVE_ADDR_INDEX_REG) return 0;
+ /* The SXTW/UXTW index-extend (L8) is encoded in the same regoff form as the
+ * plain LSL/UXTX index, so the scale legality is identical: scale 0 (no
+ * shift) or scale == access size (S bit). */
if (addr->log2_scale == 0) return 1;
sz = size_idx(mem.size ? mem.size : 8u);
return addr->log2_scale == sz;
}
+/* O1-PATTERNS rider capabilities (see native_target.h). aa64 emits both folded
+ * forms: the shifted-register ALU op (L7, aa_binop) and the SXTW/UXTW
+ * index-extend addressing mode (L8, aa_emit_mem). */
+static int aa_can_fold_shift_into_alu(NativeTarget* t) {
+ (void)t;
+ return 1;
+}
+
+static int aa_can_fold_extend_into_addr(NativeTarget* t) {
+ (void)t;
+ return 1;
+}
+
/* True if `mul Rd, Rn, #c` can be replaced by a single non-mul aarch64
* instruction using only Rn as a source (no extra scratch reg). Constants
* that match: 0, 1, -1, +/-2^k, 2^k+1, 1-2^k for k in [1..width-1]. The
@@ -1168,8 +1206,22 @@ static void aa_load_addr_from_base(AANativeTarget* a, u32 rd, u32 base, i64 off,
/* Read the index before writing rd. O1 can legally select the same physical
* register for an address result and its dead-after-use index; materializing
* the base into rd first would turn `base + index` into `base + base`. */
- aa_emit32(a->base.mc,
- aa_add_lsl(rd, base, addr->index.reg, addr->log2_scale));
+ u32 idx_reg = addr->index.reg;
+ if (addr->index_ext != NATIVE_ADDR_IDX_EXT_NONE) {
+ /* L8 extend rider on an address-materialization (LEA-shaped) use: the
+ * shifted-register add takes a full-width X index, so first widen the
+ * 32-bit W index into a scratch (sxtw/uxtw) before the scaled add. The
+ * regoff memory form (aa_emit_mem) folds the extend directly; this path is
+ * the fallback when the indexed indirect is materialized as an address. */
+ u32 tmp = aa_tmp_avoiding(rd == base ? base : rd);
+ if (tmp == base) tmp = aa_tmp_avoiding(base);
+ if (addr->index_ext == NATIVE_ADDR_IDX_EXT_SXTW)
+ aa_emit32(a->base.mc, aa_sbfm(1, tmp, idx_reg, 0, 31));
+ else
+ aa_emit32(a->base.mc, aa_ubfm(1, tmp, idx_reg, 0, 31));
+ idx_reg = tmp;
+ }
+ aa_emit32(a->base.mc, aa_add_lsl(rd, base, idx_reg, addr->log2_scale));
if (off) aa_emit_add_i64(a, rd, rd, off);
}
@@ -4631,6 +4683,8 @@ NativeTarget* aa64_native_target_new(Compiler* c, ObjBuilder* obj,
t->class_for_type = aa_class_for_type;
t->imm_legal = aa_imm_legal;
t->addr_legal = aa_addr_legal;
+ t->can_fold_shift_into_alu = aa_can_fold_shift_into_alu;
+ t->can_fold_extend_into_addr = aa_can_fold_extend_into_addr;
t->machine_op_clobbers = aa_machine_op_clobbers;
t->func_begin = aa_func_begin;
t->func_begin_known_frame = aa_func_begin_known_frame;
diff --git a/src/arch/native_target.h b/src/arch/native_target.h
@@ -254,13 +254,25 @@ typedef enum NativeImmUse {
NATIVE_IMM_ADDR_OFFSET,
} NativeImmUse;
+/* Addressing-mode index extend (O1-PATTERNS L8). A backend that advertises
+ * can_fold_extend_into_addr may receive a NATIVE_ADDR_INDEX_REG whose register
+ * names a 32-bit value the addressing mode widens to 64 bits with this extend
+ * (aa64 `[Xbase, Wm, sxtw/uxtw #log2_scale]`). NONE = full-width index, the
+ * default; every other backend only ever sees NONE. */
+typedef enum NativeAddrIndexExt {
+ NATIVE_ADDR_IDX_EXT_NONE = 0,
+ NATIVE_ADDR_IDX_EXT_SXTW = 1,
+ NATIVE_ADDR_IDX_EXT_UXTW = 2,
+} NativeAddrIndexExt;
+
typedef struct NativeAddr {
u8 base_kind; /* NativeAddrBaseKind */
u8 cls; /* NativeAllocClass for base value */
u8 index_kind; /* NativeAddrIndexKind */
u8 index_cls; /* NativeAllocClass for index value */
u8 log2_scale;
- u8 pad[3];
+ u8 index_ext; /* NativeAddrIndexExt */
+ u8 pad[2];
KitCgTypeId base_type;
KitCgTypeId index_type;
union {
@@ -448,6 +460,23 @@ struct NativeTarget {
NativeAllocClass (*class_for_type)(NativeTarget*, KitCgTypeId);
int (*imm_legal)(NativeTarget*, NativeImmUse, u32 op, KitCgTypeId, i64);
int (*addr_legal)(NativeTarget*, const NativeAddr*, MemAccess);
+ /* Optional O1-PATTERNS rider capabilities. The target-agnostic recognition
+ * passes (pass_combine / pass_addr_fold) only stamp an operand rider after
+ * the consuming backend advertises the matching capability here, so a
+ * backend that cannot emit the folded form never receives a rider it would
+ * have to silently drop. NULL == not supported (the default for every
+ * backend that has not opted in).
+ *
+ * can_fold_shift_into_alu (L7): the backend emits a single-use left shift
+ * folded into a consuming IADD/ISUB/AND/ORR/EOR as the shifted-register form
+ * (aa64 `add xD,xA,xS,lsl #k`), honoring OptOperand.shift on a binop's
+ * second source operand.
+ *
+ * can_fold_extend_into_addr (L8): the backend emits a single-use sxtw/uxtw
+ * of a load/store index folded into the addressing mode (aa64 `[Xb, Wm,
+ * sxtw/uxtw #scale]`), honoring NativeAddr.index_ext. */
+ int (*can_fold_shift_into_alu)(NativeTarget*);
+ int (*can_fold_extend_into_addr)(NativeTarget*);
/* Optional. Report the physical registers the target's encoding of `op`
* clobbers as a side effect (not its declared operands/results), one bitmask
* per NativeAllocClass. The optimizer keeps values live ACROSS the
diff --git a/src/opt/ir.h b/src/opt/ir.h
@@ -66,10 +66,27 @@ typedef enum OptOperandKind {
} OptOperandKind;
#define OPK_REG OPT_OPK_REG
+/* Operand riders (O1-PATTERNS L7/L8). Default 0 = no rider, so every existing
+ * operand and every backend that does not implement the fold keeps behaving as
+ * before. The recognition passes only stamp a rider after the consuming
+ * NativeTarget advertises the matching capability hook
+ * (can_fold_shift_into_alu / can_fold_extend_into_addr), so a backend that
+ * cannot emit the folded form never sees a rider it would have to drop. */
+typedef enum OptOpIndexExt {
+ OPT_IDX_EXT_NONE = 0, /* index used at full width (X register) */
+ OPT_IDX_EXT_SXTW = 1, /* index is a 32-bit value sign-extended to 64 (W,sxtw) */
+ OPT_IDX_EXT_UXTW = 2, /* index is a 32-bit value zero-extended to 64 (W,uxtw) */
+} OptOpIndexExt;
+
typedef struct OptOperand {
u8 kind;
u8 cls;
- u8 pad[2];
+ /* L7 shift rider: when nonzero and kind == OPK_REG, the register value is
+ * pre-shifted left by `shift` (1..4) before the consuming ALU op uses it,
+ * emitted as the shifted-register form (`add xD,xA,xS,lsl #shift`). 0 = the
+ * register is used as-is. Only stamped on a binop's second source operand. */
+ u8 shift;
+ u8 pad[1];
KitCgTypeId type;
union {
i64 imm;
@@ -84,6 +101,10 @@ typedef struct OptOperand {
Reg base;
Reg index;
u8 log2_scale;
+ /* L8 extend rider (OptOpIndexExt): when SXTW/UXTW, `index` names a 32-bit
+ * value that the addressing mode widens to 64 bits with the recorded
+ * extend (`[Xbase, Wm, sxtw #log2_scale]`). NONE = full-width X index. */
+ u8 index_ext;
i32 ofs;
} ind;
} v;
diff --git a/src/opt/pass_combine.c b/src/opt/pass_combine.c
@@ -1186,6 +1186,50 @@ static int try_addr_synth_one_op(CombineCtx* ctx, Inst* in, i32 i,
}
}
+ /* (L8) index producer is a single-use 32->64 widening convert (sxtw/uxtw):
+ * fold the extend into the addressing mode (`[Xbase, Wm, sxtw/uxtw #scale]`),
+ * one instruction instead of `sxtw xT,wS; add ...,lsl #scale; [reg]`. Runs
+ * after the ISHL fold above, so an `sxtw; lsl; [reg]` chain first collapses
+ * the shift into log2_scale and then this absorbs the sxtw producing that
+ * index. Gated behind the target capability: only a backend that emits the
+ * extended-register addressing form is handed an index_ext rider. */
+ if (op->v.ind.index != (Reg)REG_NONE &&
+ op->v.ind.index_ext == OPT_IDX_EXT_NONE && ctx->target &&
+ ctx->target->can_fold_extend_into_addr &&
+ ctx->target->can_fold_extend_into_addr(ctx->target)) {
+ Reg idx = op->v.ind.index;
+ i32 prod_idx = ctx_producer_of(ctx, RC_INT, idx);
+ if (prod_idx >= 0 && prod_idx < i) {
+ Inst* prod = &ctx->bl->insts[prod_idx];
+ u32 sb = 0, db = 0;
+ int sign_p = 0;
+ if (ext_params(prod, &sb, &db, &sign_p) && sb == 4u && db == 8u &&
+ prod->opnds[0].kind == OPK_REG && prod->opnds[0].cls == RC_INT &&
+ prod->opnds[0].v.reg == idx && prod->opnds[1].kind == OPK_REG &&
+ prod->opnds[1].cls == RC_INT) {
+ Operand prod_def = prod->opnds[0];
+ int killed = 0;
+ int uses_after = count_uses_in_live_range(ctx->f, ctx->bl, prod_idx,
+ &prod_def, &killed);
+ /* `<= 2` mirrors the ISHL fold: tolerate the degenerate
+ * [base == index] aliasing where the convert dst appears twice. The
+ * convert source must be unchanged since the convert (so the W index
+ * names the same 32-bit value), and must not alias the convert dst. */
+ if (uses_after >= 1 && uses_after <= 2 &&
+ (killed || !opt_block_live_out_has_phys_reg(
+ ctx->f, ctx->hard_live, ctx->bl->id, &prod_def)) &&
+ !producer_def_aliases_source(&prod_def, &prod->opnds[1]) &&
+ !ctx_def_changed_since(ctx, RC_INT, prod->opnds[1].v.reg,
+ prod_idx)) {
+ op->v.ind.index = prod->opnds[1].v.reg;
+ op->v.ind.index_ext =
+ sign_p ? (u8)OPT_IDX_EXT_SXTW : (u8)OPT_IDX_EXT_UXTW;
+ any = 1;
+ }
+ }
+ }
+ }
+
return any;
}
@@ -1698,6 +1742,7 @@ static int same_load_addr_operand(const Operand* a, const Operand* b) {
return a->v.ind.base == b->v.ind.base &&
a->v.ind.index == b->v.ind.index &&
a->v.ind.log2_scale == b->v.ind.log2_scale &&
+ a->v.ind.index_ext == b->v.ind.index_ext &&
a->v.ind.ofs == b->v.ind.ofs;
default:
return 0;
diff --git a/src/opt/pass_native_emit.c b/src/opt/pass_native_emit.c
@@ -312,6 +312,11 @@ static NativeAddr addr_from_operand(NativeEmitCtx* e, const OptOperand* op,
addr.index_cls = NATIVE_REG_INT;
addr.index.reg = op->v.ind.index;
addr.log2_scale = op->v.ind.log2_scale;
+ /* L8 extend rider: OPT_IDX_EXT_* and NATIVE_ADDR_IDX_EXT_* share the
+ * 0=NONE/1=SXTW/2=UXTW encoding. Only set by the addr-fold recognition
+ * when the target advertised can_fold_extend_into_addr, so a backend
+ * without the capability always sees NONE. */
+ addr.index_ext = op->v.ind.index_ext;
addr.offset = op->v.ind.ofs;
return addr;
case OPT_OPK_REG:
