[llvm] 205308d - [NFC] [MTE] Move alignAndPadAlloca to MemoryTaggingSupport.

Florian Mayer via llvm-commits llvm-commits at lists.llvm.org
Mon Feb 14 14:54:13 PST 2022


Author: Florian Mayer
Date: 2022-02-14T14:54:04-08:00
New Revision: 205308de6b37cf894951a2fa63e787641bb4739d

URL: https://github.com/llvm/llvm-project/commit/205308de6b37cf894951a2fa63e787641bb4739d
DIFF: https://github.com/llvm/llvm-project/commit/205308de6b37cf894951a2fa63e787641bb4739d.diff

LOG: [NFC] [MTE] Move alignAndPadAlloca to MemoryTaggingSupport.

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D119610

Added: 
    

Modified: 
    llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
    llvm/lib/Target/AArch64/AArch64StackTagging.cpp
    llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
    llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
index b466393012778..a6218136ce890 100644
--- a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
+++ b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
@@ -101,6 +101,9 @@ class StackInfoBuilder {
   std::function<bool(const AllocaInst &)> 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
index 59037440c2b99..88315869a6ffc 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -305,7 +305,6 @@ class AArch64StackTagging : public FunctionPass {
   }
 
   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 @@ Instruction *AArch64StackTagging::insertBaseTaggedPointer(
   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<ConstantInt>(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 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   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<DominatorTree> DeleteDT;

diff  --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index 34b0576982f5a..57197295f0340 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -1046,11 +1046,6 @@ bool HWAddressSanitizer::instrumentMemAccess(InterestingMemoryOperand &O) {
   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 @@ bool HWAddressSanitizer::instrumentStack(
                                                           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 @@ bool HWAddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
           // 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<AllocaInst *, AllocaInst *> HWAddressSanitizer::padInterestingAllocas(
   DenseMap<AllocaInst *, AllocaInst *> 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
index 132f44f88534e..da77457f42fc3 100644
--- a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
+++ b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
@@ -104,5 +104,44 @@ void StackInfoBuilder::visit(Instruction &Inst) {
     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<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;
+}
+
 } // namespace memtag
 } // namespace llvm


        


More information about the llvm-commits mailing list