AArch64: make use of reg-reg-extend amode.

When a load/store instruction needs an address of the form `v0 +
uextend(v1)` or `v0 + sextend(v1)` (or the commuted forms thereof), we
currently generate a separate zero/sign-extend operation and then use a
plain `[rA, rB]` addressing mode. This patch extends `lower_address()`
to look at both addends of an address if it has two addends and a zero
offset, recognize extension operations, and incorporate them directly
into a `[rA, rB, UXTW]` or `[rA, rB, SXTW]` form. This should improve
our performence on WebAssembly workloads, at least, because we often see
a 64-bit linear memory base indexed by a 32-bit (Wasm) pointer value.

cranelift/codegen/src/isa/aarch64/lower.rs

@@ -550,7 +550,51 @@ pub(crate) fn lower_address<C: LowerCtx<I = Inst>>(
         return MemArg::RegOffset(reg, offset as i64, elem_ty);
     }
 
-    // Handle two regs and a zero offset, if possible.
+    // Handle two regs and a zero offset with built-in extend, if possible.
+    if addends.len() == 2 && offset == 0 {
+        // r1, r2 (to be extended), r2_bits, is_signed
+        let mut parts: Option<(Reg, Reg, usize, bool)> = None;
+        // Handle extension of either first or second addend.
+        for i in 0..2 {
+            if let Some((op, ext_insn)) =
+                maybe_input_insn_multi(ctx, addends[i], &[Opcode::Uextend, Opcode::Sextend])
+            {
+                // Non-extended addend.
+                let r1 = input_to_reg(ctx, addends[1 - i], NarrowValueMode::ZeroExtend64);
+                // Extended addend.
+                let r2 = input_to_reg(
+                    ctx,
+                    InsnInput {
+                        insn: ext_insn,
+                        input: 0,
+                    },
+                    NarrowValueMode::None,
+                );
+                let r2_bits = ty_bits(ctx.input_ty(ext_insn, 0));
+                parts = Some((
+                    r1,
+                    r2,
+                    r2_bits,
+                    /* is_signed = */ op == Opcode::Sextend,
+                ));
+                break;
+            }
+        }
+
+        if let Some((r1, r2, r2_bits, is_signed)) = parts {
+            match (r2_bits, is_signed) {
+                (32, false) => {
+                    return MemArg::RegExtended(r1, r2, ExtendOp::UXTW);
+                }
+                (32, true) => {
+                    return MemArg::RegExtended(r1, r2, ExtendOp::SXTW);
+                }
+                _ => {}
+            }
+        }
+    }
+
+    // Handle two regs and a zero offset in the general case, if possible.
     if addends.len() == 2 && offset == 0 {
         let ra = input_to_reg(ctx, addends[0], NarrowValueMode::ZeroExtend64);
         let rb = input_to_reg(ctx, addends[1], NarrowValueMode::ZeroExtend64);
@@ -812,6 +856,20 @@ pub(crate) fn maybe_input_insn<C: LowerCtx<I = Inst>>(
     None
 }
 
+/// Checks for an instance of any one of `ops` feeding the given input.
+pub(crate) fn maybe_input_insn_multi<C: LowerCtx<I = Inst>>(
+    c: &mut C,
+    input: InsnInput,
+    ops: &[Opcode],
+) -> Option<(Opcode, IRInst)> {
+    for &op in ops {
+        if let Some(inst) = maybe_input_insn(c, input, op) {
+            return Some((op, inst));
+        }
+    }
+    None
+}
+
 /// Checks for an instance of `op` feeding the given input, possibly via a conversion `conv` (e.g.,
 /// Bint or a bitcast).
 ///