[HWASAN] use common alignAndPadAlloca

[fmayer]

I am also looking whether this code in Aarch64StackTagging needs changing or whether there are different constraints (because HWASAN does more changes to debug intrinsics)

// Fixup debug intrinsics to point to the new alloca.
for (auto DVI : Info.DbgVariableIntrinsics)
  DVI->replaceVariableLocationOp(Info.OldAI, Info.AI)

If we only use this code for HWASAN, IR regression tests (alloca.ll and dbg-value-tag-offset.ll) fail.

[fmayer]

In D119614#3316412, @fmayer wrote:

I am also looking whether this code in Aarch64StackTagging needs changing or whether there are different constraints (because HWASAN does more changes to debug intrinsics)

// Fixup debug intrinsics to point to the new alloca.
for (auto DVI : Info.DbgVariableIntrinsics)
  DVI->replaceVariableLocationOp(Info.OldAI, Info.AI)

If we only use this code for HWASAN, IR regression tests (alloca.ll and dbg-value-tag-offset.ll) fail.

Actually, I need to look more at this now.

[fmayer]

update

[eugenis]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

1382

Why don't we do the same in AArch64StackTagging?

[fmayer]

don't change too much at once.

[fmayer]

In D119614#3316423, @fmayer wrote:

In D119614#3316412, @fmayer wrote:

I am also looking whether this code in Aarch64StackTagging needs changing or whether there are different constraints (because HWASAN does more changes to debug intrinsics)

// Fixup debug intrinsics to point to the new alloca.
for (auto DVI : Info.DbgVariableIntrinsics)
  DVI->replaceVariableLocationOp(Info.OldAI, Info.AI)

If we only use this code for HWASAN, IR regression tests (alloca.ll and dbg-value-tag-offset.ll) fail.

Actually, I need to look more at this now.

I've looked into this. I think it makes sense to allow replacing multiple

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

1382

I am looking into that separately. I also reverted this code to the one we had in HWASan, so we don't do too many changes in one go.

[eugenis]

LGTM

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

public:
   StackInfo &get() { return Info; };
 
 private:
   StackInfo Info;
   std::function<bool(const AllocaInst &)> IsInterestingAlloca;
 };
 
 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
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

bool AArch64StackTagging::runOnFunction(Function &Fn) {
   memtag::StackInfo &SInfo = SIB.get();
 
   if (SInfo.AllocasToInstrument.empty())
     return false;
 
   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<DominatorTree> DeleteDT;
   DominatorTree *DT = nullptr;
   if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
     DT = &P->getDomTree();
 
   if (DT == nullptr && (SInfo.AllocasToInstrument.size() > 1 ||

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

this->CompileKernel = ClEnableKhwasan.getNumOccurrences() > 0
                               ? ClEnableKhwasan
                               : CompileKernel;
 
     initializeModule();
   }
 
   void setSSI(const StackSafetyGlobalInfo *S) { SSI = S; }
 
-  DenseMap<AllocaInst *, AllocaInst *> padInterestingAllocas(
-      const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument);
   bool sanitizeFunction(Function &F,
                         llvm::function_ref<const DominatorTree &()> GetDT,
                         llvm::function_ref<const PostDominatorTree &()> GetPDT);
   void initializeModule();
   void createHwasanCtorComdat();
 
   void initializeCallbacks(Module &M);
 
if (ShouldDetectUseAfterScope && StandardLifetime) {
         TagEnd(RI);
       if (!StandardLifetime) {
         for (auto &II : Info.LifetimeStart)
           II->eraseFromParent();
         for (auto &II : Info.LifetimeEnd)
           II->eraseFromParent();
       }
     }
+    if (memtag::alignAndPadAlloca(Info, Align(Mapping.getObjectAlignment()))) {
+      for (auto DVI : Info.DbgVariableIntrinsics) {
+        SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
+                                           DVI->location_ops().end());
+        for (Value *V : LocationOps) {
+          if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
+            if (V == AI)
+              DVI->replaceVariableLocationOp(V, Info.AI);
+          }
+        }
+      }
+      AI->eraseFromParent();
+    }
   }
   for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
   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));
 }
 
-DenseMap<AllocaInst *, AllocaInst *> HWAddressSanitizer::padInterestingAllocas(
-    const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument) {
-  DenseMap<AllocaInst *, AllocaInst *> 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<ConstantInt>(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<const DominatorTree &()> GetDT,
     llvm::function_ref<const PostDominatorTree &()> GetPDT) {
   if (&F == HwasanCtorFunction)
     return false;
 
   if (!F.hasFnAttribute(Attribute::SanitizeHWAddress))
     return false;
Value *StackTag =
         ClGenerateTagsWithCalls ? nullptr : getStackBaseTag(EntryIRB);
     // 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.
     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<AllocaInst *, AllocaInst *> AllocaToPaddedAllocaMap =
-      padInterestingAllocas(SInfo.AllocasToInstrument);
-
-  if (!AllocaToPaddedAllocaMap.empty()) {
-    for (auto &Inst : instructions(F)) {
-      if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
-        SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
-                                           DVI->location_ops().end());
-        for (Value *V : LocationOps) {
-          if (auto *AI = dyn_cast_or_null<AllocaInst>(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
   // dynamic allocas.
   if (EntryIRB.GetInsertBlock() != &F.getEntryBlock()) {
     InsertPt = &*F.getEntryBlock().begin();
     for (Instruction &I :
          llvm::make_early_inc_range(*EntryIRB.GetInsertBlock())) {

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

if (ExitUntag)
     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) {
+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();
 
   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 =
       Info.AI->isArrayAllocation()
           ? ArrayType::get(
                 Info.AI->getAllocatedType(),
                 cast<ConstantInt>(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;
+  return true;
 }
 
 } // namespace memtag
 } // namespace llvm