Reapply "[MTE] add stack frame history buffer"

In the reverted change, the order of the IR was dependent on the host
compiler, because we inserted instructions in arguments to functions.
Fix that, and also fix another problem with the test.

This reverts commit 3313f28.

Author: fmayer

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);
 }
 

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 @@ class AArch64StackTagging : public FunctionPass {
                                    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 *FP = memtag::getFP(IRB);
+    Value *Tag = IRB.CreateAnd(IRB.CreatePtrToInt(Base, IntptrTy), TagMask);
+    Value *TaggedFP = IRB.CreateOr(FP, Tag);
+    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

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 [[FP]] to i64
+; INSTR:  [[BASE_INT:%.*]] = ptrtoint ptr [[BASE]] 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]]]