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 @@ -67,7 +67,7 @@ bool isStandardLifetime(const SmallVectorImpl &LifetimeStart, const SmallVectorImpl &LifetimeEnd, - const DominatorTree &DT, size_t MaxLifetimes); + const DominatorTree *DT, size_t MaxLifetimes); } // namespace llvm #endif 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 @@ -78,6 +78,12 @@ ClMergeInitSizeLimit("stack-tagging-merge-init-size-limit", cl::init(272), cl::Hidden); +static cl::opt ClMaxLifetimes( + "stack-tagging-max-lifetimes-for-alloca", cl::Hidden, cl::init(3), + cl::ReallyHidden, + cl::desc("How many lifetime ends to handle for a single alloca."), + cl::Optional); + static const Align kTagGranuleSize = Align(16); namespace { @@ -645,14 +651,17 @@ Info.AI->replaceAllUsesWith(TagPCall); TagPCall->setOperand(0, Info.AI); + bool StandardLifetime = + UnrecognizedLifetimes.empty() && + isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT, + ClMaxLifetimes); // Calls to functions that may return twice (e.g. setjmp) confuse the // postdominator analysis, and will leave us to keep memory tagged after // function return. Work around this by always untagging at every return // statement if return_twice functions are called. - if (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 && - Info.LifetimeEnd.size() == 1 && !CallsReturnTwice) { + if (UnrecognizedLifetimes.empty() && StandardLifetime && + !CallsReturnTwice) { IntrinsicInst *Start = Info.LifetimeStart[0]; - IntrinsicInst *End = Info.LifetimeEnd[0]; uint64_t Size = cast(Start->getArgOperand(0))->getZExtValue(); Size = alignTo(Size, kTagGranuleSize); @@ -661,8 +670,10 @@ auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); }; if (!DT || !PDT || !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec, - TagEnd)) - End->eraseFromParent(); + TagEnd)) { + for (auto *End : Info.LifetimeEnd) + End->eraseFromParent(); + } } else { uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8; Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy()); 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 @@ -1382,7 +1382,7 @@ }; bool StandardLifetime = UnrecognizedLifetimes.empty() && - isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, GetDT(), + isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, &GetDT(), ClMaxLifetimes); if (ShouldDetectUseAfterScope && StandardLifetime) { IntrinsicInst *Start = Info.LifetimeStart[0]; 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 @@ -15,7 +15,7 @@ namespace llvm { namespace { bool maybeReachableFromEachOther(const SmallVectorImpl &Insts, - const DominatorTree &DT, size_t MaxLifetimes) { + const DominatorTree *DT, size_t MaxLifetimes) { // If we have too many lifetime ends, give up, as the algorithm below is N^2. if (Insts.size() > MaxLifetimes) return true; @@ -23,7 +23,7 @@ for (size_t J = 0; J < Insts.size(); ++J) { if (I == J) continue; - if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT)) + if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, DT)) return true; } } @@ -33,7 +33,7 @@ bool isStandardLifetime(const SmallVectorImpl &LifetimeStart, const SmallVectorImpl &LifetimeEnd, - const DominatorTree &DT, size_t MaxLifetimes) { + const DominatorTree *DT, size_t MaxLifetimes) { // An alloca that has exactly one start and end in every possible execution. // If it has multiple ends, they have to be unreachable from each other, so // at most one of them is actually used for each execution of the function. diff --git a/llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll b/llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll @@ -0,0 +1,38 @@ +; RUN: opt -S -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 %s -o - | FileCheck %s + +target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128" +target triple = "aarch64-arm-unknown-eabi" + +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) +; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48) + br i1 %cond, label %next0, label %next1 + +next0: +; CHECK-LABEL: next0: +; CHECK: call void @llvm.aarch64.settag + call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a) + br label %exit0 + +exit0: +; CHECK-LABEL: exit0: +; CHECK-NOT: call void @llvm.aarch64.settag + ret void + +next1: +; CHECK-LABEL: next1: +; CHECK: call void @llvm.aarch64.settag + call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a) + br label %exit1 + +exit1: +; CHECK-LABEL: exit1: +; CHECK-NOT: call void @llvm.aarch64.settag + ret void +} + +declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) +declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)