Cranelift: aarch64: fix preserve-all to save full vector registers. (#12944)

It turns out that the `preserve-all` ABI was only preserving some, not
all (false advertising!): specifically, the aarch64 ABI code was
continuing to use low-64-bit loads/stores on vector/float registers, as
it does for the ordinary AAPCS (SysV) calling convention. `PreserveAll`
specifically indicates that the *entire* vector register should be
saved; so now we do that.

Author: cfallin

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

@@ -38,7 +38,10 @@ impl From<StackAMode> for AMode {
 
 // Returns the size of stack space needed to store the
 // `clobbered_callee_saved` registers.
-fn compute_clobber_size(clobbered_callee_saves: &[Writable<RealReg>]) -> u32 {
+fn compute_clobber_size(
+    call_conv: isa::CallConv,
+    clobbered_callee_saves: &[Writable<RealReg>],
+) -> u32 {
     let mut int_regs = 0;
     let mut vec_regs = 0;
     for &reg in clobbered_callee_saves {
@@ -55,16 +58,22 @@ fn compute_clobber_size(clobbered_callee_saves: &[Writable<RealReg>]) -> u32 {
 
     // Round up to multiple of 2, to keep 16-byte stack alignment.
     let int_save_bytes = (int_regs + (int_regs & 1)) * 8;
-    // The Procedure Call Standard for the Arm 64-bit Architecture
-    // (AAPCS64, including several related ABIs such as the one used by
-    // Windows) mandates saving only the bottom 8 bytes of the vector
-    // registers, so we round up the number of registers to ensure
-    // proper stack alignment (similarly to the situation with
-    // `int_reg`).
-    let vec_reg_size = 8;
-    let vec_save_padding = vec_regs & 1;
-    // FIXME: SVE: ABI is different to Neon, so do we treat all vec regs as Z-regs?
-    let vec_save_bytes = (vec_regs + vec_save_padding) * vec_reg_size;
+    let vec_save_bytes = if call_conv == isa::CallConv::PreserveAll {
+        // In the PreserveAll ABI, we save the entire vector register,
+        // i.e., all 128 bits.
+        vec_regs * 16
+    } else {
+        // The Procedure Call Standard for the Arm 64-bit Architecture
+        // (AAPCS64, including several related ABIs such as the one used by
+        // Windows) mandates saving only the bottom 8 bytes of the vector
+        // registers, so we round up the number of registers to ensure
+        // proper stack alignment (similarly to the situation with
+        // `int_reg`).
+        let vec_reg_size = 8;
+        let vec_save_padding = vec_regs & 1;
+        // FIXME: SVE: ABI is different to Neon, so do we treat all vec regs as Z-regs?
+        (vec_regs + vec_save_padding) * vec_reg_size
+    };
 
     int_save_bytes + vec_save_bytes
 }
@@ -714,7 +723,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
     }
 
     fn gen_clobber_save(
-        _call_conv: isa::CallConv,
+        call_conv: isa::CallConv,
         flags: &settings::Flags,
         frame_layout: &FrameLayout,
     ) -> SmallVec<[Inst; 16]> {
@@ -857,74 +866,91 @@ impl ABIMachineSpec for AArch64MachineDeps {
             }
         }
 
-        let store_vec_reg = |rd| Inst::FpuStore64 {
-            rd,
-            mem: AMode::SPPreIndexed {
-                simm9: SImm9::maybe_from_i64(-clobber_offset_change).unwrap(),
-            },
-            flags: MemFlags::trusted(),
-        };
-        let iter = clobbered_vec.chunks_exact(2);
+        if call_conv == isa::CallConv::PreserveAll {
+            // Store full vector registers in PreserveAll convention.
+            for reg in clobbered_vec.iter().rev() {
+                let inst = Inst::FpuStore128 {
+                    rd: reg.to_reg().into(),
+                    mem: AMode::SPPreIndexed {
+                        simm9: SImm9::maybe_from_i64(-clobber_offset_change).unwrap(),
+                    },
+                    flags: MemFlags::trusted(),
+                };
+                insts.push(inst);
+                // N.B.: no unwind info: we don't have a way to
+                // represent "full register" anyway.
+            }
+        } else {
+            let store_vec_reg_half = |rd| Inst::FpuStore64 {
+                rd,
+                mem: AMode::SPPreIndexed {
+                    simm9: SImm9::maybe_from_i64(-clobber_offset_change).unwrap(),
+                },
+                flags: MemFlags::trusted(),
+            };
+            let iter = clobbered_vec.chunks_exact(2);
 
-        if let [rd] = iter.remainder() {
-            let rd: Reg = rd.to_reg().into();
+            if let [rd] = iter.remainder() {
+                let rd: Reg = rd.to_reg().into();
 
-            debug_assert_eq!(rd.class(), RegClass::Float);
-            insts.push(store_vec_reg(rd));
+                debug_assert_eq!(rd.class(), RegClass::Float);
+                insts.push(store_vec_reg_half(rd));
 
-            if flags.unwind_info() {
-                clobber_offset -= clobber_offset_change as u32;
-                insts.push(Inst::Unwind {
-                    inst: UnwindInst::SaveReg {
-                        clobber_offset,
-                        reg: rd.to_real_reg().unwrap(),
-                    },
-                });
+                if flags.unwind_info() {
+                    clobber_offset -= clobber_offset_change as u32;
+                    insts.push(Inst::Unwind {
+                        inst: UnwindInst::SaveReg {
+                            clobber_offset,
+                            reg: rd.to_real_reg().unwrap(),
+                        },
+                    });
+                }
             }
-        }
 
-        let store_vec_reg_pair = |rt, rt2| {
-            let clobber_offset_change = 16;
+            let store_vec_reg_half_pair = |rt, rt2| {
+                let clobber_offset_change = 16;
 
-            (
-                Inst::FpuStoreP64 {
-                    rt,
-                    rt2,
-                    mem: PairAMode::SPPreIndexed {
-                        simm7: SImm7Scaled::maybe_from_i64(-clobber_offset_change, F64).unwrap(),
+                (
+                    Inst::FpuStoreP64 {
+                        rt,
+                        rt2,
+                        mem: PairAMode::SPPreIndexed {
+                            simm7: SImm7Scaled::maybe_from_i64(-clobber_offset_change, F64)
+                                .unwrap(),
+                        },
+                        flags: MemFlags::trusted(),
                     },
-                    flags: MemFlags::trusted(),
-                },
-                clobber_offset_change as u32,
-            )
-        };
-        let mut iter = iter.rev();
+                    clobber_offset_change as u32,
+                )
+            };
+            let mut iter = iter.rev();
 
-        while let Some([rt, rt2]) = iter.next() {
-            let rt: Reg = rt.to_reg().into();
-            let rt2: Reg = rt2.to_reg().into();
+            while let Some([rt, rt2]) = iter.next() {
+                let rt: Reg = rt.to_reg().into();
+                let rt2: Reg = rt2.to_reg().into();
 
-            debug_assert_eq!(rt.class(), RegClass::Float);
-            debug_assert_eq!(rt2.class(), RegClass::Float);
+                debug_assert_eq!(rt.class(), RegClass::Float);
+                debug_assert_eq!(rt2.class(), RegClass::Float);
 
-            let (inst, clobber_offset_change) = store_vec_reg_pair(rt, rt2);
+                let (inst, clobber_offset_change) = store_vec_reg_half_pair(rt, rt2);
 
-            insts.push(inst);
+                insts.push(inst);
 
-            if flags.unwind_info() {
-                clobber_offset -= clobber_offset_change;
-                insts.push(Inst::Unwind {
-                    inst: UnwindInst::SaveReg {
-                        clobber_offset,
-                        reg: rt.to_real_reg().unwrap(),
-                    },
-                });
-                insts.push(Inst::Unwind {
-                    inst: UnwindInst::SaveReg {
-                        clobber_offset: clobber_offset + clobber_offset_change / 2,
-                        reg: rt2.to_real_reg().unwrap(),
-                    },
-                });
+                if flags.unwind_info() {
+                    clobber_offset -= clobber_offset_change;
+                    insts.push(Inst::Unwind {
+                        inst: UnwindInst::SaveReg {
+                            clobber_offset,
+                            reg: rt.to_real_reg().unwrap(),
+                        },
+                    });
+                    insts.push(Inst::Unwind {
+                        inst: UnwindInst::SaveReg {
+                            clobber_offset: clobber_offset + clobber_offset_change / 2,
+                            reg: rt2.to_real_reg().unwrap(),
+                        },
+                    });
+                }
             }
         }
 
@@ -943,7 +969,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
     }
 
     fn gen_clobber_restore(
-        _call_conv: isa::CallConv,
+        call_conv: isa::CallConv,
         _flags: &settings::Flags,
         frame_layout: &FrameLayout,
     ) -> SmallVec<[Inst; 16]> {
@@ -956,40 +982,55 @@ impl ABIMachineSpec for AArch64MachineDeps {
             insts.extend(Self::gen_sp_reg_adjust(stack_size as i32));
         }
 
-        let load_vec_reg = |rd| Inst::FpuLoad64 {
-            rd,
-            mem: AMode::SPPostIndexed {
-                simm9: SImm9::maybe_from_i64(16).unwrap(),
-            },
-            flags: MemFlags::trusted(),
-        };
-        let load_vec_reg_pair = |rt, rt2| Inst::FpuLoadP64 {
-            rt,
-            rt2,
-            mem: PairAMode::SPPostIndexed {
-                simm7: SImm7Scaled::maybe_from_i64(16, F64).unwrap(),
-            },
-            flags: MemFlags::trusted(),
-        };
+        if call_conv == isa::CallConv::PreserveAll {
+            for reg in clobbered_vec.iter() {
+                let inst = Inst::FpuLoad128 {
+                    rd: reg.map(|r| r.into()),
+                    mem: AMode::SPPostIndexed {
+                        simm9: SImm9::maybe_from_i64(16).unwrap(),
+                    },
+                    flags: MemFlags::trusted(),
+                };
+                insts.push(inst);
+                // N.B.: no unwind info; we don't have a way to
+                // represent "full vector register saved" anyway.
+            }
+        } else {
+            let load_vec_reg_half = |rd| Inst::FpuLoad64 {
+                rd,
+                mem: AMode::SPPostIndexed {
+                    simm9: SImm9::maybe_from_i64(16).unwrap(),
+                },
+                flags: MemFlags::trusted(),
+            };
+            let load_vec_reg_half_pair = |rt, rt2| Inst::FpuLoadP64 {
+                rt,
+                rt2,
+                mem: PairAMode::SPPostIndexed {
+                    simm7: SImm7Scaled::maybe_from_i64(16, F64).unwrap(),
+                },
+                flags: MemFlags::trusted(),
+            };
 
-        let mut iter = clobbered_vec.chunks_exact(2);
+            let mut iter = clobbered_vec.chunks_exact(2);
 
-        while let Some([rt, rt2]) = iter.next() {
-            let rt: Writable<Reg> = rt.map(|r| r.into());
-            let rt2: Writable<Reg> = rt2.map(|r| r.into());
+            while let Some([rt, rt2]) = iter.next() {
+                let rt: Writable<Reg> = rt.map(|r| r.into());
+                let rt2: Writable<Reg> = rt2.map(|r| r.into());
 
-            debug_assert_eq!(rt.to_reg().class(), RegClass::Float);
-            debug_assert_eq!(rt2.to_reg().class(), RegClass::Float);
-            insts.push(load_vec_reg_pair(rt, rt2));
-        }
+                debug_assert_eq!(rt.to_reg().class(), RegClass::Float);
+                debug_assert_eq!(rt2.to_reg().class(), RegClass::Float);
+                insts.push(load_vec_reg_half_pair(rt, rt2));
+            }
 
-        debug_assert!(iter.remainder().len() <= 1);
+            debug_assert!(iter.remainder().len() <= 1);
 
-        if let [rd] = iter.remainder() {
-            let rd: Writable<Reg> = rd.map(|r| r.into());
+            if let [rd] = iter.remainder() {
+                let rd: Writable<Reg> = rd.map(|r| r.into());
 
-            debug_assert_eq!(rd.to_reg().class(), RegClass::Float);
-            insts.push(load_vec_reg(rd));
+                debug_assert_eq!(rd.to_reg().class(), RegClass::Float);
+                insts.push(load_vec_reg_half(rd));
+            }
         }
 
         let mut iter = clobbered_int.chunks_exact(2);
@@ -1150,7 +1191,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
         regs.sort_unstable();
 
         // Compute clobber size.
-        let clobber_size = compute_clobber_size(&regs);
+        let clobber_size = compute_clobber_size(call_conv, &regs);
 
         // Compute linkage frame size.
         let setup_area_size = if flags.preserve_frame_pointers()