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 @@ -102,7 +102,7 @@ }; uint64_t getAllocaSizeInBytes(const AllocaInst &AI); -void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Align); +bool alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Align); } // namespace memtag } // namespace llvm 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 @@ -505,7 +505,9 @@ for (auto &I : SInfo.AllocasToInstrument) { memtag::AllocaInfo &Info = I.second; assert(Info.AI && isInterestingAlloca(*Info.AI)); - memtag::alignAndPadAlloca(Info, kTagGranuleSize); + auto *PrevAI = Info.AI; + if (memtag::alignAndPadAlloca(Info, kTagGranuleSize)) + PrevAI->eraseFromParent(); } std::unique_ptr DeleteDT; 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 @@ -261,8 +261,6 @@ void setSSI(const StackSafetyGlobalInfo *S) { SSI = S; } - DenseMap padInterestingAllocas( - const MapVector &AllocasToInstrument); bool sanitizeFunction(Function &F, llvm::function_ref GetDT, llvm::function_ref GetPDT); @@ -1380,6 +1378,19 @@ II->eraseFromParent(); } } + if (memtag::alignAndPadAlloca(Info, Align(Mapping.getObjectAlignment()))) { + for (auto DVI : Info.DbgVariableIntrinsics) { + SmallDenseSet LocationOps(DVI->location_ops().begin(), + DVI->location_ops().end()); + for (Value *V : LocationOps) { + if (auto *AI = dyn_cast_or_null(V)) { + if (V == AI) + DVI->replaceVariableLocationOp(V, Info.AI); + } + } + } + AI->eraseFromParent(); + } } for (auto &I : SInfo.UnrecognizedLifetimes) I->eraseFromParent(); @@ -1404,39 +1415,6 @@ !(SSI && SSI->isSafe(AI)); } -DenseMap HWAddressSanitizer::padInterestingAllocas( - const MapVector &AllocasToInstrument) { - DenseMap AllocaToPaddedAllocaMap; - for (auto &KV : AllocasToInstrument) { - AllocaInst *AI = KV.first; - uint64_t Size = memtag::getAllocaSizeInBytes(*AI); - uint64_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment()); - AI->setAlignment( - Align(std::max(AI->getAlignment(), Mapping.getObjectAlignment()))); - if (Size != AlignedSize) { - Type *AllocatedType = AI->getAllocatedType(); - if (AI->isArrayAllocation()) { - uint64_t ArraySize = - cast(AI->getArraySize())->getZExtValue(); - AllocatedType = ArrayType::get(AllocatedType, ArraySize); - } - Type *TypeWithPadding = StructType::get( - AllocatedType, ArrayType::get(Int8Ty, AlignedSize - Size)); - auto *NewAI = new AllocaInst( - TypeWithPadding, AI->getType()->getAddressSpace(), nullptr, "", AI); - NewAI->takeName(AI); - NewAI->setAlignment(AI->getAlign()); - NewAI->setUsedWithInAlloca(AI->isUsedWithInAlloca()); - NewAI->setSwiftError(AI->isSwiftError()); - NewAI->copyMetadata(*AI); - auto *Bitcast = new BitCastInst(NewAI, AI->getType(), "", AI); - AI->replaceAllUsesWith(Bitcast); - AllocaToPaddedAllocaMap[AI] = NewAI; - } - } - return AllocaToPaddedAllocaMap; -} - bool HWAddressSanitizer::sanitizeFunction( Function &F, llvm::function_ref GetDT, llvm::function_ref GetPDT) { @@ -1509,28 +1487,6 @@ instrumentStack(DetectUseAfterScope && !SInfo.CallsReturnTwice, SIB.get(), StackTag, GetDT, GetPDT); } - // Pad and align each of the allocas that we instrumented to stop small - // uninteresting allocas from hiding in instrumented alloca's padding and so - // that we have enough space to store real tags for short granules. - DenseMap AllocaToPaddedAllocaMap = - padInterestingAllocas(SInfo.AllocasToInstrument); - - if (!AllocaToPaddedAllocaMap.empty()) { - for (auto &Inst : instructions(F)) { - if (auto *DVI = dyn_cast(&Inst)) { - SmallDenseSet LocationOps(DVI->location_ops().begin(), - DVI->location_ops().end()); - for (Value *V : LocationOps) { - if (auto *AI = dyn_cast_or_null(V)) { - if (auto *NewAI = AllocaToPaddedAllocaMap.lookup(AI)) - DVI->replaceVariableLocationOp(V, NewAI); - } - } - } - } - for (auto &P : AllocaToPaddedAllocaMap) - P.first->eraseFromParent(); - } // If we split the entry block, move any allocas that were originally in the // entry block back into the entry block so that they aren't treated as 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 @@ -109,7 +109,7 @@ return AI.getAllocationSizeInBits(DL).getValue() / 8; } -void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) { +bool alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) { const Align NewAlignment = max(MaybeAlign(Info.AI->getAlign()), Alignment); Info.AI->setAlignment(NewAlignment); auto &Ctx = Info.AI->getFunction()->getContext(); @@ -117,7 +117,7 @@ uint64_t Size = getAllocaSizeInBytes(*Info.AI); uint64_t AlignedSize = alignTo(Size, Alignment); if (Size == AlignedSize) - return; + return false; // Add padding to the alloca. Type *AllocatedType = @@ -139,8 +139,8 @@ auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI); Info.AI->replaceAllUsesWith(NewPtr); - Info.AI->eraseFromParent(); Info.AI = NewAI; + return true; } } // namespace memtag