[MTE] Pin the tagged base pointer to one of the stack slots.

[eugenis]

Summary

Pin the tagged base pointer to one of the stack slots, and (if
necessary) rewrite tag offsets so that an object that occupies that
slot has both address and tag offsets of 0. This allows ADDG
instructions for that object to be eliminated and their uses replaced
with the tagged base pointer itself.

This optimization must be done in machine instructions and not in the IR
instrumentation pass, because referring to a stack slot through an IRG
pointer would confuse the stack coloring pass.

The optimization makes a (pretty naive) attempt to find the slot that
would benefit the most by counting the uses of stack slots in the
function.

[merge_guards_bot]

Unit tests: pass. 61305 tests passed, 0 failed and 736 were skipped.

clang-tidy: pass.

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: pass. 62114 tests passed, 0 failed and 808 were skipped.

clang-tidy: pass.

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]

ping

[eugenis]

ping

[rankov]

llvm/lib/Target/AArch64/AArch64StackTaggingPreRA.cpp

256

Size of the slot might be important because of the range of addg. It can only reach 2^6 tag granules (64 * 16B) from base register. If you select a large slot it will reduce the number of slots that can be reached with addg and will result in additional add instructions.

Also, it might be beneficial to sort smaller slots near the SP.

llvm/test/CodeGen/AArch64/irg_sp_tagp.ll

53

There is no addg anymore.

60

Why not check that there is no more addg?

69

No more addg.

73

This might not be guaranteed to be x1

76

Check that addg is not there?

94

Might not be x1 here.

95

Check that there is no addg.

[eugenis]

llvm/lib/Target/AArch64/AArch64StackTaggingPreRA.cpp

256

I've tried sorting by size and/or access density (size / number of references) in D72366, but could not show any significant improvement. There are conflicting requirements there between minimizing the average ld/st immediate offset, maximizing the number of references to the object with tag 0 & offset 0, and keeping the objects that are untagged at the same time together (for STG merging).

I may return to this problem at some later time.

[eugenis]

address comments

[eugenis]

rebase

[eugenis]

rebase

[danielkiss]

LGTM, just please fix the lint issues in the llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

[eugenis]

Sorry, I don't see any lint issues. Could you be more specific?

[danielkiss]

Those are gone from the view, so ignore what I said.

llvm/include/llvm/IR/IntrinsicsAArch64.td

 // Generate a randomly tagged stack base pointer.
 def int_aarch64_irg_sp   : Intrinsic<[llvm_ptr_ty], [llvm_i64_ty],
     [IntrNoMem, IntrHasSideEffects]>;
 
 // Transfer pointer tag with offset.
 // ptr1 = tagp(ptr0, baseptr, tag_offset) returns a pointer where
 // * address is the address in ptr0
 // * tag is a function of (tag in baseptr, tag_offset).
+// ** Beware, this is not the same function as implemented by the ADDG instruction!
+//    Backend optimizations may change tag_offset; the only guarantee is that calls
+//    to tagp with the same pair of (baseptr, tag_offset) will produce pointers
+//    with the same tag value, assuming the set of excluded tags has not changed.
 // Address bits in baseptr and tag bits in ptr0 are ignored.
 // When offset between ptr0 and baseptr is a compile time constant, this can be emitted as
 //   ADDG ptr1, baseptr, (ptr0 - baseptr), tag_offset
 // It is intended that ptr0 is an alloca address, and baseptr is the direct output of llvm.aarch64.irg.sp.
 def int_aarch64_tagp : Intrinsic<[llvm_anyptr_ty], [LLVMMatchType<0>, llvm_ptr_ty, llvm_i64_ty],
     [IntrNoMem, ImmArg<2>]>;
 
 // Update allocation tags for the memory range to match the tag in the pointer argument.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
         .setMIFlags(MachineInstr::FrameSetup);
   }
 
   // All calls are tail calls in GHC calling conv, and functions have no
   // prologue/epilogue.
   if (MF.getFunction().getCallingConv() == CallingConv::GHC)
     return;
 
-  // Set tagged base pointer to the bottom of the stack frame.
+  // Set tagged base pointer to the requested stack slot.
   // Ideally it should match SP value after prologue.
+  Optional<int> TBPI = AFI->getTaggedBasePointerIndex();
+  if (TBPI)
+    AFI->setTaggedBasePointerOffset(-MFI.getObjectOffset(*TBPI));
+  else
     AFI->setTaggedBasePointerOffset(MFI.getStackSize());
 
   const StackOffset &SVEStackSize = getSVEStackSize(MF);
 
   // getStackSize() includes all the locals in its size calculation. We don't
   // include these locals when computing the stack size of a funclet, as they
   // are allocated in the parent's stack frame and accessed via the frame
   // pointer from the funclet.  We only save the callee saved registers in the
   // funclet, which are really the callee saved registers of the parent

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// Initialized during frame lowering, unless the function has the noredzone
   /// attribute, in which case it is set to false at construction.
   Optional<bool> HasRedZone;
 
   /// ForwardedMustTailRegParms - A list of virtual and physical registers
   /// that must be forwarded to every musttail call.
   SmallVector<ForwardedRegister, 1> ForwardedMustTailRegParms;
 
-  // Offset from SP-at-entry to the tagged base pointer.
-  // Tagged base pointer is set up to point to the first (lowest address) tagged
-  // stack slot.
-  unsigned TaggedBasePointerOffset = 0;
+  /// FrameIndex for the tagged base pointer.
+  Optional<int> TaggedBasePointerIndex;
+
+  /// Offset from SP-at-entry to the tagged base pointer.
+  /// Tagged base pointer is set up to point to the first (lowest address) tagged
+  /// stack slot.
+  unsigned TaggedBasePointerOffset;
 
 public:
   AArch64FunctionInfo() = default;
 
   explicit AArch64FunctionInfo(MachineFunction &MF) {
     (void)MF;
 
     // If we already know that the function doesn't have a redzone, set
void addLOHDirective(MCLOHType Kind, MILOHArgs Args) {
     LOHContainerSet.push_back(MILOHDirective(Kind, Args));
     LOHRelated.insert(Args.begin(), Args.end());
   }
 
   SmallVectorImpl<ForwardedRegister> &getForwardedMustTailRegParms() {
     return ForwardedMustTailRegParms;
   }
 
+  Optional<int> getTaggedBasePointerIndex() const {
+    return TaggedBasePointerIndex;
+  }
+  void setTaggedBasePointerIndex(int Index) {
+    TaggedBasePointerIndex = Index;
+  }
+
   unsigned getTaggedBasePointerOffset() const {
     return TaggedBasePointerOffset;
   }
   void setTaggedBasePointerOffset(unsigned Offset) {
     TaggedBasePointerOffset = Offset;
   }
 
 private:

llvm/lib/Target/AArch64/AArch64StackTaggingPreRA.cpp

 //===-- AArch64StackTaggingPreRA.cpp --- Stack Tagging for AArch64 -----===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 
 #include "AArch64.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64InstrInfo.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SetVector.h"
-#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
cl::desc(
         "stack frames, or in the presence of variable sized allocas)."),
     cl::values(
         clEnumValN(UncheckedNever, "never", "never apply unchecked-ld-st"),
         clEnumValN(
             UncheckedSafe, "safe",
             "apply unchecked-ld-st when the target is definitely within range"),
         clEnumValN(UncheckedAlways, "always", "always apply unchecked-ld-st")));
 
+static cl::opt<bool>
+    ClFirstSlot("stack-tagging-first-slot-opt", cl::Hidden, cl::init(true),
+                cl::ZeroOrMore,
+                cl::desc("Apply first slot optimization for stack tagging "
+                         "(eliminate ADDG Rt, Rn, 0, 0)."));
+
 namespace {
 
 class AArch64StackTaggingPreRA : public MachineFunctionPass {
   MachineFunction *MF;
   AArch64FunctionInfo *AFI;
   MachineFrameInfo *MFI;
   MachineRegisterInfo *MRI;
   const AArch64RegisterInfo *TRI;
   const AArch64InstrInfo *TII;
 
   SmallVector<MachineInstr*, 16> ReTags;
 
 public:
   static char ID;
   AArch64StackTaggingPreRA() : MachineFunctionPass(ID) {
     initializeAArch64StackTaggingPreRAPass(*PassRegistry::getPassRegistry());
   }
 
   bool mayUseUncheckedLoadStore();
   void uncheckUsesOf(unsigned TaggedReg, int FI);
   void uncheckLoadsAndStores();
+  Optional<int> findFirstSlotCandidate();
 
   bool runOnMachineFunction(MachineFunction &Func) override;
   StringRef getPassName() const override {
     return "AArch64 Stack Tagging PreRA";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();

 void AArch64StackTaggingPreRA::uncheckLoadsAndStores() {
   for (auto *I : ReTags) {
     unsigned TaggedReg = I->getOperand(0).getReg();
     int FI = I->getOperand(1).getIndex();
     uncheckUsesOf(TaggedReg, FI);
   }
 }
 
+struct SlotWithTag {
+  int FI;
+  int Tag;
+  SlotWithTag(int FI, int Tag) : FI(FI), Tag(Tag) {}
+  explicit SlotWithTag(const MachineInstr &MI)
+      : FI(MI.getOperand(1).getIndex()), Tag(MI.getOperand(4).getImm()) {}
+  bool operator==(const SlotWithTag &Other) const {
+    return FI == Other.FI && Tag == Other.Tag;
+  }
+};
+
+namespace llvm {
+template <> struct DenseMapInfo<SlotWithTag> {
+  static inline SlotWithTag getEmptyKey() { return {-2, -2}; }
+  static inline SlotWithTag getTombstoneKey() { return {-3, -3}; }
+  static unsigned getHashValue(const SlotWithTag &V) {
+    return hash_combine(DenseMapInfo<int>::getHashValue(V.FI),
+                        DenseMapInfo<int>::getHashValue(V.Tag));
+  }
+  static bool isEqual(const SlotWithTag &A, const SlotWithTag &B) {
+    return A == B;
+  }
+};
+} // namespace llvm
+
+static bool isSlotPreAllocated(MachineFrameInfo *MFI, int FI) {
+  return MFI->getUseLocalStackAllocationBlock() &&
+         MFI->isObjectPreAllocated(FI);
+}
+
+// Pin one of the tagged slots to offset 0 from the tagged base pointer.
+// This would make its address available in a virtual register (IRG's def), as
+// opposed to requiring an ADDG instruction to materialize. This effectively
+// eliminates a vreg (by replacing it with direct uses of IRG, which is usually
+// live almost everywhere anyway), and therefore needs to happen before
+// regalloc.
+Optional<int> AArch64StackTaggingPreRA::findFirstSlotCandidate() {
+  // Find the best (FI, Tag) pair to pin to offset 0.
+  // Looking at the possible uses of a tagged address, the advantage of pinning
+  // is:
+  // - COPY to physical register.
+  //   Does not matter, this would trade a MOV instruction for an ADDG.
+  // - ST*G matter, but those mostly appear near the function prologue where all
+  //   the tagged addresses need to be materialized anyway; also, counting ST*G
+  //   uses would overweight large allocas that require more than one ST*G
+  //   instruction.
+  // - Load/Store instructions in the address operand do not require a tagged
+  //   pointer, so they also do not benefit. These operands have already been
+  //   eliminated (see uncheckLoadsAndStores) so all remaining load/store
+  //   instructions count.
+  // - Any other instruction may benefit from being pinned to offset 0.
+  LLVM_DEBUG(dbgs() << "AArch64StackTaggingPreRA::findFirstSlotCandidate\n");
+  if (!ClFirstSlot)
+    return None;
+
+  DenseMap<SlotWithTag, int> RetagScore;
+  SlotWithTag MaxScoreST{-1, -1};
+  int MaxScore = -1;
+  for (auto *I : ReTags) {
+    SlotWithTag ST{*I};
+    if (isSlotPreAllocated(MFI, ST.FI))
+      continue;
+
+    Register RetagReg = I->getOperand(0).getReg();
+    if (!Register::isVirtualRegister(RetagReg))
+      continue;
+
+    int Score = 0;
+    SmallVector<Register, 8> WorkList;
+    WorkList.push_back(RetagReg);
+
+    while (!WorkList.empty()) {
+      Register UseReg = WorkList.back();
+      WorkList.pop_back();
+      for (auto &UseI : MRI->use_instructions(UseReg)) {
+        unsigned Opcode = UseI.getOpcode();
+        if (Opcode == AArch64::STGOffset || Opcode == AArch64::ST2GOffset ||
+            Opcode == AArch64::STZGOffset || Opcode == AArch64::STZ2GOffset ||
+            Opcode == AArch64::STGPi || Opcode == AArch64::STGloop ||
+            Opcode == AArch64::STZGloop || Opcode == AArch64::STGloop_wback ||
+            Opcode == AArch64::STZGloop_wback)
+          continue;
+        if (UseI.isCopy()) {
+          Register DstReg = UseI.getOperand(0).getReg();
+          if (Register::isVirtualRegister(DstReg))
+            WorkList.push_back(DstReg);
+          continue;
+        }
+        LLVM_DEBUG(dbgs() << "[" << ST.FI << ":" << ST.Tag << "] use of %"
+                          << Register::virtReg2Index(UseReg) << " in " << UseI
+                          << "\n");
+        Score++;
+      }
+    }
+
+    int TotalScore = RetagScore[ST] += Score;
+    if (TotalScore > MaxScore ||
+        (TotalScore == MaxScore && ST.FI > MaxScoreST.FI)) {
+      MaxScore = TotalScore;
+      MaxScoreST = ST;
+    }
+  }
+
+  if (MaxScoreST.FI < 0)
+    return None;
+
+  // If FI's tag is already 0, we are done.
+  if (MaxScoreST.Tag == 0)
+    return MaxScoreST.FI;
+
+  // Otherwise, find a random victim pair (FI, Tag) where Tag == 0.
+  SlotWithTag SwapST{-1, -1};
+  for (auto *I : ReTags) {
+    SlotWithTag ST{*I};
+    if (ST.Tag == 0) {
+      SwapST = ST;
+      break;
+    }
+  }
+
+  // Swap tags between the victim and the highest scoring pair.
+  // If SwapWith is still (-1, -1), that's fine, too - we'll simply take tag for
+  // the highest score slot without changing anything else.
+  for (auto *&I : ReTags) {
+    SlotWithTag ST{*I};
+    MachineOperand &TagOp = I->getOperand(4);
+    if (ST == MaxScoreST) {
+      TagOp.setImm(0);
+    } else if (ST == SwapST) {
+      TagOp.setImm(MaxScoreST.Tag);
+    }
+  }
+  return MaxScoreST.FI;
+}
+
 bool AArch64StackTaggingPreRA::runOnMachineFunction(MachineFunction &Func) {
   MF = &Func;
   MRI = &MF->getRegInfo();
   AFI = MF->getInfo<AArch64FunctionInfo>();
   TII = static_cast<const AArch64InstrInfo *>(MF->getSubtarget().getInstrInfo());
   TRI = static_cast<const AArch64RegisterInfo *>(
       MF->getSubtarget().getRegisterInfo());
   MFI = &MF->getFrameInfo();
for (auto &I : BB) {
         int FI = I.getOperand(1).getIndex();
         TaggedSlots.insert(FI);
         // There should be no offsets in TAGP yet.
         assert(I.getOperand(2).getImm() == 0);
       }
     }
   }
 
+  // Take over from SSP. It does nothing for tagged slots, and should not really
+  // have been enabled in the first place.
+  for (int FI : TaggedSlots)
+    MFI->setObjectSSPLayout(FI, MachineFrameInfo::SSPLK_None);
+
   if (ReTags.empty())
     return false;
 
   if (mayUseUncheckedLoadStore())
     uncheckLoadsAndStores();
 
+  // Find a slot that is used with zero tag offset, like ADDG #fi, 0.
+  // If the base tagged pointer is set up to the address of this slot,
+  // the ADDG instruction can be eliminated.
+  Optional<int> BaseSlot = findFirstSlotCandidate();
+  if (BaseSlot)
+    AFI->setTaggedBasePointerIndex(*BaseSlot);
+
+  for (auto *I : ReTags) {
+    int FI = I->getOperand(1).getIndex();
+    int Tag = I->getOperand(4).getImm();
+    Register Base = I->getOperand(3).getReg();
+    if (Tag == 0 && FI == BaseSlot) {
+      BuildMI(*I->getParent(), I, {}, TII->get(AArch64::COPY),
+              I->getOperand(0).getReg())
+          .addReg(Base);
+      I->eraseFromParent();
+    }
+  }
+
   return true;
 }

llvm/test/CodeGen/AArch64/irg_sp_tagp.ll

 ; RUN: llc < %s -mtriple=aarch64 -mattr=+mte | FileCheck %s
 
 define i8* @small_alloca() {
 entry:
 ; CHECK-LABEL: small_alloca:
+; CHECK:      irg  x0, sp{{$}}
+; CHECK:      ret
+  %a = alloca i8, align 16
+  %q = call i8* @llvm.aarch64.irg.sp(i64 0)
+  %q1 = call i8* @llvm.aarch64.tagp.p0i8(i8* %a, i8* %q, i64 1)
+  ret i8* %q1
+}
+
+@sink = global i8* null, align 8
+
+; Check that IRG is pinned to %b because the store instruction needs
+; the address in a non-fixed physical register and can benefit from it
+; being equal to the base tagged pointer.
+define i8* @small_allocas() {
+entry:
+; CHECK-LABEL: small_allocas:
 ; CHECK:      irg  [[R:x[0-9]+]], sp{{$}}
-; CHECK-NEXT: addg x0, [[R]], #0, #1
+; CHECK:      addg x0, [[R]], #16, #1
+; CHECK:      str  [[R]], {{.*}}sink
 ; CHECK:      ret
   %a = alloca i8, align 16
+  %b = alloca i8, align 16
   %q = call i8* @llvm.aarch64.irg.sp(i64 0)
   %q1 = call i8* @llvm.aarch64.tagp.p0i8(i8* %a, i8* %q, i64 1)
+  %q2 = call i8* @llvm.aarch64.tagp.p0i8(i8* %b, i8* %q, i64 2)
+  store i8* %q2, i8** @sink, align 8
   ret i8* %q1
 }
 
 ; Two large allocas. One's offset overflows addg immediate.
 define void @huge_allocas() {
 entry:
 ; CHECK-LABEL: huge_allocas:
-; CHECK:      irg  [[R:x[0-9]+]], sp{{$}}
-; CHECK:      add  [[TMP:x[0-9]+]], [[R]], #3088
+; CHECK:      irg  x1, sp{{$}}
+; CHECK:      add  [[TMP:x[0-9]+]], x1, #3088
 ; CHECK:      addg x0, [[TMP]], #1008, #1
-; CHECK:      addg x1, [[R]], #0, #2
 ; CHECK:      bl use2
   %a = alloca i8, i64 4096, align 16
   %b = alloca i8, i64 4096, align 16
   %base = call i8* @llvm.aarch64.irg.sp(i64 0)
   %a_t = call i8* @llvm.aarch64.tagp.p0i8(i8* %a, i8* %base, i64 1)
-  %b_t = call i8* @llvm.aarch64.tagp.p0i8(i8* %b, i8* %base, i64 2)
+  %b_t = call i8* @llvm.aarch64.tagp.p0i8(i8* %b, i8* %base, i64 0)
   call void @use2(i8* %a_t, i8* %b_t)
   ret void
 }
 
 ; Realigned stack frame. IRG uses value of SP after realignment,
 ; ADDG for the first stack allocation has offset 0.
 define void @realign() {
 entry:
 ; CHECK-LABEL: realign:
 ; CHECK:      mov  x29, sp
 ; CHECK:      and  sp, x{{[0-9]*}}, #0xffffffffffffffc0
-; CHECK:      irg  [[R:x[0-9]+]], sp{{$}}
-; CHECK:      addg x0, [[R]], #0, #1
+; CHECK:      irg  x0, sp{{$}}
 ; CHECK:      bl use
   %a = alloca i8, i64 4096, align 64
   %base = call i8* @llvm.aarch64.irg.sp(i64 0)
   %a_t = call i8* @llvm.aarch64.tagp.p0i8(i8* %a, i8* %base, i64 1)
   call void @use(i8* %a_t)
   ret void
 }
 
 ; With a dynamic alloca, IRG has to use FP with non-zero offset.
 ; ADDG offset for the single static alloca is still zero.
 define void @dynamic_alloca(i64 %size) {
 entry:
 ; CHECK-LABEL: dynamic_alloca:
-; CHECK:      sub  [[R:x[0-9]+]], x29, #[[OFS:[0-9]+]]
-; CHECK:      irg  [[R]], [[R]]
-; CHECK:      addg x1, [[R]], #0, #1
-; CHECK:      sub  x0, x29, #[[OFS]]
+; CHECK:      sub  x1, x29, #[[OFS:[0-9]+]]
+; CHECK:      irg  x1, x1
+; CHECK-DAG:  sub  x0, x29, #[[OFS]]
 ; CHECK:      bl   use2
   %base = call i8* @llvm.aarch64.irg.sp(i64 0)
   %a = alloca i128, i64 %size, align 16
   %b = alloca i8, i64 16, align 16
   %b_t = call i8* @llvm.aarch64.tagp.p0i8(i8* %b, i8* %base, i64 1)
   call void @use2(i8* %b, i8* %b_t)
   ret void
 }
 
 ; Both dynamic alloca and realigned frame.
 ; After initial realignment, generate the base pointer.
 ; IRG uses the base pointer w/o offset.
 ; Offsets for tagged and untagged pointers to the same alloca match.
 define void @dynamic_alloca_and_realign(i64 %size) {
 entryz:
 ; CHECK-LABEL: dynamic_alloca_and_realign:
 ; CHECK:      and  sp, x{{.*}}, #0xffffffffffffffc0
 ; CHECK:      mov  x19, sp
-; CHECK:      irg  [[R:x[0-9]+]], x19
-; CHECK:      addg x1, [[R]], #[[OFS:[0-9]+]], #1
-; CHECK:      add  x0, x19, #[[OFS]]
+; CHECK:      add  x1, x19, #[[OFS:[0-9]+]]
+; CHECK:      irg  x1, x1
+; CHECK-DAG:  add  x0, x19, #[[OFS]]
 ; CHECK:      bl   use2
   %base = call i8* @llvm.aarch64.irg.sp(i64 0)
   %a = alloca i128, i64 %size, align 64
   %b = alloca i8, i64 16, align 16
   %b_t = call i8* @llvm.aarch64.tagp.p0i8(i8* %b, i8* %base, i64 1)
   call void @use2(i8* %b, i8* %b_t)
   ret void
 }
 
 declare void @use(i8*)
 declare void @use2(i8*, i8*)
 
 declare i8* @llvm.aarch64.irg.sp(i64 %exclude)
 declare i8* @llvm.aarch64.tagp.p0i8(i8* %p, i8* %tag, i64 %ofs)