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

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

@@ -209,6 +209,27 @@
         (from_bits u8)
         (to_bits u8))
 
+        ;; A bitfield move instruction, either UBFM or SBFM.
+        ;; The BFM instruction is instead encoded as BitfieldMoveMod.
+        (BitfieldMove
+         (size OperandSize)
+         (bfm_op BfmOp)
+         (rd WritableReg)
+         (rn Reg)
+         (immr UImm6)
+         (imms UImm6))
+
+        ;; The bitfield move instruction that modifies its `rd`, i.e. BFM.
+        ;; The input state of `rd` is represented as `ri` (for "input"),
+        ;; and is constrained with `rd` so they use the same register.
+        (BitfieldMoveMod
+         (size OperandSize)
+         (rd WritableReg)
+         (ri Reg)
+         (rn Reg)
+         (immr UImm6)
+         (imms UImm6))
+
        ;; A conditional-select operation.
        (CSel
         (rd WritableReg)
@@ -1204,6 +1225,7 @@
 
 ;; Leaves currently without specs: excluded from verification via todo.
 (attr MInst.Udf (tag TODO))
+(attr MInst.BitfieldMove (tag TODO))
 
 ;; An ALU operation. This can be paired with several instruction formats
 ;; below (see `Inst`) in any combination.
@@ -1263,6 +1285,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 +1312,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 +2987,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,6 +417,15 @@ fn enc_ccmp_imm(size: OperandSize, rn: Reg, imm: UImm5, nzcv: NZCV, cond: Cond)
         | nzcv.bits()
 }
 
+impl BfmOp {
+    fn opc(self) -> u8 {
+        match self {
+            BfmOp::UBfm => 0b10,
+            BfmOp::SBfm => 0b00,
+        }
+    }
+}
+
 fn enc_bfm(opc: u8, size: OperandSize, rd: Writable<Reg>, rn: Reg, immr: u8, imms: u8) -> u32 {
     match size {
         OperandSize::Size64 => {
@@ -2920,13 +2929,36 @@ 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)
                 };
+                let opc = bfm_op.opc();
                 sink.put4(enc_bfm(opc, size, rd, rn, 0, from_bits - 1));
             }
+            &Inst::BitfieldMove {
+                size,
+                bfm_op,
+                rd,
+                rn,
+                immr,
+                imms,
+            } => {
+                let opc = bfm_op.opc();
+                sink.put4(enc_bfm(opc, size, rd, rn, immr.value(), imms.value()));
+            }
+            &Inst::BitfieldMoveMod {
+                size,
+                rd,
+                ri,
+                rn,
+                immr,
+                imms,
+            } => {
+                debug_assert_eq!(rd.to_reg(), ri);
+                sink.put4(enc_bfm(0b01, size, rd, rn, immr.value(), imms.value()));
+            }
             &Inst::Jump { ref dest } => {
                 let off = sink.cur_offset();
                 // Indicate that the jump uses a label, if so, so that a fixup can occur later.

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,17 @@ fn aarch64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_def(rd);
             collector.reg_use(rn);
         }
+        Inst::BitfieldMove { rd, rn, .. } => {
+            // The UBFM and SBFM instructions overwrite all bits in `rd`,
+            // unlike BFM which is represented as `BitfieldMoveMod` instead.
+            collector.reg_def(rd);
+            collector.reg_use(rn);
+        }
+        Inst::BitfieldMoveMod { rd, ri, rn, .. } => {
+            collector.reg_reuse_def(rd, 1); // `rd` == `ri`.
+            collector.reg_use(ri);
+            collector.reg_use(rn);
+        }
         Inst::Args { args } => {
             for ArgPair { vreg, preg } in args {
                 collector.reg_fixed_def(vreg, *preg);
@@ -2594,6 +2615,36 @@ 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::BitfieldMoveMod {
+                size,
+                rd,
+                ri,
+                rn,
+                immr,
+                imms,
+            } => {
+                let rd = pretty_print_ireg(rd.to_reg(), size);
+                let ri = pretty_print_ireg(ri, size);
+                let rn = pretty_print_ireg(rn, size);
+                let immr = immr.pretty_print(0);
+                let imms = imms.pretty_print(0);
+                format!("bfm {rd}, {ri}, {rn}, {immr}, {imms}")
+            }
             &Inst::Call { ref info } => {
                 let try_call = info
                     .try_call_info

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

@@ -1699,6 +1699,23 @@
 (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 (bfm_immr ty a b) (bfm_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 (bfm_immr ty a b) (bfm_imms ty a b)))
+
+;; Helper methods for constructing the correct `immr` and `imms` immediates.
+(spec (bfm_immr ty a b) (provide true))
+(decl bfm_immr (Type u64 u64) UImm6)
+(extern constructor bfm_immr bfm_immr)
+(spec (bfm_imms ty a b) (provide true))
+(decl bfm_imms (Type u64 u64) UImm6)
+(extern constructor bfm_imms bfm_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 bfm_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 bfm_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
-