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 @@ -101,6 +101,9 @@ std::function IsInterestingAlloca; }; +uint64_t getAllocaSizeInBytes(const AllocaInst &AI); +void 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 @@ -305,7 +305,6 @@ } bool isInterestingAlloca(const AllocaInst &AI); - void alignAndPadAlloca(memtag::AllocaInfo &Info); void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr, uint64_t Size); @@ -482,40 +481,6 @@ return Base; } -void AArch64StackTagging::alignAndPadAlloca(memtag::AllocaInfo &Info) { - const Align NewAlignment = - max(MaybeAlign(Info.AI->getAlign()), kTagGranuleSize); - Info.AI->setAlignment(NewAlignment); - - uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8; - uint64_t AlignedSize = alignTo(Size, kTagGranuleSize); - if (Size == AlignedSize) - return; - - // Add padding to the alloca. - Type *AllocatedType = - Info.AI->isArrayAllocation() - ? ArrayType::get( - Info.AI->getAllocatedType(), - cast(Info.AI->getArraySize())->getZExtValue()) - : Info.AI->getAllocatedType(); - Type *PaddingType = - ArrayType::get(Type::getInt8Ty(F->getContext()), AlignedSize - Size); - Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType); - auto *NewAI = new AllocaInst( - TypeWithPadding, Info.AI->getType()->getAddressSpace(), nullptr, "", Info.AI); - NewAI->takeName(Info.AI); - NewAI->setAlignment(Info.AI->getAlign()); - NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca()); - NewAI->setSwiftError(Info.AI->isSwiftError()); - NewAI->copyMetadata(*Info.AI); - - auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI); - Info.AI->replaceAllUsesWith(NewPtr); - Info.AI->eraseFromParent(); - Info.AI = NewAI; -} - // FIXME: check for MTE extension bool AArch64StackTagging::runOnFunction(Function &Fn) { if (!Fn.hasFnAttribute(Attribute::SanitizeMemTag)) @@ -540,7 +505,7 @@ for (auto &I : SInfo.AllocasToInstrument) { memtag::AllocaInfo &Info = I.second; assert(Info.AI && isInterestingAlloca(*Info.AI)); - alignAndPadAlloca(Info); + memtag::alignAndPadAlloca(Info, kTagGranuleSize); } 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 @@ -1046,11 +1046,6 @@ return true; } -static uint64_t getAllocaSizeInBytes(const AllocaInst &AI) { - auto DL = AI.getModule()->getDataLayout(); - return AI.getAllocationSizeInBits(DL).getValue() / 8; -} - void HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size) { size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment()); @@ -1353,7 +1348,7 @@ NewOps, LocNo)); } - size_t Size = getAllocaSizeInBytes(*AI); + size_t Size = memtag::getAllocaSizeInBytes(*AI); size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment()); auto TagEnd = [&](Instruction *Node) { IRB.SetInsertPoint(Node); @@ -1396,7 +1391,7 @@ // FIXME: instrument dynamic allocas, too AI.isStaticAlloca() && // alloca() may be called with 0 size, ignore it. - getAllocaSizeInBytes(AI) > 0 && + memtag::getAllocaSizeInBytes(AI) > 0 && // We are only interested in allocas not promotable to registers. // Promotable allocas are common under -O0. !isAllocaPromotable(&AI) && @@ -1414,7 +1409,7 @@ DenseMap AllocaToPaddedAllocaMap; for (auto &KV : AllocasToInstrument) { AllocaInst *AI = KV.first; - uint64_t Size = getAllocaSizeInBytes(*AI); + uint64_t Size = memtag::getAllocaSizeInBytes(*AI); uint64_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment()); AI->setAlignment( Align(std::max(AI->getAlignment(), Mapping.getObjectAlignment()))); 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 @@ -104,5 +104,44 @@ Info.RetVec.push_back(ExitUntag); } +uint64_t getAllocaSizeInBytes(const AllocaInst &AI) { + auto DL = AI.getModule()->getDataLayout(); + return AI.getAllocationSizeInBits(DL).getValue() / 8; +} + +void 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(); + + uint64_t Size = getAllocaSizeInBytes(*Info.AI); + uint64_t AlignedSize = alignTo(Size, Alignment); + if (Size == AlignedSize) + return; + + // Add padding to the alloca. + Type *AllocatedType = + Info.AI->isArrayAllocation() + ? ArrayType::get( + Info.AI->getAllocatedType(), + cast(Info.AI->getArraySize())->getZExtValue()) + : Info.AI->getAllocatedType(); + Type *PaddingType = ArrayType::get(Type::getInt8Ty(Ctx), AlignedSize - Size); + Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType); + auto *NewAI = + new AllocaInst(TypeWithPadding, Info.AI->getType()->getAddressSpace(), + nullptr, "", Info.AI); + NewAI->takeName(Info.AI); + NewAI->setAlignment(Info.AI->getAlign()); + NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca()); + NewAI->setSwiftError(Info.AI->isSwiftError()); + NewAI->copyMetadata(*Info.AI); + + auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI); + Info.AI->replaceAllUsesWith(NewPtr); + Info.AI->eraseFromParent(); + Info.AI = NewAI; +} + } // namespace memtag } // namespace llvm