[MTE] add stack frame history buffer (#86356)upstream/users/fmayer/spr/main.wip-mte-stack-record

this will allow us to find offending objects in a symbolization step,
like we can do with hwasan.

needs matching changes in AOSP:

https://android-review.git.corp.google.com/q/topic:%22stackhistorybuffer%22

3 files changed, 134 insertions, 2 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index dc7759367687..cd532671f501 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -2500,7 +2500,8 @@ AArch64FrameLowering::getFrameIndexReference(const MachineFunction &MF, int FI,
   return resolveFrameIndexReference(
       MF, FI, FrameReg,
       /*PreferFP=*/
-      MF.getFunction().hasFnAttribute(Attribute::SanitizeHWAddress),
+      MF.getFunction().hasFnAttribute(Attribute::SanitizeHWAddress) ||
+          MF.getFunction().hasFnAttribute(Attribute::SanitizeMemTag),
       /*ForSimm=*/false);
 }
 
diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index aabc5d5d22e2..eab3a90e57e2 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -11,6 +11,7 @@
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
 #include "AArch64TargetMachine.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
@@ -21,6 +22,7 @@
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
+#include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/CodeGen/LiveRegUnits.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -82,6 +84,26 @@ static cl::opt<size_t> ClMaxLifetimes(
     cl::desc("How many lifetime ends to handle for a single alloca."),
     cl::Optional);
 
+// Mode for selecting how to insert frame record info into the stack ring
+// buffer.
+enum RecordStackHistoryMode {
+  // Do not record frame record info.
+  none,
+
+  // Insert instructions into the prologue for storing into the stack ring
+  // buffer directly.
+  instr,
+};
+
+static cl::opt<RecordStackHistoryMode> ClRecordStackHistory(
+    "stack-tagging-record-stack-history",
+    cl::desc("Record stack frames with tagged allocations in a thread-local "
+             "ring buffer"),
+    cl::values(clEnumVal(none, "Do not record stack ring history"),
+               clEnumVal(instr, "Insert instructions into the prologue for "
+                                "storing into the stack ring buffer")),
+    cl::Hidden, cl::init(none));
+
 static const Align kTagGranuleSize = Align(16);
 
 namespace {
@@ -309,6 +331,7 @@ public:
                                    uint64_t Size, InitializerBuilder &IB);
 
   Instruction *insertBaseTaggedPointer(
+      const Module &M,
       const MapVector<AllocaInst *, memtag::AllocaInfo> &Allocas,
       const DominatorTree *DT);
   bool runOnFunction(Function &F) override;
@@ -437,6 +460,7 @@ void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
 }
 
 Instruction *AArch64StackTagging::insertBaseTaggedPointer(
+    const Module &M,
     const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument,
     const DominatorTree *DT) {
   BasicBlock *PrologueBB = nullptr;
@@ -458,6 +482,41 @@ Instruction *AArch64StackTagging::insertBaseTaggedPointer(
   Instruction *Base =
       IRB.CreateCall(IRG_SP, {Constant::getNullValue(IRB.getInt64Ty())});
   Base->setName("basetag");
+  auto TargetTriple = Triple(M.getTargetTriple());
+  // This is not a stable ABI for now, so only allow in dev builds with API
+  // level 10000.
+  // The ThreadLong format is the same as with HWASan, but the entries for
+  // stack MTE take two slots (16 bytes).
+  if (ClRecordStackHistory == instr && TargetTriple.isAndroid() &&
+      TargetTriple.isAArch64() && !TargetTriple.isAndroidVersionLT(10000) &&
+      !AllocasToInstrument.empty()) {
+    constexpr int StackMteSlot = -3;
+    constexpr uint64_t TagMask = 0xFULL << 56;
+
+    auto *IntptrTy = IRB.getIntPtrTy(M.getDataLayout());
+    Value *SlotPtr = memtag::getAndroidSlotPtr(IRB, StackMteSlot);
+    auto *ThreadLong = IRB.CreateLoad(IntptrTy, SlotPtr);
+    Value *TaggedFP = IRB.CreateOr(
+        memtag::getFP(IRB),
+        IRB.CreateAnd(IRB.CreatePtrToInt(Base, IntptrTy), TagMask));
+    Value *PC = memtag::getPC(TargetTriple, IRB);
+    Value *RecordPtr = IRB.CreateIntToPtr(ThreadLong, IRB.getPtrTy(0));
+    IRB.CreateStore(PC, RecordPtr);
+    IRB.CreateStore(TaggedFP, IRB.CreateConstGEP1_64(IntptrTy, RecordPtr, 1));
+    // Update the ring buffer. Top byte of ThreadLong defines the size of the
+    // buffer in pages, it must be a power of two, and the start of the buffer
+    // must be aligned by twice that much. Therefore wrap around of the ring
+    // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
+    // The use of AShr instead of LShr is due to
+    //   https://bugs.llvm.org/show_bug.cgi?id=39030
+    // Runtime library makes sure not to use the highest bit.
+    Value *WrapMask = IRB.CreateXor(
+        IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
+        ConstantInt::get(IntptrTy, (uint64_t)-1));
+    Value *ThreadLongNew = IRB.CreateAnd(
+        IRB.CreateAdd(ThreadLong, ConstantInt::get(IntptrTy, 16)), WrapMask);
+    IRB.CreateStore(ThreadLongNew, SlotPtr);
+  }
   return Base;
 }
 
@@ -513,7 +572,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
-  Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
+  Instruction *Base =
+      insertBaseTaggedPointer(*Fn.getParent(), SInfo.AllocasToInstrument, DT);
 
   int NextTag = 0;
   for (auto &I : SInfo.AllocasToInstrument) {
@@ -575,6 +635,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
       for (auto *II : Info.LifetimeEnd)
         II->eraseFromParent();
     }
+
+    memtag::annotateDebugRecords(Info, static_cast<unsigned long>(Tag));
   }
 
   // If we have instrumented at least one alloca, all unrecognized lifetime
diff --git a/llvm/test/CodeGen/AArch64/stack-tagging-prologue.ll b/llvm/test/CodeGen/AArch64/stack-tagging-prologue.ll
new file mode 100644
index 000000000000..3f55f3cc9a2e
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/stack-tagging-prologue.ll
@@ -0,0 +1,69 @@
+; RUN: opt < %s -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 -S -o - | FileCheck %s --check-prefixes=CHECK
+; RUN: opt < %s -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 -S -stack-tagging-record-stack-history=instr -o - | FileCheck %s --check-prefixes=INSTR
+; RUN llc -mattr=+mte -stack-tagging-use-stack-safety=0 -stack-tagging-record-stack-history=instr %s -o - | FileCheck %s --check-prefixes=ASMINSTR
+
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64--linux-android10000"
+
+declare void @use8(ptr)
+declare void @use32(ptr)
+declare void @llvm.lifetime.start.p0(i64, ptr nocapture)
+declare void @llvm.lifetime.end.p0(i64, ptr nocapture)
+
+define dso_local void @noUse32(ptr) sanitize_memtag {
+entry:
+  ret void
+}
+
+define void @OneVar() sanitize_memtag {
+entry:
+  %x = alloca i32, align 4
+  call void @use32(ptr %x)
+  ret void
+}
+
+; CHECK-LABEL: define void @OneVar(
+; CHECK:  [[BASE:%.*]] = call ptr @llvm.aarch64.irg.sp(i64 0)
+; CHECK:  [[X:%.*]] = alloca { i32, [12 x i8] }, align 16
+; CHECK:  [[TX:%.*]] = call ptr @llvm.aarch64.tagp.{{.*}}(ptr [[X]], ptr [[BASE]], i64 0)
+; CHECK:  ret void
+
+; INSTR-LABEL: define void @OneVar(
+; INSTR:  [[BASE:%.*]] = call ptr @llvm.aarch64.irg.sp(i64 0)
+; INSTR:  [[TLS:%.*]] = call ptr @llvm.thread.pointer()
+; INSTR:  [[TLS_SLOT:%.*]] = getelementptr i8, ptr [[TLS]], i32 -24
+; INSTR:  [[TLS_VALUE:%.*]] = load i64, ptr %1, align 8
+; INSTR:  [[FP:%.*]] = call ptr @llvm.frameaddress.p0(i32 0)
+; INSTR:  [[FP_INT:%.*]] = ptrtoint ptr %3 to i64
+; INSTR:  [[BASE_INT:%.*]] = ptrtoint ptr %basetag to i64
+; INSTR:  [[BASE_TAG:%.*]] = and i64 [[BASE_INT]], 1080863910568919040
+; INSTR:  [[TAGGED_FP:%.*]] = or i64 [[FP_INT]], [[BASE_TAG]]
+; INSTR:  [[PC:%.*]] = call i64 @llvm.read_register.i64(metadata !0)
+; INSTR:  [[TLS_VALUE_PTR:%.*]] = inttoptr i64 [[TLS_VALUE]] to ptr
+; INSTR:  store i64 [[PC]], ptr [[TLS_VALUE_PTR]], align 8
+; INSTR:  [[SECOND_SLOT:%.*]] = getelementptr i64, ptr [[TLS_VALUE_PTR]], i64 1
+; INSTR:  store i64 [[TAGGED_FP]], ptr [[SECOND_SLOT]], align 8
+; INSTR:  [[SIZE_IN_PAGES:%.*]] = ashr i64 [[TLS_VALUE]], 56
+; INSTR:  [[WRAP_MASK_INTERMEDIARY:%.*]] = shl nuw nsw i64 [[SIZE_IN_PAGES]], 12
+; INSTR:  [[WRAP_MASK:%.*]] = xor i64 [[WRAP_MASK_INTERMEDIARY]], -1
+; INSTR:  [[NEXT_TLS_VALUE_BEFORE_WRAP:%.*]] = add i64 [[TLS_VALUE]], 16
+; INSTR:  [[NEXT_TLS_VALUE:%.*]] = and i64 [[NEXT_TLS_VALUE_BEFORE_WRAP]], [[WRAP_MASK]]
+; INSTR:  store i64 [[NEXT_TLS_VALUE]], ptr [[TLS_SLOT]], align 8
+; INSTR:  [[X:%.*]] = alloca { i32, [12 x i8] }, align 16
+; INSTR:  [[TX:%.*]] = call ptr @llvm.aarch64.tagp.{{.*}}(ptr [[X]], ptr [[BASE]], i64 0)
+; INSTR:  [[PC:!.*]] = !{!"pc"}
+
+; ASMINSTR-LABEL: OneVar:
+; ASMINSTR:  mrs	[[TLS:x.*]], TPIDR_EL0
+; ASMINSTR:  irg	[[BASE:x.*]], sp
+; ASMINSTR:  adr	[[PC:x.*]], #0
+; ASMINSTR:  ldur	[[TLS_SLOT:x.*]], [[[TLS]], #-24]
+; ASMINSTR:  and	[[SP_TAG:x.*]], [[BASE]], #0xf00000000000000
+; ASMINSTR:  orr	[[TAGGED_FP]], x29, [[SP_TAG]]
+; ASMINSTR:  asr	[[TLS_SIZE:x.*]], [[TLS_SLOT]], #56
+; ASMINSTR:  add	[[NEXT_TLS_VALUE_BEFORE_WRAP:x.*]], [[TLS_SLOT]], #16
+; ASMINSTR:  stp	[[PC]], [[TAGGED_FP]], [[[TLS_SLOT]]]
+; ASMINSTR:  bic	[[NEXT_TLS_VALUE:x.*]], [[NEXT_TLS_VALUE_BEFORE_WRAP]], [[TLS_SIZE]], lsl #12
+; ASMINSTR:  stur	[[NEXT_TLS_VALUE]], [[[TLS]], #-24]
+; ASMINSTR:  stg	[[BASE]], [[[BASE]]]