[AArch64][SVE] Correctly allocate scavenging slot in presence of SVE.

[sdesmalen]

Summary

This patch addresses two issues:

• Forces the availability of the base-pointer (x19) when the frame has both scalable vectors and variable-length arrays. Otherwise it will be expensive to access non-SVE locals.

• In presence of SVE stack objects, it will allocate the emergency scavenging slot close to the SP, so that they can be accessed from the SP or BP if available. If accessed from the frame-pointer, it will otherwise need an extra register to access the scavenging slot because of mixed scalable/non-scalable addressing modes.

[efriedma]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

393

The problem here, I guess, is the spill slot can't be addressed relative to SP? That makes sense, I guess... but do we also need to check for stack realignment here?

[david-arm]

This fixes an issue building SLEEF with SVE at -O0. It needs rebasing though as it doesn't apply cleanly to HEAD.

[fpetrogalli]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

2589–2590

I think you should update the description of the method to reflect the new interface.

2658–2662

I think you can remove these two and pass them directly inside the call as [](int, int64_t){} and [](int, unsigned){}

You could actually make these two values the default argument of the method, so you don't have to bother specifying them here? (up to you, though)

2666–2672

Nit: these should also be passed directly as argument to the call, no need to set up variables that are not used anywhere else.

I say this is a nit because the first lambda is 4 lines long, so it might not look beautiful once placed inside the function parameters... :)

[sdesmalen]

• Rebased patch.
• Updated framelayout-sve.mir
• Addressed comments.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

2666–2672

Yes this was deliberate as I think it reads a bit easier than one big expression.

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

393

Apologies for my delayed repsonse, I completely forgot about this patch.

The problem here, I guess, is the spill slot can't be addressed relative to SP?

In that case it will be allocated close to the BP. If there are variable-sized objects the choices are either BP or FP, otherwise the choices are SP or FP. Because the emergency spillslot needs to be accessed using the normal (non-SVE) addressing modes in one go (i.e. without requiring another register to materialize the address to the emergency slot), it needs to be able to access it directly from Pointer + <offset>, so BP or SP are really the only options. Otherwise it would need a temporary register to calculate FP - sizeof(SVE area), which itself would require the emergency slot to materialise the temporary.

do we also need to check for stack realignment here?

I don't think that's specifically needed, because hasFP() will always evaluate to true if the stack requires realignment. Do you want me to add an assert for that?

[efriedma]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

393

I think you're misunderstanding what I'm saying about stack realignment. I think the interaction with stack realignment is actually broken (and nobody has run into it because stack realignment is rare on AArch64). I'm thinking that the condition here should just be TFI.hasFP(MF) && !hasBasePointer(MF).

[sdesmalen]

• Changed condition to useFPForScavengingIndex
• Added new test framelayout-scavengingslot.mir to guard that stack realignment (without SVE) allocates the scavenging slot in the right place.
• Simplified patch, no need to refactor estimateSVEStackObjectOffsets
• Renamed misleading/wrong title of this patch.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

393

Thanks for clarifying, I see what you mean now. The realignment case was already covered by code in PEI so it's not currently broken, but I've added the condition to this function anyway. I've also created framelayout-scavengingslot.mir to guard the correct behaviour.

I don't think we need to check for !hasBasePointer(MF) because the availability of a base pointer doesn't really matter. AIUI the FP is always desired unless there is some runtime-sized area between the FP and the emergency slot.

+-----------------+
|   Callee Saves  |
+-----------------+ <- FP
\ unknown size at \    |
/  compile time   /    | (realignment gap or SVE area)
+-----------------+    /
|<emergency slot> | <-' This always requires more than 1 instruction from FP.
| - - - - - - - - |  
|     Locals      |  
| - - - - - - - - |
|                 | <- Better allocated here, can always access from either SP or BP.
+-----------------+ < BP/SP
\  [Optional VLA] \
/                 /
+-----------------+ < SP

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

393

The realignment case was already covered by code in PEI

Oh, hmm, the caller checks for stack realignment. Makes sense, I guess.

llvm/lib/Target/AArch64/AArch64FrameLowering.h

default:
       return false;
     case TargetStackID::Default:
     case TargetStackID::SVEVector:
     case TargetStackID::NoAlloc:
       return true;
     }
   }
 
+  int64_t estimateSVEStackObjectOffsets(const MachineFrameInfo &MF) const;
+
 private:
   bool shouldCombineCSRLocalStackBump(MachineFunction &MF,
                                       uint64_t StackBumpBytes) const;
 
-  int64_t estimateSVEStackObjectOffsets(MachineFrameInfo &MF) const;
   int64_t assignSVEStackObjectOffsets(MachineFrameInfo &MF,
                                       int &MinCSFrameIndex,
                                       int &MaxCSFrameIndex) const;
   bool shouldCombineCSRLocalStackBumpInEpilogue(MachineBasicBlock &MBB,
                                                 unsigned StackBumpBytes) const;
 };
 
 } // End llvm namespace
 
 #endif

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include <cassert>
 #include <cstdint>
+#include <functional>
 #include <iterator>
 #include <vector>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "frame-info"
 
 static cl::opt<bool> EnableRedZone("aarch64-redzone",
if (AArch64::ZPRRegClass.contains(CS.getReg()) ||
              "SVE CalleeSaves are not consecutive");
 
       Min = std::min(Min, CS.getFrameIdx());
       Max = std::max(Max, CS.getFrameIdx());
     }
   }
   return Min != std::numeric_limits<int>::max();
 }
 
-// Process all the SVE stack objects and determine offsets for each
-// object. If AssignOffsets is true, the offsets get assigned.
-// Fills in the first and last callee-saved frame indices into
-// Min/MaxCSFrameIndex, respectively.
-// Returns the size of the stack.
-static int64_t determineSVEStackObjectOffsets(MachineFrameInfo &MFI,
-                                              int &MinCSFrameIndex,
-                                              int &MaxCSFrameIndex,
-                                              bool AssignOffsets) {
+/// Process all the SVE stack objects and determine offsets for each
+/// object. \p AssignFn can be used to assign the object offsets, e.g.
+/// using `MF.setObjectOffset()`, similarly AlignFn can be used to
+/// set the alignment of an object.
+/// This methods returns the size of the stack, and also returns the frame
+/// indices of the first- and last callee-saves in MinCSFrameIndex and
+/// MaxCSFrameIndex respectively.
+static int64_t
+determineSVEStackObjectOffsets(const MachineFrameInfo &MFI,
+                               int &MinCSFrameIndex, int &MaxCSFrameIndex,
+                               std::function<void(int, int64_t)> AssignFn,
+                               std::function<void(int, Align)> AlignFn) {
   // First process all fixed stack objects.
   int64_t Offset = 0;
   for (int I = MFI.getObjectIndexBegin(); I != 0; ++I)
     if (MFI.getStackID(I) == TargetStackID::SVEVector) {
       int64_t FixedOffset = -MFI.getObjectOffset(I);
       if (FixedOffset > Offset)
         Offset = FixedOffset;
     }
 
-  auto Assign = [&MFI](int FI, int64_t Offset) {
-    LLVM_DEBUG(dbgs() << "alloc FI(" << FI << ") at SP[" << Offset << "]\n");
-    MFI.setObjectOffset(FI, Offset);
-  };
-
   // Then process all callee saved slots.
   if (getSVECalleeSaveSlotRange(MFI, MinCSFrameIndex, MaxCSFrameIndex)) {
     // Make sure to align the last callee save slot.
-    MFI.setObjectAlignment(MaxCSFrameIndex, Align(16));
+    AlignFn(MaxCSFrameIndex, Align(16U));
 
     // Assign offsets to the callee save slots.
     for (int I = MinCSFrameIndex; I <= MaxCSFrameIndex; ++I) {
       Offset += MFI.getObjectSize(I);
       Offset = alignTo(Offset, MFI.getObjectAlign(I));
-      if (AssignOffsets)
-        Assign(I, -Offset);
+      AssignFn(I, -Offset);
     }
   }
 
   // Create a buffer of SVE objects to allocate and sort it.
   SmallVector<int, 8> ObjectsToAllocate;
   for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
     unsigned StackID = MFI.getStackID(I);
     if (StackID != TargetStackID::SVEVector)
for (unsigned FI : ObjectsToAllocate) {
     // FIXME: Given that the length of SVE vectors is not necessarily a power of
     // two, we'd need to align every object dynamically at runtime if the
     // alignment is larger than 16. This is not yet supported.
     if (Alignment > Align(16))
       report_fatal_error(
           "Alignment of scalable vectors > 16 bytes is not yet supported");
 
     Offset = alignTo(Offset + MFI.getObjectSize(FI), Alignment);
-    if (AssignOffsets)
-      Assign(FI, -Offset);
+    AssignFn(FI, -Offset);
   }
 
   return Offset;
 }
 
 int64_t AArch64FrameLowering::estimateSVEStackObjectOffsets(
-    MachineFrameInfo &MFI) const {
+    const MachineFrameInfo &MFI) const {
   int MinCSFrameIndex, MaxCSFrameIndex;
-  return determineSVEStackObjectOffsets(MFI, MinCSFrameIndex, MaxCSFrameIndex, false);
+  return determineSVEStackObjectOffsets(
+      MFI, MinCSFrameIndex, MaxCSFrameIndex,
+      /* AssignFn */ [](int FI, int64_t Offset) {},
+      /* AlignFn */ [](int FI, Align A) {});
 }
 
 int64_t AArch64FrameLowering::assignSVEStackObjectOffsets(
     MachineFrameInfo &MFI, int &MinCSFrameIndex, int &MaxCSFrameIndex) const {
+  auto AssignFn = [&MFI](int FI, int64_t Offset) {
+    LLVM_DEBUG(dbgs() << "alloc FI(" << FI << ") at SP[" << Offset << "]\n");
+    MFI.setObjectOffset(FI, Offset);
+  };
+  auto AlignFn = [&MFI](int FI, Align A) {
+    MFI.setObjectAlignment(FI, A);
+  };
   return determineSVEStackObjectOffsets(MFI, MinCSFrameIndex, MaxCSFrameIndex,
-                                        true);
+                                        AssignFn, AlignFn);
 }
 
 void AArch64FrameLowering::processFunctionBeforeFrameFinalized(
     MachineFunction &MF, RegScavenger *RS) const {
   MachineFrameInfo &MFI = MF.getFrameInfo();
 
   assert(getStackGrowthDirection() == TargetFrameLowering::StackGrowsDown &&
          "Upwards growing stack unsupported");

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

bool AArch64RegisterInfo::hasBasePointer(const MachineFunction &MF) const {
   // from the other direction like the SP normally works.
   //
   // Furthermore, if both variable sized objects are present, and the
   // stack needs to be dynamically re-aligned, the base pointer is the only
   // reliable way to reference the locals.
   if (MFI.hasVarSizedObjects() || MF.hasEHFunclets()) {
     if (needsStackRealignment(MF))
       return true;
+
+    if (MF.getSubtarget<AArch64Subtarget>().hasSVE()) {
+      const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+      // Frames that have variable sized objects and scalable SVE objects,
+      // should always use a basepointer.
+      if (!AFI->hasCalculatedStackSizeSVE() || AFI->getStackSizeSVE())
+        return true;
+    }
+
     // Conservatively estimate whether the negative offset from the frame
     // pointer will be sufficient to reach. If a function has a smallish
     // frame, it's less likely to have lots of spills and callee saved
     // space, so it's all more likely to be within range of the frame pointer.
     // If it's wrong, we'll materialize the constant and still get to the
     // object; it's just suboptimal. Negative offsets use the unscaled
     // load/store instructions, which have a 9-bit signed immediate.
     return MFI.getLocalFrameSize() >= 256;

 
 bool
 AArch64RegisterInfo::useFPForScavengingIndex(const MachineFunction &MF) const {
   // This function indicates whether the emergency spillslot should be placed
   // close to the beginning of the stackframe (closer to FP) or the end
   // (closer to SP).
   //
   // The beginning works most reliably if we have a frame pointer.
+  // In the presence of SVE however, it is better to use SP or BP.
   const AArch64FrameLowering &TFI = *getFrameLowering(MF);
-  return TFI.hasFP(MF);
+  return TFI.hasFP(MF) && !TFI.estimateSVEStackObjectOffsets(MF.getFrameInfo());
 }
 
 bool AArch64RegisterInfo::requiresFrameIndexScavenging(
     const MachineFunction &MF) const {
   return true;
 }
 
 bool

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

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog -mattr=+sve %s -o - | FileCheck %s
+---
+# This test verifies that the basepointer is available in presence of SVE stack objects.
+name: hasBasepointer
+# CHECK-LABEL: name: hasBasepointer
+# CHECK: bb.0:
+# CHECK:      $sp = frame-setup ADDVL_XXI $sp, -1
+# CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0
+# CHECK-NEXT: $x19 = ADDXri $sp, 0, 0
+# CHECK:      STRXui $x0, $x19, 0
+tracksRegLiveness: true
+frameInfo:
+  isFrameAddressTaken: true
+stack:
+  - { id: 0, type: variable-sized,  alignment: 1 }
+  - { id: 1, name: '', size: 16, alignment: 8 }
+  - { id: 2, stack-id: sve-vec, size: 16, alignment: 16 }
+body: |
+  bb.0:
+    liveins: $x0
+    STRXui $x0, %stack.1, 0
+    RET_ReallyLR
+...

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

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog -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
+# CHECK:      STRXui killed $[[SCRATCH:x[0-9]+]], $sp, 1
+# CHECK-NEXT: $[[SCRATCH]] = ADDVL_XXI $fp, -1
+# CHECK-NEXT: STRXui $x0, killed $[[SCRATCH]], 0
+# CHECK: bb.1:
+tracksRegLiveness: true
+frameInfo:
+  isFrameAddressTaken: true
+stack:
+  - { id: 0, name: '', size: 32761, alignment: 8 }
+  - { id: 1, stack-id: sve-vec, size: 16, alignment: 16 }
+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
+...

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

-# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog %s -o - | FileCheck %s
+# RUN: llc -mattr=+sve -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog %s -o - | FileCheck %s
 #
 # Test allocation and deallocation of SVE objects on the stack,
 # as well as using a combination of scalable and non-scalable
 # offsets to access the SVE on the stack.
 #
 # SVE objects are allocated below the (scalar) callee saves,
 # and above spills/locals and the alignment gap, e.g.
 #
--- |
 
   define void @test_allocate_sve() nounwind { entry: unreachable }
   define void @test_allocate_sve_gpr_callee_saves() nounwind { entry: unreachable }
   define void @test_allocate_sve_gpr_realigned() nounwind { entry: unreachable }
   define void @test_address_sve() nounwind { entry: unreachable }
   define void @test_address_sve_fp() nounwind { entry: unreachable }
   define void @test_stack_arg_sve() nounwind { entry: unreachable }
   define void @test_address_sve_out_of_range() nounwind { entry: unreachable }
-  define void @test_address_gpr_vla_nobp() nounwind { entry: unreachable }
+  define void @test_address_gpr_vla() nounwind { entry: unreachable }
   define aarch64_sve_vector_pcs void @save_restore_pregs_sve() nounwind { entry: unreachable }
   define aarch64_sve_vector_pcs void @save_restore_zregs_sve() nounwind { entry: unreachable }
   define aarch64_sve_vector_pcs void @save_restore_sve() nounwind { entry: unreachable }
   define aarch64_sve_vector_pcs void @save_restore_sve_realign() nounwind { entry: unreachable }
   define aarch64_sve_vector_pcs void @frame_layout() nounwind { entry: unreachable }
 
 ...
 # +----------+
bb.0.entry:
     liveins: $z0, $p0
 
     STR_ZXI $z0, %fixed-stack.0, 0
     STR_PXI $p0, %fixed-stack.1, 0
 
     RET_ReallyLR
 ---
 ...
-# Test that non-SVE objects are accessed from FP when there is no BP,
-# but the SP cannot be used because of variable-length arrays.
+# Test that non-SVE objects are accessed from BP when there are
+# variable length arrays, because it will be more expensive to
+# access from the FP when there are also SVE objects on the stack.
 #
 # +----------+ <- FP
 # | %fstack.0|  // 16 scalable bytes
 # +----------+ <- @FP - 16 scalable bytes
 # | %stack.0 |  // 16 bytes
-# +----------+ <- @FP - 16 scalable bytes - 16b
+# +----------+ <- @BP
 # : %stack.1 :  // variable length
 # +----------+ <- SP
 
-# CHECK-LABEL: name: test_address_gpr_vla_nobp
-# CHECK:      bb.0.entry:
-# CHECK:      $[[TMP:x[0-9]+]] = ADDVL_XXI $fp, -1
-# CHECK-NEXT: STURXi $xzr, killed $[[TMP]], -16
-# CHECK:      RET_ReallyLR
-name:            test_address_gpr_vla_nobp
+# CHECK-LABEL: name: test_address_gpr_vla
+# CHECK: bb.0.entry:
+# CHECK: STRXui $xzr, $x19, 0
+# CHECK: RET_ReallyLR
+name:            test_address_gpr_vla
 frameInfo:
   maxAlignment:  16
 fixedStack:
   - { id: 0, stack-id: sve-vec, size: 16, alignment: 8, offset: -16 }
 stack:
   - { id: 0, stack-id: default, size: 16, alignment: 8 }
   - { id: 1, stack-id: default, type: variable-sized }
 body:             |