[SVE] Don't use LocalStackAllocation for SVE objects

[david-arm]

Summary

[SVE] Don't use LocalStackAllocation for SVE objects

I have introduced a new TargetFrameLowering query function:

isStackIdSafeForLocalArea

that queries whether or not it is safe for objects of a given stack
id to be bundled into the local area. The default behaviour is to
always bundle regardless of the stack id, however for AArch64 this is
overriden so that it's only safe for fixed-size stack objects.
There is future work here to extend this algorithm for multiple local
areas so that SVE stack objects can be bundled together and accessed
from their own virtual base-pointer.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64FrameLowering.h

108

To me the name implies something that is safe either way, whereas you're making the change to fix a bug. Would isStackIdSafeForLocalArea be a better fit?

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

2

I forget the exact requirement but does the use of -debug-only mean this tests needs

REQUIRES: Asserts

or something. It is not possible to test the expected frame/stack layout instead, rather the rely on debug?

[paulwalker-arm]

When generating the test mir file did you use -simplify-mir. If not then I recommend regenerating it to remove redundant fields and thus have a simpler test input.

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

2

You can use -run-only= here instead of the start/stop combo.

[paulwalker-arm]

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

19

Perhaps worth adding a CHECK for the name line, just in case somebody adds additional tests to this file.

llvm/include/llvm/CodeGen/TargetFrameLowering.h

bool isStackRealignable() const {
     return StackRealignable;
   }
 
   /// Return the skew that has to be applied to stack alignment under
   /// certain conditions (e.g. stack was adjusted before function \p MF
   /// was called).
   virtual unsigned getStackAlignmentSkew(const MachineFunction &MF) const;
 
+  /// This method returns whether or not it is safe for an object with the
+  /// given stack id to be bundled into the local area.
+  virtual bool isStackIdSafeForLocalArea(unsigned StackId) const {
+    return true;
+  }
+
   /// getOffsetOfLocalArea - This method returns the offset of the local area
   /// from the stack pointer on entrance to a function.
   ///
   int getOffsetOfLocalArea() const { return LocalAreaOffset; }
 
   /// isFPCloseToIncomingSP - Return true if the frame pointer is close to
   /// the incoming stack pointer, false if it is close to the post-prologue
   /// stack pointer.

llvm/lib/CodeGen/LocalStackSlotAllocation.cpp

if (MFI.hasStackProtectorIndex()) {
     AdjustStackOffset(MFI, StackProtectorFI, Offset, StackGrowsDown, MaxAlign);
 
     // Assign large stack objects first.
     for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
       if (MFI.isDeadObjectIndex(i))
         continue;
       if (StackProtectorFI == (int)i)
         continue;
+      if (!TFI.isStackIdSafeForLocalArea(MFI.getStackID(i)))
+        continue;
 
       switch (MFI.getObjectSSPLayout(i)) {
       case MachineFrameInfo::SSPLK_None:
         continue;
       case MachineFrameInfo::SSPLK_SmallArray:
         SmallArrayObjs.insert(i);
         continue;
       case MachineFrameInfo::SSPLK_AddrOf:
void LocalStackSlotPass::calculateFrameObjectOffsets(MachineFunction &Fn) {
   // callee saved registers.
   for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
     if (MFI.isDeadObjectIndex(i))
       continue;
     if (MFI.getStackProtectorIndex() == (int)i)
       continue;
     if (ProtectedObjs.count(i))
       continue;
+    if (!TFI.isStackIdSafeForLocalArea(MFI.getStackID(i)))
+      continue;
 
     AdjustStackOffset(MFI, i, Offset, StackGrowsDown, MaxAlign);
   }
 
   // Remember how big this blob of stack space is
   MFI.setLocalFrameSize(Offset);
   MFI.setLocalFrameMaxAlign(MaxAlign);
 }

llvm/lib/Target/AArch64/AArch64FrameLowering.h

default:
       return false;
     case TargetStackID::Default:
     case TargetStackID::SVEVector:
     case TargetStackID::NoAlloc:
       return true;
     }
   }
 
+  bool isStackIdSafeForLocalArea(unsigned StackId) const override {
+    // We don't support putting SVE objects into the pre-allocated local
+    // frame block at the moment.
+    return StackId != TargetStackID::SVEVector;
+  }
+
 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/test/CodeGen/AArch64/sve-localstackalloc.mir

+# RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve -run-pass=localstackalloc -o - %s | FileCheck %s
+
+--- |
+  ; ModuleID = '<stdin>'
+  source_filename = "<stdin>"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-unknown-linux-gnu"
+
+  define <vscale x 32 x i8> @insert_32i8_idx(<vscale x 32 x i8> %a, i8 %elt, i64 %idx) #0 {
+    %ins = insertelement <vscale x 32 x i8> %a, i8 %elt, i64 %idx
+    ret <vscale x 32 x i8> %ins
+  }
+
+  attributes #0 = { "target-features"="+sve" }
+
+...
+---
+name:            insert_32i8_idx
+alignment:       4
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: zpr, preferred-register: '' }
+  - { id: 1, class: zpr, preferred-register: '' }
+  - { id: 2, class: gpr32, preferred-register: '' }
+  - { id: 3, class: gpr64, preferred-register: '' }
+  - { id: 5, class: ppr_3b, preferred-register: '' }
+  - { id: 6, class: gpr64sp, preferred-register: '' }
+  - { id: 7, class: zpr, preferred-register: '' }
+  - { id: 8, class: zpr, preferred-register: '' }
+liveins:
+  - { reg: '$z0', virtual-reg: '%0' }
+  - { reg: '$z1', virtual-reg: '%1' }
+  - { reg: '$w0', virtual-reg: '%2' }
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+# CHECK-LABEL: name: insert_32i8_idx
+# CHECK: localFrameSize:  0
+stack:
+  - { id: 0, name: '', type: default, offset: 0, size: 32, alignment: 16,
+      stack-id: sve-vec, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+machineFunctionInfo: {}
+body:             |
+  bb.0 (%ir-block.0):
+    liveins: $z0, $z1, $w0
+
+    %2:gpr32 = COPY $w0
+    %1:zpr = COPY $z1
+    %0:zpr = COPY $z0
+    %5:ppr_3b = PTRUE_B 31
+    %6:gpr64sp = ADDXri %stack.0, 0, 0
+    ST1B_IMM %1, %5, %6, 1 :: (store unknown-size, align 16)
+    ST1B_IMM %0, %5, %stack.0, 0 :: (store unknown-size into %stack.0, align 16)
+    %7:zpr = LD1B_IMM %5, %6, 1 :: (load unknown-size from %stack.0 + 16, align 16)
+    %8:zpr = LD1B_IMM %5, %stack.0, 0 :: (load unknown-size from %stack.0, align 16)
+    $z0 = COPY %8
+    $z1 = COPY %7
+    RET_ReallyLR implicit $z0, implicit $z1
+
+...