[AArch64] Allocate emergency spillslot when using fixed-length SVE.

[sdesmalen]

Summary

The SVE reg+imm addressing modes require the immediate to be a multiple
of VL. When it is not a multiple of VL, it can use 0 for the immediate
and use an add instruction to materialise the base-register.

When using SVE instructions to access fixed-length vectors on the stack
that are not VL-aligned, the frame should anticipate that it might require
an emergency spill-slot to accomodate for an extra 'add' to materialize
the base-address of the load when there are no free registers available.

Fixes https://github.com/llvm/llvm-project/issues/60918

[peterwaller-arm]

Makes sense to me, one minor suggestion. I note that this introduces a few loads/stores to the stack for fixed length vectorization, but having it work correctly is the priority, those may be possible to remove later.

llvm/test/CodeGen/AArch64/framelayout-sve-scavengingslot.mir

30

Since this test lacks check lines I presume it is meant to check only that llc exits successfully. Could you have checklines or a comment explaining what failure it is meant to catch?

[sdesmalen]

Added missing CHECK lines

[sdesmalen]

llvm/test/CodeGen/AArch64/framelayout-sve-scavengingslot.mir

30

Sorry I forgot to put in the CHECK lines. Thanks for catching this!

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

2979

Is this initialisation required?

2985–2986

This seems a bit "belt and braces". You sure there's no other solution?

2992

Is the StackID of this frame index also relevant? As in, you only care about accesses to the non-scalable stack region?

[efriedma]

Does using an SVE instruction to access a fixed stack slot ever work? I mean, I guess if you get lucky, the fixed offset is zero, but that should happen rarely. Maybe we'd get better code if we expand out the arithmetic earlier. (Doesn't need to block this patch, just something to think about.)

[sdesmalen]

Abandoning since the alternative fix has now landed (D146056)

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

for (const RegPairInfo &RPI : RegPairs) {
         continue;
       (void)EmitMI(RPI);
     }
   }
 
   return true;
 }
 
+/// Returns true when SVE load/store instructions are used to access
+/// fixed-length locals on the stack.
+///
+/// This information can be used to determine whether an emergency spill-slot
+/// may be required to cater for an explicit 'add' for materializing the SVE
+/// load/store base-address when the address being loaded from cannot be
+/// represented with <basereg> + imm, when `imm` is not a multiple of VL-bytes.
+bool useSVEAddressingModesForFixedLengthLocals(const MachineFunction &MF) {
+  const auto &STI = MF.getSubtarget<AArch64Subtarget>();
+  if (!STI.useSVEForFixedLengthVectors())
+    return false;
+
+  TypeSize Scale = TypeSize::Fixed(1);
+  unsigned Width;
+  int64_t MinOffset, MaxOffset;
+  const AArch64InstrInfo *TII =
+      static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
+
+  for (auto MBBI = MF.begin(), MBBE = MF.end(); MBBI != MBBE; ++MBBI) {
+    for (auto I = MBBI->begin(), E = MBBI->end(); I != E; ++I) {
+      unsigned Opc = I->getOpcode();
+      if (!TII->getMemOpInfo(Opc, Scale, Width, MinOffset, MaxOffset) ||
+          !Scale.isScalable())
+        continue;
+      if (any_of(I->operands(),
+                 [](const MachineOperand &Op) { return Op.isFI(); }))
+        return true;
+    }
+  }
+  return false;
+}
+
 void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
                                                 BitVector &SavedRegs,
                                                 RegScavenger *RS) const {
   // All calls are tail calls in GHC calling conv, and functions have no
   // prologue/epilogue.
   if (MF.getFunction().getCallingConv() == CallingConv::GHC)
     return;
 
void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
   // for our own arguments or during a call. Include that in calculating whether
   // we have complicated addressing concerns.
   int64_t CalleeStackUsed = 0;
   for (int I = MFI.getObjectIndexBegin(); I != 0; ++I) {
     int64_t FixedOff = MFI.getObjectOffset(I);
     if (FixedOff > CalleeStackUsed) CalleeStackUsed = FixedOff;
   }
 
-  // Conservatively always assume BigStack when there are SVE spills.
-  bool BigStack = SVEStackSize || (EstimatedStackSize + CSStackSize +
-                                   CalleeStackUsed) > EstimatedStackSizeLimit;
+  bool BigStack = SVEStackSize ||
+                  (EstimatedStackSize + CSStackSize + CalleeStackUsed) >
+                      EstimatedStackSizeLimit ||
+                  useSVEAddressingModesForFixedLengthLocals(MF);
   if (BigStack || !CanEliminateFrame || RegInfo->cannotEliminateFrame(MF))
     AFI->setHasStackFrame(true);
 
   // Estimate if we might need to scavenge a register at some point in order
   // to materialize a stack offset. If so, either spill one additional
   // callee-saved register or reserve a special spill slot to facilitate
   // register scavenging. If we already spilled an extra callee-saved register
   // above to keep the number of spills even, we don't need to do anything else

llvm/test/CodeGen/AArch64/framelayout-sve-scavengingslot.mir

-# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog -mattr=+sve %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog -force-streaming-compatible-sve -mattr=+sve %s -o - | FileCheck %s
 ---
 # This test verifies that the emergency scavenging slot is located near the SP/BP.
 name: LateScavengingSlot
 # CHECK-LABEL: name: LateScavengingSlot
 # CHECK: bb.0:
 # CHECK:      $sp = frame-setup ADDVL_XXI $sp, -1
 # CHECK-NEXT: $sp = frame-setup SUBXri $sp, 8, 12
 # CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0
body: |
   bb.0:
     liveins: $x0, $x8
     STRXui $x0, %stack.1, 0
     B %bb.1
   bb.1:
     liveins: $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x16, $x17, $x18, $x19, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $lr
     RET_ReallyLR implicit $x19, implicit $x20, implicit $x21, implicit $x22, implicit $x23, implicit $x24, implicit $x25, implicit $x26, implicit $x27, implicit $x28, implicit $lr
 ...
+---
+# This test verifies that an emergency scavenging slot is allocated when using fixed-length SVE
+name: RequiresScaveningSlotFixedLengthSVE
+# CHECK-LABEL: name: RequiresScaveningSlotFixedLengthSVE
+# CHECK:        $sp = frame-setup SUBXri $sp, 32, 0
+# CHECK:        frame-setup STPXi killed $fp, killed $lr, $sp, 2
+# CHECK-NEXT:   $fp = frame-setup ADDXri $sp, 16, 0
+# CHECK:        STRXui killed $x8, $sp, 1
+# CHECK-NEXT:   $x8 = ADDXri $fp, 16, 0
+# CHECK-NEXT:   ST1W_IMM $z0, $p0, killed $x8, 0
+# CHECK-NEXT:   $x8 = LDRXui $sp, 1
+tracksRegLiveness: true
+frameInfo:
+  isFrameAddressTaken: true
+fixedStack:
+  - { id: 0, stack-id: default, size: 16, alignment: 8 }
+body: |
+  bb.0:
+    liveins: $z0, $p0
+    ST1W_IMM $z0, $p0, %fixed-stack.0, 0
+    B %bb.1
+  bb.1:
+    liveins: $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x16, $x17, $x18, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $lr
+    RET_ReallyLR  implicit $x0, implicit $x1, implicit $x2, implicit $x3, implicit $x4, implicit $x5, implicit $x6, implicit $x7, implicit $x8, implicit $x9, implicit $x10, implicit $x11, implicit $x12, implicit $x13, implicit $x14, implicit $x15, implicit $x16, implicit $x17, implicit $x18, implicit $x20, implicit $x21, implicit $x22, implicit $x23, implicit $x24, implicit $x25, implicit $x26, implicit $x27, implicit $x28, implicit $lr
+...

llvm/test/CodeGen/AArch64/sve-fixed-ld2-alloca.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s | FileCheck %s
 
 target triple = "aarch64-unknown-linux-gnu"
 
 declare void @def(ptr)
 
 define void @st1d_fixed(ptr %ptr) #0 {
 ; CHECK-LABEL: st1d_fixed:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub sp, sp, #144
-; CHECK-NEXT:    stp x30, x19, [sp, #128] // 16-byte Folded Spill
+; CHECK-NEXT:    sub sp, sp, #160
+; CHECK-NEXT:    stp x30, x19, [sp, #144] // 16-byte Folded Spill
 ; CHECK-NEXT:    mov x19, x0
 ; CHECK-NEXT:    mov x0, sp
+; CHECK-NEXT:    str x29, [sp, #128] // 8-byte Folded Spill
 ; CHECK-NEXT:    bl def
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    ld2d { z0.d, z1.d }, p0/z, [sp]
+; CHECK-NEXT:    ldr x29, [sp, #128] // 8-byte Folded Reload
 ; CHECK-NEXT:    st1d { z0.d }, p0, [x19]
-; CHECK-NEXT:    ldp x30, x19, [sp, #128] // 16-byte Folded Reload
-; CHECK-NEXT:    add sp, sp, #144
+; CHECK-NEXT:    ldp x30, x19, [sp, #144] // 16-byte Folded Reload
+; CHECK-NEXT:    add sp, sp, #160
 ; CHECK-NEXT:    ret
   %alloc = alloca [16 x double]
   call void @def(ptr %alloc)
   %load = load <8 x double>, ptr %alloc
   %strided.vec = shufflevector <8 x double> %load, <8 x double> poison, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
   store <4 x double> %strided.vec, ptr %ptr
   ret void
 }
 
 attributes #0 = { "target-features"="+sve" vscale_range(2,2) nounwind }

llvm/test/CodeGen/AArch64/sve-streaming-mode-fixed-length-ld2-alloca.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -force-streaming-compatible-sve < %s | FileCheck %s
 
 target triple = "aarch64-unknown-linux-gnu"
 
 declare void @def(ptr)
 
 define void @alloc_v4i8(ptr %st_ptr) #0 {
 ; CHECK-LABEL: alloc_v4i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub sp, sp, #32
-; CHECK-NEXT:    stp x30, x19, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    sub sp, sp, #48
+; CHECK-NEXT:    stp x30, x19, [sp, #32] // 16-byte Folded Spill
 ; CHECK-NEXT:    mov x19, x0
-; CHECK-NEXT:    add x0, sp, #12
+; CHECK-NEXT:    add x0, sp, #28
+; CHECK-NEXT:    str x29, [sp, #16] // 8-byte Folded Spill
 ; CHECK-NEXT:    bl def
-; CHECK-NEXT:    add x8, sp, #12
+; CHECK-NEXT:    add x8, sp, #28
 ; CHECK-NEXT:    ptrue p0.b, vl2
 ; CHECK-NEXT:    ld2b { z0.b, z1.b }, p0/z, [x8]
 ; CHECK-NEXT:    ptrue p0.s, vl2
+; CHECK-NEXT:    ldr x29, [sp, #16] // 8-byte Folded Reload
 ; CHECK-NEXT:    mov z2.b, z0.b[1]
 ; CHECK-NEXT:    fmov w8, s0
 ; CHECK-NEXT:    fmov w9, s2
-; CHECK-NEXT:    stp w8, w9, [sp]
-; CHECK-NEXT:    ldr d0, [sp]
+; CHECK-NEXT:    stp w8, w9, [sp, #8]
+; CHECK-NEXT:    ldr d0, [sp, #8]
 ; CHECK-NEXT:    st1b { z0.s }, p0, [x19]
-; CHECK-NEXT:    ldp x30, x19, [sp, #16] // 16-byte Folded Reload
-; CHECK-NEXT:    add sp, sp, #32
+; CHECK-NEXT:    ldp x30, x19, [sp, #32] // 16-byte Folded Reload
+; CHECK-NEXT:    add sp, sp, #48
 ; CHECK-NEXT:    ret
   %alloc = alloca [4 x i8]
   call void @def(ptr %alloc)
   %load = load <4 x i8>, ptr %alloc
   %strided.vec = shufflevector <4 x i8> %load, <4 x i8> poison, <2 x i32> <i32 0, i32 2>
   store <2 x i8> %strided.vec, ptr %st_ptr
   ret void
 }
 
 define void @alloc_v6i8(ptr %st_ptr) #0 {
 ; CHECK-LABEL: alloc_v6i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sub sp, sp, #48
 ; CHECK-NEXT:    stp x30, x19, [sp, #32] // 16-byte Folded Spill
 ; CHECK-NEXT:    mov x19, x0
 ; CHECK-NEXT:    add x0, sp, #24
+; CHECK-NEXT:    str x29, [sp, #16] // 8-byte Folded Spill
 ; CHECK-NEXT:    bl def
 ; CHECK-NEXT:    add x8, sp, #24
 ; CHECK-NEXT:    ptrue p0.b, vl3
 ; CHECK-NEXT:    ld2b { z0.b, z1.b }, p0/z, [x8]
 ; CHECK-NEXT:    ptrue p0.h, vl4
 ; CHECK-NEXT:    fmov w8, s1
 ; CHECK-NEXT:    mov z2.b, z1.b[3]
 ; CHECK-NEXT:    mov z3.b, z1.b[2]
 ; CHECK-NEXT:    mov z0.b, z1.b[1]
 ; CHECK-NEXT:    fmov w9, s2
 ; CHECK-NEXT:    fmov w10, s3
-; CHECK-NEXT:    strh w8, [sp, #8]
+; CHECK-NEXT:    strh w8, [sp]
 ; CHECK-NEXT:    fmov w8, s0
-; CHECK-NEXT:    strh w9, [sp, #14]
-; CHECK-NEXT:    strh w10, [sp, #12]
-; CHECK-NEXT:    strh w8, [sp, #10]
-; CHECK-NEXT:    add x8, sp, #20
-; CHECK-NEXT:    ldr d0, [sp, #8]
+; CHECK-NEXT:    strh w9, [sp, #6]
+; CHECK-NEXT:    strh w10, [sp, #4]
+; CHECK-NEXT:    strh w8, [sp, #2]
+; CHECK-NEXT:    add x8, sp, #12
+; CHECK-NEXT:    ldr d0, [sp]
 ; CHECK-NEXT:    st1b { z0.h }, p0, [x8]
-; CHECK-NEXT:    ldrh w8, [sp, #20]
+; CHECK-NEXT:    ldrh w8, [sp, #12]
 ; CHECK-NEXT:    strb w10, [x19, #2]
+; CHECK-NEXT:    ldr x29, [sp, #16] // 8-byte Folded Reload
 ; CHECK-NEXT:    strh w8, [x19]
 ; CHECK-NEXT:    ldp x30, x19, [sp, #32] // 16-byte Folded Reload
 ; CHECK-NEXT:    add sp, sp, #48
 ; CHECK-NEXT:    ret
   %alloc = alloca [6 x i8]
   call void @def(ptr %alloc)
   %load = load <6 x i8>, ptr %alloc
   %strided.vec = shufflevector <6 x i8> %load, <6 x i8> poison, <3 x i32> <i32 1, i32 3, i32 5>

 
 define void @alloc_v8f64(ptr %st_ptr) #0 {
 ; CHECK-LABEL: alloc_v8f64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sub sp, sp, #96
 ; CHECK-NEXT:    stp x20, x19, [sp, #80] // 16-byte Folded Spill
 ; CHECK-NEXT:    mov x19, x0
 ; CHECK-NEXT:    mov x0, sp
-; CHECK-NEXT:    str x30, [sp, #64] // 8-byte Folded Spill
+; CHECK-NEXT:    stp x29, x30, [sp, #64] // 16-byte Folded Spill
 ; CHECK-NEXT:    mov x20, sp
 ; CHECK-NEXT:    bl def
 ; CHECK-NEXT:    mov x8, #4
 ; CHECK-NEXT:    ptrue p0.d, vl2
 ; CHECK-NEXT:    ld2d { z0.d, z1.d }, p0/z, [sp]
 ; CHECK-NEXT:    ld2d { z2.d, z3.d }, p0/z, [x20, x8, lsl #3]
-; CHECK-NEXT:    ldr x30, [sp, #64] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp x29, x30, [sp, #64] // 16-byte Folded Reload
 ; CHECK-NEXT:    stp q0, q2, [x19]
 ; CHECK-NEXT:    ldp x20, x19, [sp, #80] // 16-byte Folded Reload
 ; CHECK-NEXT:    add sp, sp, #96
 ; CHECK-NEXT:    ret
   %alloc = alloca [8 x double]
   call void @def(ptr %alloc)
   %load = load <8 x double>, ptr %alloc
   %strided.vec = shufflevector <8 x double> %load, <8 x double> poison, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
   store <4 x double> %strided.vec, ptr %st_ptr
   ret void
 }
 
 attributes #0 = { "target-features"="+sve" nounwind}