aarch64: lower ishl + ushr/sshr pairs to a single ubfm/sbfm instruction

cranelift/codegen/src/isa/aarch64/inst.isle

@@ -209,6 +209,17 @@
         (from_bits u8)
         (to_bits u8))
 
+        ;; A bitfield move instruction, either UBFM or SBFM.
+        ;; BFM has been intentionally excluded as it leaves some bits
+        ;; in `rd` unchanged, rather than overwriting all of them.
+        (BitfieldMove
+         (size OperandSize)
+         (bfm_op BfmOp)
+         (rd WritableReg)
+         (rn Reg)
+         (immr UImm6)
+         (imms UImm6))
+
        ;; A conditional-select operation.
        (CSel
         (rd WritableReg)
@@ -1263,6 +1274,14 @@
     (MovN)
 ))
 
+;; A bitfield move operation.
+;; Note that BFM is excluded as it modifies rather than overwrites `rd`.
+(type BfmOp
+  (enum
+    (UBfm)
+    (SBfm)
+))
+
 (model UImm5 (type (bv 5)))
 (type UImm5 (primitive UImm5))
 
@@ -1282,6 +1301,9 @@
 (model ImmShift (type (bv 6)))
 (type ImmShift (primitive ImmShift))
 
+(model UImm6 (type (bv 6)))
+(type UImm6 (primitive UImm6))
+
 (model ShiftOpAndAmt
       (type
             (struct
@@ -2954,6 +2976,14 @@
             (_ Unit (emit (MInst.Extend dst rn signed from_bits to_bits))))
         dst))
 
+;; Helper for emitting `MInst.BitfieldMove` instructions.
+(attr bitfield_move (veri chain))
+(decl bitfield_move (Type BfmOp Reg UImm6 UImm6) Reg)
+(rule (bitfield_move ty bfm_op rn immr imms)
+      (let ((dst WritableReg (temp_writable_reg ty))
+            (_ Unit (emit (MInst.BitfieldMove (operand_size ty) bfm_op dst rn immr imms))))
+        dst))
+
 ;; Helper for emitting `MInst.FpuExtend` instructions.
 (attr fpu_extend (veri chain))
 (decl fpu_extend (Reg ScalarSize) Reg)

cranelift/codegen/src/isa/aarch64/inst/emit.rs

@@ -417,7 +417,14 @@ fn enc_ccmp_imm(size: OperandSize, rn: Reg, imm: UImm5, nzcv: NZCV, cond: Cond)
         | nzcv.bits()
 }
 
-fn enc_bfm(opc: u8, size: OperandSize, rd: Writable<Reg>, rn: Reg, immr: u8, imms: u8) -> u32 {
+fn enc_bfm(
+    bfm_op: BfmOp,
+    size: OperandSize,
+    rd: Writable<Reg>,
+    rn: Reg,
+    immr: u8,
+    imms: u8,
+) -> u32 {
     match size {
         OperandSize::Size64 => {
             debug_assert!(immr <= 63);
@@ -428,11 +435,15 @@ fn enc_bfm(opc: u8, size: OperandSize, rd: Writable<Reg>, rn: Reg, immr: u8, imm
             debug_assert!(imms <= 31);
         }
     }
-    debug_assert_eq!(opc & 0b11, opc);
+    let opc = match bfm_op {
+        BfmOp::UBfm => 0b10,
+        BfmOp::SBfm => 0b00,
+        // Note: BFM (`01`) is intentionally excluded
+    };
     let n_bit = size.sf_bit();
     0b0_00_100110_0_000000_000000_00000_00000
         | size.sf_bit() << 31
-        | u32::from(opc) << 29
+        | opc << 29
         | n_bit << 22
         | u32::from(immr) << 16
         | u32::from(imms) << 10
@@ -2920,12 +2931,22 @@ impl MachInstEmit for Inst {
                 from_bits,
                 to_bits,
             } => {
-                let (opc, size) = if signed {
-                    (0b00, OperandSize::from_bits(to_bits))
+                let (bfm_op, size) = if signed {
+                    (BfmOp::SBfm, OperandSize::from_bits(to_bits))
                 } else {
-                    (0b10, OperandSize::Size32)
+                    (BfmOp::UBfm, OperandSize::Size32)
                 };
-                sink.put4(enc_bfm(opc, size, rd, rn, 0, from_bits - 1));
+                sink.put4(enc_bfm(bfm_op, size, rd, rn, 0, from_bits - 1));
+            }
+            &Inst::BitfieldMove {
+                size,
+                bfm_op,
+                rd,
+                rn,
+                immr,
+                imms,
+            } => {
+                sink.put4(enc_bfm(bfm_op, size, rd, rn, immr.value(), imms.value()));
             }
             &Inst::Jump { ref dest } => {
                 let off = sink.cur_offset();

cranelift/codegen/src/isa/aarch64/inst/imms.rs

@@ -547,11 +547,26 @@ pub struct ImmShift {
 impl ImmShift {
     /// Create an ImmShift from raw bits, if possible.
     pub fn maybe_from_u64(val: u64) -> Option<ImmShift> {
-        if val < 64 {
-            Some(ImmShift { imm: val as u8 })
-        } else {
-            None
-        }
+        (val < 64).then_some(ImmShift { imm: val as u8 })
+    }
+
+    /// Get the immediate value.
+    pub fn value(&self) -> u8 {
+        self.imm
+    }
+}
+
+/// A 6-bit immediate used by the `immr` and `imms` fields of bitfield move instructions.
+#[derive(Copy, Clone, Debug)]
+pub struct UImm6 {
+    /// 6-bit immediate.
+    pub imm: u8,
+}
+
+impl UImm6 {
+    /// Create a UImm6 from raw bits, if possible.
+    pub fn maybe_from_u8(val: u8) -> Option<UImm6> {
+        (val < 64).then_some(UImm6 { imm: val })
     }
 
     /// Get the immediate value.
@@ -915,6 +930,12 @@ impl PrettyPrint for ImmShift {
     }
 }
 
+impl PrettyPrint for UImm6 {
+    fn pretty_print(&self, _: u8) -> String {
+        format!("#{}", self.imm)
+    }
+}
+
 impl PrettyPrint for MoveWideConst {
     fn pretty_print(&self, _: u8) -> String {
         if self.shift == 0 {

cranelift/codegen/src/isa/aarch64/inst/mod.rs

@@ -34,10 +34,10 @@ mod emit_tests;
 // Instructions (top level): definition
 
 pub use crate::isa::aarch64::lower::isle::generated_code::{
-    ALUOp, ALUOp3, AMode, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BitOp, BranchTargetType, FPUOp1,
-    FPUOp2, FPUOp3, FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUModOp,
-    VecALUOp, VecExtendOp, VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp, VecRRNarrowOp,
-    VecRRPairLongOp, VecRRRLongModOp, VecRRRLongOp, VecShiftImmModOp, VecShiftImmOp,
+    ALUOp, ALUOp3, AMode, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BfmOp, BitOp, BranchTargetType,
+    FPUOp1, FPUOp2, FPUOp3, FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp,
+    VecALUModOp, VecALUOp, VecExtendOp, VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp,
+    VecRRNarrowOp, VecRRPairLongOp, VecRRRLongModOp, VecRRRLongOp, VecShiftImmModOp, VecShiftImmOp,
 };
 
 /// A floating-point unit (FPU) operation with two args, a register and an immediate.
@@ -60,6 +60,16 @@ pub enum FPUOpRIMod {
     Sli64(FPULeftShiftImm),
 }
 
+impl BfmOp {
+    /// Get the assembly mnemonic for this opcode.
+    pub fn op_str(&self) -> &'static str {
+        match self {
+            BfmOp::UBfm => "ubfm",
+            BfmOp::SBfm => "sbfm",
+        }
+    }
+}
+
 impl BitOp {
     /// Get the assembly mnemonic for this opcode.
     pub fn op_str(&self) -> &'static str {
@@ -791,6 +801,14 @@ fn aarch64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_def(rd);
             collector.reg_use(rn);
         }
+        Inst::BitfieldMove { rd, rn, .. } => {
+            // BFM has been excluded from this instruction format
+            // as it can leave some bits of `rd` unchanged.
+            // In contrast, the UBFM and SBFM instructions always
+            // replace all bits in `rd`, making it a true def.
+            collector.reg_def(rd);
+            collector.reg_use(rn);
+        }
         Inst::Args { args } => {
             for ArgPair { vreg, preg } in args {
                 collector.reg_fixed_def(vreg, *preg);
@@ -2594,6 +2612,21 @@ impl Inst {
                     format!("{op} {rd}, {rn}")
                 }
             }
+            &Inst::BitfieldMove {
+                size,
+                bfm_op,
+                rd,
+                rn,
+                immr,
+                imms,
+            } => {
+                let op = bfm_op.op_str();
+                let rd = pretty_print_ireg(rd.to_reg(), size);
+                let rn = pretty_print_ireg(rn, size);
+                let immr = immr.pretty_print(0);
+                let imms = imms.pretty_print(0);
+                format!("{op} {rd}, {rn}, {immr}, {imms}")
+            }
             &Inst::Call { ref info } => {
                 let try_call = info
                     .try_call_info

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

@@ -1699,6 +1699,21 @@
 (rule sshr_64 (lower (sshr $I64 x y))
       (do_shift (ALUOp.Asr) $I64 (put_in_reg_sext64 x) y))
 
+;; Specialized lowerings to generate a single `ubfm`/`sbfm` instruction from
+;; an appropriate pair of `ishl` and `ushr`/`sshr` operations.
+(rule sbfm 1 (lower
+  (sshr (ty_32_or_64 ty) (ishl _ x (u64_from_iconst a)) (u64_from_iconst b)))
+  (bitfield_move ty (BfmOp.SBfm) x (sbfm_immr ty a b) (sbfm_imms ty a b)))
+(rule ubfm 1 (lower
+  (ushr (ty_32_or_64 ty) (ishl _ x (u64_from_iconst a)) (u64_from_iconst b)))
+  (bitfield_move ty (BfmOp.UBfm) x (sbfm_immr ty a b) (sbfm_imms ty a b)))
+
+;; Helper methods for constructing the correct `immr` and `imms` immediates.
+(decl sbfm_immr (Type u64 u64) UImm6)
+(extern constructor sbfm_immr sbfm_immr)
+(decl sbfm_imms (Type u64 u64) UImm6)
+(extern constructor sbfm_imms sbfm_imms)
+
 ;; Shift for i128.
 (rule (lower (sshr $I128 x y))
       (lower_sshr128 x (value_regs_get y 0)))

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

@@ -9,7 +9,7 @@ use super::{
     ASIMDFPModImm, ASIMDMovModImm, BranchTarget, CallInfo, Cond, CondBrKind, ExtendOp, FPUOpRI,
     FPUOpRIMod, FloatCC, Imm12, ImmLogic, ImmShift, Inst as MInst, IntCC, MachLabel, MemLabel,
     MoveWideConst, MoveWideOp, NZCV, Opcode, OperandSize, Reg, SImm9, ScalarSize, ShiftOpAndAmt,
-    UImm5, UImm12Scaled, VecMisc2, VectorSize, fp_reg, lower_condcode, stack_reg,
+    UImm5, UImm6, UImm12Scaled, VecMisc2, VectorSize, fp_reg, lower_condcode, stack_reg,
     writable_link_reg, writable_zero_reg, zero_reg,
 };
 use crate::ir::{ArgumentExtension, condcodes};
@@ -241,6 +241,29 @@ impl Context for IsleContext<'_, '_, MInst, AArch64Backend> {
         ImmShift::maybe_from_u64(n.into()).unwrap()
     }
 
+    /// Compute the `immr` value for an `sbfm` instruction,
+    /// derived by fusing an `ishl` by amount `a`, with an `sshr` by amount `b`.
+    fn sbfm_immr(&mut self, ty: Type, a: u64, b: u64) -> UImm6 {
+        let w = ty.lane_bits() as u8;
+        debug_assert!(w <= 64);
+
+        let a = (a as u8) & (w - 1);
+        let b = (b as u8) & (w - 1);
+        let result = if a <= b { b - a } else { w - (a - b) };
+        UImm6::maybe_from_u8(result).expect("result is always less than 64")
+    }
+
+    /// Compute the `imms` value for an `sbfm` instruction,
+    /// derived by fusing an `ishl` by amount `a`, with an `sshr` by amount `b`.
+    fn sbfm_imms(&mut self, ty: Type, a: u64, _b: u64) -> UImm6 {
+        let w = ty.lane_bits() as u8;
+        debug_assert!(w <= 64);
+
+        let a = (a as u8) & (w - 1);
+        let result = w - 1 - (a & (w - 1));
+        UImm6::maybe_from_u8(result).expect("result is always less than 64")
+    }
+
     fn lshr_from_u64(&mut self, ty: Type, n: u64) -> Option<ShiftOpAndAmt> {
         let shiftimm = ShiftOpShiftImm::maybe_from_shift(n)?;
         if let Ok(bits) = u8::try_from(ty_bits(ty)) {

cranelift/filetests/filetests/isa/aarch64/shift-op.clif

@@ -36,4 +36,3 @@ block0(v0: i32):
 ; block0: ; offset 0x0
 ;   lsl w0, w0, #0x15
 ;   ret
-