[MTE] [HWASan] unify isInterestingAlloca

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D134779

llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h

@@ -16,6 +16,7 @@
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Support/Alignment.h"
 
 namespace llvm {
@@ -62,15 +63,15 @@ struct StackInfo {
 
 class StackInfoBuilder {
 public:
-  StackInfoBuilder(std::function<bool(const AllocaInst &)> IsInterestingAlloca)
-      : IsInterestingAlloca(IsInterestingAlloca) {}
+  StackInfoBuilder(const StackSafetyGlobalInfo *SSI) : SSI(SSI) {}
 
   void visit(Instruction &Inst);
+  bool isInterestingAlloca(const AllocaInst &AI);
   StackInfo &get() { return Info; };
 
 private:
   StackInfo Info;
-  std::function<bool(const AllocaInst &)> IsInterestingAlloca;
+  const StackSafetyGlobalInfo *SSI;
 };
 
 uint64_t getAllocaSizeInBytes(const AllocaInst &AI);

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

@@ -306,8 +306,6 @@ class AArch64StackTagging : public FunctionPass {
     initializeAArch64StackTaggingPass(*PassRegistry::getPassRegistry());
   }
 
-  bool isInterestingAlloca(const AllocaInst &AI);
-
   void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr,
                  uint64_t Size);
   void untagAlloca(AllocaInst *AI, Instruction *InsertBefore, uint64_t Size);
@@ -413,22 +411,6 @@ Instruction *AArch64StackTagging::collectInitializers(Instruction *StartInst,
   return LastInst;
 }
 
-bool AArch64StackTagging::isInterestingAlloca(const AllocaInst &AI) {
-  // FIXME: support dynamic allocas
-  bool IsInteresting =
-      AI.getAllocatedType()->isSized() && AI.isStaticAlloca() &&
-      // alloca() may be called with 0 size, ignore it.
-      *AI.getAllocationSizeInBits(*DL) > 0 &&
-      // inalloca allocas are not treated as static, and we don't want
-      // dynamic alloca instrumentation for them as well.
-      !AI.isUsedWithInAlloca() &&
-      // swifterror allocas are register promoted by ISel
-      !AI.isSwiftError() &&
-      // safe allocas are not interesting
-      !(SSI && SSI->isSafe(AI));
-  return IsInteresting;
-}
-
 void AArch64StackTagging::tagAlloca(AllocaInst *AI, Instruction *InsertBefore,
                                     Value *Ptr, uint64_t Size) {
   auto SetTagZeroFunc =
@@ -495,8 +477,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   if (MergeInit)
     AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
-  memtag::StackInfoBuilder SIB(
-      [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
+  memtag::StackInfoBuilder SIB(SSI);
   for (Instruction &I : instructions(F))
     SIB.visit(I);
   memtag::StackInfo &SInfo = SIB.get();
@@ -541,7 +522,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   int NextTag = 0;
   for (auto &I : SInfo.AllocasToInstrument) {
     memtag::AllocaInfo &Info = I.second;
-    assert(Info.AI && isInterestingAlloca(*Info.AI));
+    assert(Info.AI && SIB.isInterestingAlloca(*Info.AI));
     TrackingVH<Instruction> OldAI = Info.AI;
     memtag::alignAndPadAlloca(Info, kTagGranuleSize);
     AllocaInst *AI = Info.AI;

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -308,7 +308,6 @@ class HWAddressSanitizer {
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
-  bool isInterestingAlloca(const AllocaInst &AI);
   void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
@@ -1397,24 +1396,6 @@ bool HWAddressSanitizer::instrumentStack(memtag::StackInfo &SInfo,
   return true;
 }
 
-bool HWAddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
-  return (AI.getAllocatedType()->isSized() &&
-          // FIXME: instrument dynamic allocas, too
-          AI.isStaticAlloca() &&
-          // alloca() may be called with 0 size, ignore it.
-          memtag::getAllocaSizeInBytes(AI) > 0 &&
-          // We are only interested in allocas not promotable to registers.
-          // Promotable allocas are common under -O0.
-          !isAllocaPromotable(&AI) &&
-          // inalloca allocas are not treated as static, and we don't want
-          // dynamic alloca instrumentation for them as well.
-          !AI.isUsedWithInAlloca() &&
-          // swifterror allocas are register promoted by ISel
-          !AI.isSwiftError()) &&
-         // safe allocas are not interesting
-         !(SSI && SSI->isSafe(AI));
-}
-
 bool HWAddressSanitizer::sanitizeFunction(Function &F,
                                           FunctionAnalysisManager &FAM) {
   if (&F == HwasanCtorFunction)
@@ -1429,8 +1410,7 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F,
   SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
   SmallVector<Instruction *, 8> LandingPadVec;
 
-  memtag::StackInfoBuilder SIB(
-      [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
+  memtag::StackInfoBuilder SIB(SSI);
   for (auto &Inst : instructions(F)) {
     if (InstrumentStack) {
       SIB.visit(Inst);

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

@@ -14,9 +14,11 @@
 
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/Transforms/Utils/PromoteMemToReg.h"
 
 namespace llvm {
 namespace memtag {
@@ -114,7 +116,7 @@ void StackInfoBuilder::visit(Instruction &Inst) {
     }
   }
   if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
-    if (IsInterestingAlloca(*AI)) {
+    if (isInterestingAlloca(*AI)) {
       Info.AllocasToInstrument[AI].AI = AI;
     }
     return;
@@ -127,7 +129,7 @@ void StackInfoBuilder::visit(Instruction &Inst) {
       Info.UnrecognizedLifetimes.push_back(&Inst);
       return;
     }
-    if (!IsInterestingAlloca(*AI))
+    if (!isInterestingAlloca(*AI))
       return;
     if (II->getIntrinsicID() == Intrinsic::lifetime_start)
       Info.AllocasToInstrument[AI].LifetimeStart.push_back(II);
@@ -138,7 +140,7 @@ void StackInfoBuilder::visit(Instruction &Inst) {
   if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
     for (Value *V : DVI->location_ops()) {
       if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
-        if (!IsInterestingAlloca(*AI))
+        if (!isInterestingAlloca(*AI))
           continue;
         AllocaInfo &AInfo = Info.AllocasToInstrument[AI];
         auto &DVIVec = AInfo.DbgVariableIntrinsics;
@@ -152,6 +154,24 @@ void StackInfoBuilder::visit(Instruction &Inst) {
     Info.RetVec.push_back(ExitUntag);
 }
 
+bool StackInfoBuilder::isInterestingAlloca(const AllocaInst &AI) {
+  return (AI.getAllocatedType()->isSized() &&
+          // FIXME: instrument dynamic allocas, too
+          AI.isStaticAlloca() &&
+          // alloca() may be called with 0 size, ignore it.
+          memtag::getAllocaSizeInBytes(AI) > 0 &&
+          // We are only interested in allocas not promotable to registers.
+          // Promotable allocas are common under -O0.
+          !isAllocaPromotable(&AI) &&
+          // inalloca allocas are not treated as static, and we don't want
+          // dynamic alloca instrumentation for them as well.
+          !AI.isUsedWithInAlloca() &&
+          // swifterror allocas are register promoted by ISel
+          !AI.isSwiftError()) &&
+         // safe allocas are not interesting
+         !(SSI && SSI->isSafe(AI));
+}
+
 uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
   auto DL = AI.getModule()->getDataLayout();
   return *AI.getAllocationSizeInBits(DL) / 8;

llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll

@@ -3,11 +3,14 @@
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-arm-unknown-eabi"
 
+declare void @use8(i8*)
+
 define  void @f(i1 %cond) local_unnamed_addr sanitize_memtag {
 start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %next1
 
@@ -39,6 +42,7 @@ start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %next1
 
@@ -65,6 +69,7 @@ start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %next1
 
@@ -90,6 +95,7 @@ start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %start1
 
@@ -125,6 +131,7 @@ start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %start1
 

llvm/test/CodeGen/AArch64/stack-tagging.ll

@@ -1,5 +1,5 @@
 ; RUN: opt < %s -aarch64-stack-tagging -S -o - | FileCheck %s --check-prefixes=CHECK,SSI
-; RUN: opt < %s -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 -S -o - | FileCheck %s --check-prefixes=CHECK,NOSSI
+; RUN: opt < %s -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 -S -o - | FileCheck %s --check-prefixes=CHECK
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-android"
@@ -56,10 +56,7 @@ entry:
 ; CHECK:  alloca { [11 x i32], [4 x i8] }, align 16
 ; CHECK:  call { [11 x i32], [4 x i8] }* @llvm.aarch64.tagp.{{.*}}({ [11 x i32], [4 x i8] }* {{.*}}, i64 2)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 48)
-; SSI:    alloca i32, align 4
-; NOSSI:  alloca { i32, [12 x i8] }, align 16
-; NOSSI: @llvm.aarch64.tagp.
-; NOSSI: call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
+; CHECK:  alloca i32, align 4
 ; SSI-NOT: @llvm.aarch64.tagp
 ; SSI-NOT: @llvm.aarch64.settag
 
@@ -70,7 +67,6 @@ entry:
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 48)
-; NOSSI:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK-NEXT:  ret void