aarch64: lower AtomicCAS128 to caspal when LSE is supported

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

@@ -313,6 +313,10 @@
          (ty Type)
          (flags MemFlagsData))
 
+       ;; A 128-bit atomic compare-and-swap operation, such as `caspal`. This
+       ;; instruction requires LSE support (FEAT_LSE) and has acquire-release semantics.
+       (AtomicCAS128 (args BoxAtomicCAS128Args))
+
        ;; Read `access_ty` bits from address `rt`, either 8, 16, 32 or 64-bits, and put
        ;; it in `rn`, optionally zero-extending to fill a word or double word result.
        ;; This instruction is sequentially consistent.
@@ -1341,6 +1345,7 @@
 (type BoxReturnCallIndInfo (primitive BoxReturnCallIndInfo))
 (type BranchTarget (primitive BranchTarget))
 (type BoxJTSequenceInfo (primitive BoxJTSequenceInfo))
+(type BoxAtomicCAS128Args (primitive BoxAtomicCAS128Args))
 (type CodeOffset (primitive CodeOffset))
 (type VecMachLabel extern (enum))
 
@@ -4441,6 +4446,26 @@
           )
           dst))
 
+;; Helper for emitting `MInst.AtomicCAS128` instructions.
+(decl atomic_cas_128_args (WritableReg WritableReg Reg Reg Reg Reg Reg MemFlagsData)
+      BoxAtomicCAS128Args)
+(extern constructor atomic_cas_128_args atomic_cas_128_args)
+
+(decl lse_atomic_cas_128 (Reg ValueRegs ValueRegs MemFlagsData) ValueRegs)
+(rule (lse_atomic_cas_128 addr expect replace flags)
+      (let (
+            (dst_lo WritableReg (temp_writable_reg $I64))
+            (dst_hi WritableReg (temp_writable_reg $I64))
+            (args BoxAtomicCAS128Args (atomic_cas_128_args dst_lo dst_hi
+                                                           (value_regs_get expect 0)
+                                                           (value_regs_get expect 1)
+                                                           (value_regs_get replace 0)
+                                                           (value_regs_get replace 1)
+                                                           addr flags))
+            (_ Unit (emit (MInst.AtomicCAS128 args)))
+          )
+          (value_regs dst_lo dst_hi)))
+
 ;; Helper for emitting `MInst.AtomicRMWLoop` instructions.
 ;; - Make sure that both args are in virtual regs, since in effect
 ;; we have to do a parallel copy to get them safely to the AtomicRMW input

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

@@ -713,6 +713,25 @@ impl APIKey {
     }
 }
 
+/// "Package" of the arguments for the instruction `AtomicCAS128` to avoid
+/// making the `Inst` enum massive.
+///
+/// `rd_lo`/`rd_hi` are really `rs`/`rs+1` in the encoded instruction (so
+/// `rd_lo == rs_lo` and `rd_hi == rs_hi`); they are separated here to have
+/// separate use and def vregs for regalloc.
+#[derive(Clone, Debug)]
+#[expect(missing_docs, reason = "self-describing fields")]
+pub struct AtomicCAS128Args {
+    pub rd_lo: Writable<Reg>,
+    pub rd_hi: Writable<Reg>,
+    pub rs_lo: Reg,
+    pub rs_hi: Reg,
+    pub rt_lo: Reg,
+    pub rt_hi: Reg,
+    pub rn: Reg,
+    pub flags: MemFlagsData,
+}
+
 pub use crate::isa::aarch64::lower::isle::generated_code::TestBitAndBranchKind;
 
 impl TestBitAndBranchKind {

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

@@ -646,6 +646,16 @@ fn enc_cas(size: u32, rs: Writable<Reg>, rt: Reg, rn: Reg) -> u32 {
         | machreg_to_gpr(rt)
 }
 
+fn enc_casp(rs: Writable<Reg>, rt: Reg, rn: Reg) -> u32 {
+    debug_assert_eq!(machreg_to_gpr(rs.to_reg()) & 1, 0);
+    debug_assert_eq!(machreg_to_gpr(rt) & 1, 0);
+
+    0b0_1_0010000_1_1_00000_1_11111_00000_00000
+        | machreg_to_gpr(rs.to_reg()) << 16
+        | machreg_to_gpr(rn) << 5
+        | machreg_to_gpr(rt)
+}
+
 fn enc_asimd_mod_imm(rd: Writable<Reg>, q_op: u32, cmode: u32, imm: u8) -> u32 {
     let abc = (imm >> 5) as u32;
     let defgh = (imm & 0b11111) as u32;
@@ -1666,6 +1676,30 @@ impl MachInstEmit for Inst {
 
                 sink.put4(enc_cas(size, rd, rt, rn));
             }
+            Inst::AtomicCAS128 { args } => {
+                let &AtomicCAS128Args {
+                    rd_lo,
+                    rd_hi,
+                    rs_lo,
+                    rs_hi,
+                    rt_lo,
+                    rt_hi,
+                    rn,
+                    flags,
+                } = &**args;
+                debug_assert_eq!(rd_lo.to_reg(), rs_lo);
+                debug_assert_eq!(rd_hi.to_reg(), rs_hi);
+
+                // These should be pinned to pairs that `casp` requires.
+                debug_assert_eq!(rs_hi, xreg(machreg_to_gpr(rs_lo) as u8 + 1));
+                debug_assert_eq!(rt_hi, xreg(machreg_to_gpr(rt_lo) as u8 + 1));
+
+                if let Some(trap_code) = flags.trap_code() {
+                    sink.add_trap(trap_code);
+                }
+
+                sink.put4(enc_casp(rd_lo, rt_lo, rn));
+            }
             &Inst::AtomicCASLoop { ty, flags, .. } => {
                 /* Emit this:
                     again:

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

@@ -488,6 +488,27 @@ fn aarch64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_use(rt);
             collector.reg_use(rn);
         }
+        Inst::AtomicCAS128 { args } => {
+            let AtomicCAS128Args {
+                rd_lo,
+                rd_hi,
+                rs_lo,
+                rs_hi,
+                rt_lo,
+                rt_hi,
+                rn,
+                flags: _,
+            } = &mut **args;
+            // `casp` requires two consecutive even-aligned register pairs,
+            // which regalloc2 cannot express, so pin everything down.
+            collector.reg_fixed_use(rs_lo, xreg(24));
+            collector.reg_fixed_use(rs_hi, xreg(25));
+            collector.reg_fixed_def(rd_lo, xreg(24));
+            collector.reg_fixed_def(rd_hi, xreg(25));
+            collector.reg_fixed_use(rt_lo, xreg(26));
+            collector.reg_fixed_use(rt_hi, xreg(27));
+            collector.reg_fixed_use(rn, xreg(28));
+        }
         Inst::AtomicCASLoop {
             addr,
             expected,
@@ -1634,6 +1655,28 @@ impl Inst {
 
                 format!("{op} {rd}, {rs}, {rt}, [{rn}]")
             }
+            Inst::AtomicCAS128 { args } => {
+                let &AtomicCAS128Args {
+                    rd_lo,
+                    rd_hi,
+                    rs_lo,
+                    rs_hi,
+                    rt_lo,
+                    rt_hi,
+                    rn,
+                    flags: _,
+                } = &**args;
+                let size = OperandSize::Size64;
+                let rd_lo = pretty_print_ireg(rd_lo.to_reg(), size);
+                let rd_hi = pretty_print_ireg(rd_hi.to_reg(), size);
+                let rs_lo = pretty_print_ireg(rs_lo, size);
+                let rs_hi = pretty_print_ireg(rs_hi, size);
+                let rt_lo = pretty_print_ireg(rt_lo, size);
+                let rt_hi = pretty_print_ireg(rt_hi, size);
+                let rn = pretty_print_ireg(rn, size);
+
+                format!("caspal {rd_lo}, {rd_hi}, {rs_lo}, {rs_hi}, {rt_lo}, {rt_hi}, [{rn}]")
+            }
             &Inst::AtomicCASLoop {
                 ty,
                 addr,

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

@@ -2342,6 +2342,11 @@
       (atomic_rmw_loop (AtomicRMWLoopOp.Xchg) addr src ty flags))
 
 ;;;; Rules for `AtomicCAS` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(rule 2 (lower (and (use_lse)
+                  (atomic_cas $I128 (little_or_native_endian flags) addr expect replace)))
+      (lse_atomic_cas_128 addr expect replace flags))
+
 (rule 1 (lower (and (use_lse)
                   (atomic_cas (valid_atomic_transaction ty) (little_or_native_endian flags) addr src1 src2)))
       (lse_atomic_cas addr src1 src2 ty flags))

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

@@ -25,7 +25,7 @@ use crate::{
     },
     isa::aarch64::abi::AArch64MachineDeps,
     isa::aarch64::inst::SImm7Scaled,
-    isa::aarch64::inst::args::{ShiftOp, ShiftOpShiftImm},
+    isa::aarch64::inst::args::{AtomicCAS128Args, ShiftOp, ShiftOpShiftImm},
     machinst::{
         CallArgList, CallRetList, InstOutput, MachInst, VCodeConstant, VCodeConstantData,
         abi::ArgPair, ty_bits,
@@ -42,6 +42,7 @@ type BoxReturnCallIndInfo = Box<ReturnCallInfo<Reg>>;
 type VecMachLabel = Vec<MachLabel>;
 type BoxExternalName = Box<ExternalName>;
 type VecArgPair = Vec<ArgPair>;
+type BoxAtomicCAS128Args = Box<AtomicCAS128Args>;
 
 /// The main entry point for lowering with ISLE.
 pub(crate) fn lower(
@@ -178,6 +179,29 @@ impl Context for IsleContext<'_, '_, MInst, AArch64Backend> {
         }
     }
 
+    fn atomic_cas_128_args(
+        &mut self,
+        rd_lo: WritableReg,
+        rd_hi: WritableReg,
+        rs_lo: Reg,
+        rs_hi: Reg,
+        rt_lo: Reg,
+        rt_hi: Reg,
+        rn: Reg,
+        flags: MemFlagsData,
+    ) -> BoxAtomicCAS128Args {
+        Box::new(AtomicCAS128Args {
+            rd_lo,
+            rd_hi,
+            rs_lo,
+            rs_hi,
+            rt_lo,
+            rt_hi,
+            rn,
+            flags,
+        })
+    }
+
     fn use_dotprod(&mut self, _: Inst) -> Option<()> {
         if self.backend.isa_flags.has_dotprod() {
             Some(())

cranelift/filetests/filetests/isa/aarch64/atomic-cas-128-lse.clif

@@ -0,0 +1,104 @@
+test compile precise-output
+target aarch64 has_lse
+
+function %atomic_cas_i128(i64, i128, i128) -> i128 {
+block0(v0: i64, v1: i128, v2: i128):
+    v3 = atomic_cas.i128 v0, v1, v2
+    return v3
+}
+
+; VCode:
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+;   str x28, [sp, #-16]!
+;   stp x26, x27, [sp, #-16]!
+;   stp x24, x25, [sp, #-16]!
+; block0:
+;   mov x24, x2
+;   mov x25, x3
+;   mov x26, x4
+;   mov x27, x5
+;   mov x28, x0
+;   caspal x24, x25, x24, x25, x26, x27, [x28]
+;   mov x0, x24
+;   mov x1, x25
+;   ldp x24, x25, [sp], #16
+;   ldp x26, x27, [sp], #16
+;   ldr x28, [sp], #16
+;   ldp fp, lr, [sp], #16
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   stp x29, x30, [sp, #-0x10]!
+;   mov x29, sp
+;   str x28, [sp, #-0x10]!
+;   stp x26, x27, [sp, #-0x10]!
+;   stp x24, x25, [sp, #-0x10]!
+; block1: ; offset 0x14
+;   mov x24, x2
+;   mov x25, x3
+;   mov x26, x4
+;   mov x27, x5
+;   mov x28, x0
+;   caspal x24, x25, x26, x27, [x28] ; trap: heap_oob
+;   mov x0, x24
+;   mov x1, x25
+;   ldp x24, x25, [sp], #0x10
+;   ldp x26, x27, [sp], #0x10
+;   ldr x28, [sp], #0x10
+;   ldp x29, x30, [sp], #0x10
+;   ret
+
+function %atomic_cas_i128_success(i64, i128, i128) -> i8 {
+block0(v0: i64, v1: i128, v2: i128):
+    v3 = atomic_cas.i128 v0, v1, v2
+    v4 = icmp eq v3, v1
+    return v4
+}
+
+; VCode:
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+;   str x28, [sp, #-16]!
+;   stp x26, x27, [sp, #-16]!
+;   stp x24, x25, [sp, #-16]!
+; block0:
+;   mov x26, x4
+;   mov x27, x5
+;   mov x28, x0
+;   mov x24, x2
+;   mov x25, x3
+;   caspal x24, x25, x24, x25, x26, x27, [x28]
+;   subs xzr, x24, x2
+;   ccmp x25, x3, #nzcv, eq
+;   cset x0, eq
+;   ldp x24, x25, [sp], #16
+;   ldp x26, x27, [sp], #16
+;   ldr x28, [sp], #16
+;   ldp fp, lr, [sp], #16
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   stp x29, x30, [sp, #-0x10]!
+;   mov x29, sp
+;   str x28, [sp, #-0x10]!
+;   stp x26, x27, [sp, #-0x10]!
+;   stp x24, x25, [sp, #-0x10]!
+; block1: ; offset 0x14
+;   mov x26, x4
+;   mov x27, x5
+;   mov x28, x0
+;   mov x24, x2
+;   mov x25, x3
+;   caspal x24, x25, x26, x27, [x28] ; trap: heap_oob
+;   cmp x24, x2
+;   ccmp x25, x3, #0, eq
+;   cset x0, eq
+;   ldp x24, x25, [sp], #0x10
+;   ldp x26, x27, [sp], #0x10
+;   ldr x28, [sp], #0x10
+;   ldp x29, x30, [sp], #0x10
+;   ret
+

cranelift/filetests/filetests/runtests/atomic-128-cas-lse.clif

@@ -0,0 +1,32 @@
+test interpret
+test run
+set enable_llvm_abi_extensions
+set enable_multi_ret_implicit_sret
+target aarch64 has_lse
+
+function %atomic_cas(i128, i128, i128) -> i128, i128 {
+    ss0 = explicit_slot 16
+
+block0(v0: i128, v1: i128, v2: i128):
+    v6 = stack_addr.i64 ss0
+    store.i128 notrap v0, v6
+    v3 = stack_addr.i64 ss0
+    v4 = atomic_cas.i128 v3, v1, v2
+    v7 = stack_addr.i64 ss0
+    v5 = load.i128 notrap v7
+    return v5, v4
+}
+
+; run: %atomic_cas(0, 0, 2) == [2, 0]
+; run: %atomic_cas(1, 0, 2) == [1, 1]
+; run: %atomic_cas(0, 1, 2) == [0, 0]
+; run: %atomic_cas(0, 0xC0FFEEEE_ABCDEF01_00000000_00000000, 0xDECAFFFF_12345678) == [0, 0]
+; run: %atomic_cas(0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678, 0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678, 0xFEDCBA98_76543210_F7E6D5C4_B3A29180) == [0xFEDCBA98_76543210_F7E6D5C4_B3A29180, 0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678]
+
+; Only the high half differs, so a lowering that swapped the low/high halves of
+; either operand pair would still report success here.
+; run: %atomic_cas(0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678, 0x11111111_11111111_DECAFFFF_12345678, 0) == [0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678, 0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678]
+; Only the low half differs.
+; run: %atomic_cas(0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678, 0xC0FFEEEE_ABCDEF01_11111111_11111111, 0) == [0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678, 0xC0FFEEEE_ABCDEF01_DECAFFFF_12345678]
+; The replacement's halves must not be swapped on a successful swap.
+; run: %atomic_cas(0, 0, 0xFEDCBA98_76543210_F7E6D5C4_B3A29180) == [0xFEDCBA98_76543210_F7E6D5C4_B3A29180, 0]