[AArch64] Stack frame reordering.

[eugenis]

Summary

Implement stack frame reordering in the AArch64 backend.

Unlike the X86 implementation, AArch64 does not seem to benefit from
"access density" based frame reordering, mainly because it has a much
smaller variety of addressing modes, and the fact that all instructions
are 4 bytes so each frame object is either in range of an instruction
(and then the access is "free") or not (and that has a code size cost
of 4 bytes).

This change improves Memory Tagging codegen by

• Placing an object that has been chosen as the base tagged pointer of

the function at SP + 0. This saves one instruction to setup the pointer
(IRG does not have an offset immediate), and more because that object
can now be referenced without materializing its tagged address in a
scratch register.

• Placing objects that go out of scope simultaneously together. This

exposes opportunities for instruction merging in tryMergeAdjacentSTG.

[eugenis]

This revision and its two dependencies improve MTE code size of chromium on aarch64-linux-android by slightly more than 1%.

[merge_guards_bot]

Unit tests: fail. 61307 tests passed, 1 failed and 736 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_sharedtimedmutex_requirements/thread_sharedtimedmutex_class/try_lock.pass.cpp

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[eugenis]

rebase

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[eugenis]

ping

[eugenis]

ping

[danielkiss]

Otherwise I like the idea.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3113

IMHO if the data structure would be vector of vectors of FrameObjects then maybe the whole logic would be a bit simpler, but maybe I missed something.

[efriedma]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3175

It's usually more concise to write this sort of thing with std::tuple: return std::make_tuple(!A.IsValid, A.ObjectFirst, A.GroupFirst, A.GroupIndex, A.ObjectIndex) < std::make_tuple(!B.IsValid, B.ObjectFirst, B.GroupFirst, B.GroupIndex, B.ObjectIndex).

3181

Can we skip running this code when stack tagging is disabled? It doesn't do anything in that case, I think, but it would avoid running a bunch of unnecessary code, and make the intent more clear.

[eugenis]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3113

Do you mean a list of groups, where each group is a list of objects?
In that case we would have to deal with the possibility that one object participates in multiple groups, and resolve that somehow when building the final frame order. In the current implementation, the object always abandons its old group when joining a new one - it's not optimal, but simple and seems to work fine.

Also, FrameObjectCompare needs a fast way to determine when two objects are in the same group, hence group id is a member of FrameObject, and not vice versa.

3181

Sure, we can look for the sanitize_memtag attribute. Stack tagging instructions can be added some other way of course (ex. with intrinsic calls in the C source).
We could check if MTE is in the subtarget's feature bits, but that would not help much.

[danielkiss]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3113

Do you mean a list of groups, where each group is a list of objects?

Yes, like that.

In that case we would have to deal with the possibility that one object participates in multiple groups

Right, that I missed :)

[eugenis]

address comments

[eugenis]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3181

I'm not sure how to do this. As settag-merge-order.ll test shows, it's pretty easy to generate code that benefits from this transformation w/o a sanitize_memtag attribute.
I could fold this check into some other place where we scan all instructions in a function, like the StackTaggingPreRA pass, and store in AArch64FunctionInfo, but that looks like a premature optimization to me.

[eugenis]

rebase

[eugenis]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3181

I've measured compilation time (pdfium_test in chromium without tagging, default config vs -aarch64-order-frame-objects=0) and did not see any change due to this patch.

[eugenis]

Could anyone take another look at this and https://reviews.llvm.org/D72365 ?

[efriedma]

this and https://reviews.llvm.org/D72365 ?

The two patches seem to overlap?

[eugenis]

rebase

[eugenis]

rebase

[eugenis]

In D72366#2322650, @efriedma wrote:

this and https://reviews.llvm.org/D72365 ?

The two patches seem to overlap?

Indeed. Arc is full of surprises.

I've reuploaded the patches separately.

[efriedma]

LGTM

[danielkiss]

LGTM

[nickdesaulniers]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3254–3255

What happened with the formatting here?

[eugenis]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3254–3255

clang-format happened :)
Does it change if you reformat it?

I can move it to a helper function.

llvm/lib/Target/AArch64/AArch64FrameLowering.h

default:
       return false;
     case TargetStackID::Default:
     case TargetStackID::SVEVector:
     case TargetStackID::NoAlloc:
       return true;
     }
   }
 
+  void
+  orderFrameObjects(const MachineFunction &MF,
+                    SmallVectorImpl<int> &ObjectsToAllocate) const override;
+
 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

ReverseCSRRestoreSeq("reverse-csr-restore-seq",
                          cl::desc("reverse the CSR restore sequence"),
                          cl::init(false), cl::Hidden);
 
 static cl::opt<bool> StackTaggingMergeSetTag(
     "stack-tagging-merge-settag",
     cl::desc("merge settag instruction in function epilog"), cl::init(true),
     cl::Hidden);
 
+static cl::opt<bool> OrderFrameObjects("aarch64-order-frame-objects",
+                                   cl::desc("sort stack allocations"),
+                                   cl::init(true), cl::Hidden);
+
 STATISTIC(NumRedZoneFunctions, "Number of functions using red zone");
 
 /// This is the biggest offset to the stack pointer we can encode in aarch64
 /// instructions (without using a separate calculation and a temp register).
 /// Note that the exception here are vector stores/loads which cannot encode any
 /// displacements (see estimateRSStackSizeLimit(), isAArch64FrameOffsetLegal()).
 static const unsigned DefaultSafeSPDisplacement = 255;
 
bool AArch64FrameLowering::shouldCombineCSRLocalStackBumpInEpilogue(
   default:
     return true;
   }
   llvm_unreachable("unreachable");
 }
 
 // Given a load or a store instruction, generate an appropriate unwinding SEH
 // code on Windows.
 static MachineBasicBlock::iterator InsertSEH(MachineBasicBlock::iterator MBBI,
                                              const TargetInstrInfo &TII,
                                              MachineInstr::MIFlag Flag) {
   unsigned Opc = MBBI->getOpcode();
   MachineBasicBlock *MBB = MBBI->getParent();
   MachineFunction &MF = *MBB->getParent();
   DebugLoc DL = MBBI->getDebugLoc();
   unsigned ImmIdx = MBBI->getNumOperands() - 1;
   int Imm = MBBI->getOperand(ImmIdx).getImm();
public:
   void emitCode(MachineBasicBlock::iterator &InsertI,
                 const AArch64FrameLowering *TFI, bool IsLast);
 };
 
 void TagStoreEdit::emitUnrolled(MachineBasicBlock::iterator InsertI) {
   const AArch64InstrInfo *TII =
       MF->getSubtarget<AArch64Subtarget>().getInstrInfo();
 
   const int64_t kMinOffset = -256 * 16;
   const int64_t kMaxOffset = 255 * 16;
 
   Register BaseReg = FrameReg;
   int64_t BaseRegOffsetBytes = FrameRegOffset.getBytes();
   if (BaseRegOffsetBytes < kMinOffset ||
       BaseRegOffsetBytes + (Size - Size % 32) > kMaxOffset) {
     Register ScratchReg = MRI->createVirtualRegister(&AArch64::GPR64RegClass);
     emitFrameOffset(*MBB, InsertI, DL, ScratchReg, BaseReg,
                     {BaseRegOffsetBytes, MVT::i8}, TII);
bool canMergeRegUpdate(MachineBasicBlock::iterator II, unsigned Reg,
   if ((MI.getOpcode() == AArch64::ADDXri ||
        MI.getOpcode() == AArch64::SUBXri) &&
       MI.getOperand(0).getReg() == Reg && MI.getOperand(1).getReg() == Reg) {
     unsigned Shift = AArch64_AM::getShiftValue(MI.getOperand(3).getImm());
     int64_t Offset = MI.getOperand(2).getImm() << Shift;
     if (MI.getOpcode() == AArch64::SUBXri)
       Offset = -Offset;
     int64_t AbsPostOffset = std::abs(Offset - Size);
     const int64_t kMaxOffset =
         0xFFF; // Max encoding for unshifted ADDXri / SUBXri
     if (AbsPostOffset <= kMaxOffset && AbsPostOffset % 16 == 0) {
       *TotalOffset = Offset;
       return true;
     }
   }
   return false;
 }
void TagStoreEdit::emitCode(MachineBasicBlock::iterator &InsertI,
   mergeMemRefs(TagStores, CombinedMemRefs);
 
   LLVM_DEBUG(dbgs() << "Replacing adjacent STG instructions:\n";
              for (const auto &Instr
                   : TagStores) { dbgs() << "  " << *Instr.MI; });
 
   // Size threshold where a loop becomes shorter than a linear sequence of
   // tagging instructions.
   const int kSetTagLoopThreshold = 176;
   if (Size < kSetTagLoopThreshold) {
     if (TagStores.size() < 2)
       return;
     emitUnrolled(InsertI);
   } else {
     MachineInstr *UpdateInstr = nullptr;
     int64_t TotalOffset;
     if (IsLast) {
MachineBasicBlock::iterator tryMergeAdjacentSTG(MachineBasicBlock::iterator II,
   if (&MI == &MBB->instr_back())
     return II;
   if (!isMergeableStackTaggingInstruction(MI, Offset, Size, FirstZeroData))
     return II;
 
   SmallVector<TagStoreInstr, 4> Instrs;
   Instrs.emplace_back(&MI, Offset, Size);
 
   constexpr int kScanLimit = 10;
   int Count = 0;
   for (MachineBasicBlock::iterator E = MBB->end();
        NextI != E && Count < kScanLimit; ++NextI) {
     MachineInstr &MI = *NextI;
     bool ZeroData;
     int64_t Size, Offset;
     // Collect instructions that update memory tags with a FrameIndex operand
     // and (when applicable) constant size, and whose output registers are dead
unsigned AArch64FrameLowering::getWinEHFuncletFrameSize(
     const MachineFunction &MF) const {
   // This is the size of the pushed CSRs.
   unsigned CSSize =
       MF.getInfo<AArch64FunctionInfo>()->getCalleeSavedStackSize();
   // This is the amount of stack a funclet needs to allocate.
   return alignTo(CSSize + MF.getFrameInfo().getMaxCallFrameSize(),
                  getStackAlignment());
 }
+
+namespace {
+struct FrameObject {
+  bool IsValid = false;
+  // Index of the object in MFI.
+  int ObjectIndex = 0;
+  // Group ID this object belongs to.
+  int GroupIndex = -1;
+  // This object should be placed first (closest to SP).
+  bool ObjectFirst = false;
+  // This object's group (which always contains the object with
+  // ObjectFirst==true) should be placed first.
+  bool GroupFirst = false;
+};
+
+class GroupBuilder {
+  SmallVector<int, 8> CurrentMembers;
+  int NextGroupIndex = 0;
+  std::vector<FrameObject> &Objects;
+
+public:
+  GroupBuilder(std::vector<FrameObject> &Objects) : Objects(Objects) {}
+  void AddMember(int Index) { CurrentMembers.push_back(Index); }
+  void EndCurrentGroup() {
+    if (CurrentMembers.size() > 1) {
+      // Create a new group with the current member list. This might remove them
+      // from their pre-existing groups. That's OK, dealing with overlapping
+      // groups is too hard and unlikely to make a difference.
+      LLVM_DEBUG(dbgs() << "group:");
+      for (int Index : CurrentMembers) {
+        Objects[Index].GroupIndex = NextGroupIndex;
+        LLVM_DEBUG(dbgs() << " " << Index);
+      }
+      LLVM_DEBUG(dbgs() << "\n");
+      NextGroupIndex++;
+    }
+    CurrentMembers.clear();
+  }
+};
+
+bool FrameObjectCompare(const FrameObject &A, const FrameObject &B) {
+  // For consistency in our comparison, all invalid objects are placed
+  // at the end. This also allows us to stop walking when we hit the
+  // first invalid item after it's all sorted.
+  if (!A.IsValid)
+    return false;
+  if (!B.IsValid)
+    return true;
+
+  // Objects at a lower index are closer to FP; objects at a higher index are
+  // closer to SP.
+
+  // The "first" object goes first (closest to SP).
+  if (A.ObjectFirst != B.ObjectFirst)
+    return B.ObjectFirst;
+
+  // The "first" group members follow.
+  if (A.GroupFirst != B.GroupFirst)
+    return B.GroupFirst;
+
+  // Order the rest by the group ID first to keep the groups together.
+  // Higher numbered groups are more likely to be around longer (i.e. untagged
+  // in the function epilogue and not at some earlier point). Place them closer
+  // to SP.
+  if (A.GroupIndex != B.GroupIndex)
+    return A.GroupIndex < B.GroupIndex;
+
+  // Otherwise, keep the objects in the original order.
+  return A.ObjectIndex < B.ObjectIndex;
+}
+} // namespace
+
+void AArch64FrameLowering::orderFrameObjects(
+    const MachineFunction &MF, SmallVectorImpl<int> &ObjectsToAllocate) const {
+  if (!OrderFrameObjects || ObjectsToAllocate.empty())
+    return;
+
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  std::vector<FrameObject> FrameObjects(MFI.getObjectIndexEnd());
+  for (auto &Obj : ObjectsToAllocate) {
+    FrameObjects[Obj].IsValid = true;
+    FrameObjects[Obj].ObjectIndex = Obj;
+  }
+
+  // Identify stack slots that are tagged at the same time.
+  GroupBuilder GB(FrameObjects);
+  for (auto &MBB : MF) {
+    for (auto &MI : MBB) {
+      if (MI.isDebugInstr())
+        continue;
+      int OpIndex;
+      switch (MI.getOpcode()) {
+      case AArch64::STGloop:
+      case AArch64::STZGloop:
+        OpIndex = 3;
+        break;
+      case AArch64::STGOffset:
+      case AArch64::STZGOffset:
+      case AArch64::ST2GOffset:
+      case AArch64::STZ2GOffset:
+        OpIndex = 1;
+        break;
+      default:
+        OpIndex = -1;
+      }
+
+      int TaggedFI = -1;
+      if (OpIndex >= 0) {
+        const MachineOperand &MO = MI.getOperand(OpIndex);
+        if (MO.isFI()) {
+          int FI = MO.getIndex();
+          if (FI >= 0 && FI < MFI.getObjectIndexEnd() &&
+              FrameObjects[FI].IsValid)
+            TaggedFI = FI;
+        }
+      }
+
+      // If this is a stack tagging instruction for a slot that is not part of a
+      // group yet, either start a new group or add it to the current one.
+      if (TaggedFI >= 0)
+        GB.AddMember(TaggedFI);
+      else
+        GB.EndCurrentGroup();
+    }
+    // Groups should never span multiple basic blocks.
+    GB.EndCurrentGroup();
+  }
+
+  // If the function's tagged base pointer is pinned to a stack slot, we want to
+  // put that slot first when possible. This will likely place it at SP + 0,
+  // and save one instruction when generating the base pointer because IRG does
+  // not allow an immediate offset.
+  const AArch64FunctionInfo &AFI = *MF.getInfo<AArch64FunctionInfo>();
+  Optional<int> TBPI = AFI.getTaggedBasePointerIndex();
+  if (TBPI) {
+    FrameObjects[*TBPI].ObjectFirst = true;
+    FrameObjects[*TBPI].GroupFirst = true;
+    int FirstGroupIndex = FrameObjects[*TBPI].GroupIndex;
+    if (FirstGroupIndex >= 0)
+      for (FrameObject &Object : FrameObjects)
+        if (Object.GroupIndex == FirstGroupIndex)
+          Object.GroupFirst = true;
+  }
+
+  llvm::stable_sort(FrameObjects, FrameObjectCompare);
+
+  int i = 0;
+  for (auto &Obj : FrameObjects) {
+    // All invalid items are sorted at the end, so it's safe to stop.
+    if (!Obj.IsValid)
+      break;
+    ObjectsToAllocate[i++] = Obj.ObjectIndex;
+  }
+
+  LLVM_DEBUG(dbgs() << "Final frame order:\n"; for (auto &Obj
+                                                    : FrameObjects) {
+    if (!Obj.IsValid)
+      break;
+    dbgs() << "  " << Obj.ObjectIndex << ": group " << Obj.GroupIndex;
+    if (Obj.ObjectFirst)
+      dbgs() << ", first";
+    if (Obj.GroupFirst)
+      dbgs() << ", group-first";
+    dbgs() << "\n";
+  });
+}

llvm/test/CodeGen/AArch64/settag-merge-order.ll

+; RUN: llc < %s -mtriple=aarch64 -mattr=+mte -aarch64-order-frame-objects=1 | FileCheck %s
+
+declare void @use(i8* %p)
+declare void @llvm.aarch64.settag(i8* %p, i64 %a)
+declare void @llvm.aarch64.settag.zero(i8* %p, i64 %a)
+
+; Two loops of size 256; the second loop updates SP.
+; After frame reordering, two loops can be merged into one.
+define void @stg128_128_gap_128_128() {
+entry:
+; CHECK-LABEL: stg128_128_gap_128_128:
+; CHECK: mov     x8, #512
+; CHECK: st2g    sp, [sp], #32
+; CHECK: sub     x8, x8, #32
+; CHECK: cbnz    x8,
+; CHECK: ret
+  %a = alloca i8, i32 128, align 16
+  %a2 = alloca i8, i32 128, align 16
+  %b = alloca i8, i32 32, align 16
+  %c = alloca i8, i32 128, align 16
+  %c2 = alloca i8, i32 128, align 16
+  call void @use(i8* %b)
+  call void @llvm.aarch64.settag(i8* %a, i64 128)
+  call void @llvm.aarch64.settag(i8* %a2, i64 128)
+  call void @llvm.aarch64.settag(i8* %c, i64 128)
+  call void @llvm.aarch64.settag(i8* %c2, i64 128)
+  ret void
+}
+
+define void @stg2(i1 %flag) {
+entry:
+; CHECK-LABEL: stg2:
+  %a = alloca i8, i32 160, align 16
+  %a2 = alloca i8, i32 160, align 16
+  %b = alloca i8, i32 32, align 16
+  %c = alloca i8, i32 128, align 16
+  %c2 = alloca i8, i32 128, align 16
+  call void @use(i8* %b)
+  br i1 %flag, label %if.then, label %if.else
+
+if.then:
+; CHECK: mov     x8, #320
+; CHECK: st2g    x9, [x9], #32
+; CHECK: sub     x8, x8, #32
+; CHECK: cbnz    x8,
+  call void @llvm.aarch64.settag(i8* %a, i64 160)
+  call void @llvm.aarch64.settag(i8* %a2, i64 160)
+  br label %if.end
+
+if.else:
+; CHECK: mov     x8, #256
+; CHECK: st2g    x9, [x9], #32
+; CHECK: sub     x8, x8, #32
+; CHECK: cbnz    x8,
+  call void @llvm.aarch64.settag(i8* %c, i64 128)
+  call void @llvm.aarch64.settag(i8* %c2, i64 128)
+  br label %if.end
+
+if.end:
+; CHECK: mov     x8, #576
+; CHECK: st2g    sp, [sp], #32
+; CHECK: sub     x8, x8, #32
+; CHECK: cbnz    x8,
+  call void @llvm.aarch64.settag(i8* %a, i64 160)
+  call void @llvm.aarch64.settag(i8* %a2, i64 160)
+  call void @llvm.aarch64.settag(i8* %c, i64 128)
+  call void @llvm.aarch64.settag(i8* %c2, i64 128)
+
+; CHECK: ret
+  ret void
+}

llvm/test/CodeGen/AArch64/settag-merge.ll

-; RUN: llc < %s -mtriple=aarch64 -mattr=+mte | FileCheck %s
+; RUN: llc < %s -mtriple=aarch64 -mattr=+mte -aarch64-order-frame-objects=0 | FileCheck %s
 
 declare void @use(i8* %p)
 declare void @llvm.aarch64.settag(i8* %p, i64 %a)
 declare void @llvm.aarch64.settag.zero(i8* %p, i64 %a)
 
 define void @stg16_16() {
 entry:
 ; CHECK-LABEL: stg16_16: