diff --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h --- a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h +++ b/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 @@ class StackInfoBuilder { public: - StackInfoBuilder(std::function 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 IsInterestingAlloca; + const StackSafetyGlobalInfo *SSI; }; uint64_t getAllocaSizeInBytes(const AllocaInst &AI); diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp --- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp +++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp @@ -306,8 +306,6 @@ 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 @@ 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 @@ if (MergeInit) AA = &getAnalysis().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 @@ 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 OldAI = Info.AI; memtag::alignAndPadAlloca(Info, kTagGranuleSize); AllocaInst *AI = Info.AI; diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp @@ -308,7 +308,6 @@ void getInterestingMemoryOperands( Instruction *I, SmallVectorImpl &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 @@ 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 @@ SmallVector IntrinToInstrument; SmallVector 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); diff --git a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp --- a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp +++ b/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 @@ } } if (AllocaInst *AI = dyn_cast(&Inst)) { - if (IsInterestingAlloca(*AI)) { + if (isInterestingAlloca(*AI)) { Info.AllocasToInstrument[AI].AI = AI; } return; @@ -127,7 +129,7 @@ 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 @@ if (auto *DVI = dyn_cast(&Inst)) { for (Value *V : DVI->location_ops()) { if (auto *AI = dyn_cast_or_null(V)) { - if (!IsInterestingAlloca(*AI)) + if (!isInterestingAlloca(*AI)) continue; AllocaInfo &AInfo = Info.AllocasToInstrument[AI]; auto &DVIVec = AInfo.DbgVariableIntrinsics; @@ -152,6 +154,24 @@ 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; diff --git a/llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll b/llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll --- a/llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll +++ b/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 @@ ; 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 @@ ; 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 @@ ; 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 @@ ; 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 diff --git a/llvm/test/CodeGen/AArch64/stack-tagging.ll b/llvm/test/CodeGen/AArch64/stack-tagging.ll --- a/llvm/test/CodeGen/AArch64/stack-tagging.ll +++ b/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 @@ ; 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 @@ ; 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