contrib/llvm-project/llvm/lib/Target/AArch64/AArch64StackTagging.cpp

   1 //===- AArch64StackTagging.cpp - Stack tagging in IR --===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //===----------------------------------------------------------------------===//
  10
  11 #include "AArch64.h"
  12 #include "AArch64InstrInfo.h"
  13 #include "AArch64Subtarget.h"
  14 #include "AArch64TargetMachine.h"
  15 #include "llvm/ADT/DenseMap.h"
  16 #include "llvm/ADT/DepthFirstIterator.h"
  17 #include "llvm/ADT/MapVector.h"
  18 #include "llvm/ADT/None.h"
  19 #include "llvm/ADT/Optional.h"
  20 #include "llvm/ADT/SmallVector.h"
  21 #include "llvm/ADT/Statistic.h"
  22 #include "llvm/Analysis/LoopInfo.h"
  23 #include "llvm/Analysis/ScalarEvolution.h"
  24 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
  25 #include "llvm/Analysis/ValueTracking.h"
  26 #include "llvm/CodeGen/LiveRegUnits.h"
  27 #include "llvm/CodeGen/MachineBasicBlock.h"
  28 #include "llvm/CodeGen/MachineFunction.h"
  29 #include "llvm/CodeGen/MachineFunctionPass.h"
  30 #include "llvm/CodeGen/MachineInstr.h"
  31 #include "llvm/CodeGen/MachineInstrBuilder.h"
  32 #include "llvm/CodeGen/MachineLoopInfo.h"
  33 #include "llvm/CodeGen/MachineOperand.h"
  34 #include "llvm/CodeGen/MachineRegisterInfo.h"
  35 #include "llvm/CodeGen/TargetPassConfig.h"
  36 #include "llvm/CodeGen/TargetRegisterInfo.h"
  37 #include "llvm/IR/DebugLoc.h"
  38 #include "llvm/IR/Dominators.h"
  39 #include "llvm/IR/Function.h"
  40 #include "llvm/IR/GetElementPtrTypeIterator.h"
  41 #include "llvm/IR/Instruction.h"
  42 #include "llvm/IR/Instructions.h"
  43 #include "llvm/IR/IntrinsicInst.h"
  44 #include "llvm/IR/Metadata.h"
  45 #include "llvm/Pass.h"
  46 #include "llvm/Support/Casting.h"
  47 #include "llvm/Support/Debug.h"
  48 #include "llvm/Support/raw_ostream.h"
  49 #include "llvm/Transforms/Utils/Local.h"
  50 #include <cassert>
  51 #include <iterator>
  52 #include <utility>
  53
  54 using namespace llvm;
  55
  56 #define DEBUG_TYPE "stack-tagging"
  57
  58 static constexpr unsigned kTagGranuleSize = 16;
  59
  60 namespace {
  61 class AArch64StackTagging : public FunctionPass {
  62   struct AllocaInfo {
  63     AllocaInst *AI;
  64     SmallVector<IntrinsicInst *, 2> LifetimeStart;
  65     SmallVector<IntrinsicInst *, 2> LifetimeEnd;
  66     SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;
  67     int Tag; // -1 for non-tagged allocations
  68   };
  69
  70 public:
  71   static char ID; // Pass ID, replacement for typeid
  72
  73   AArch64StackTagging() : FunctionPass(ID) {
  74     initializeAArch64StackTaggingPass(*PassRegistry::getPassRegistry());
  75   }
  76
  77   bool isInterestingAlloca(const AllocaInst &AI);
  78   void alignAndPadAlloca(AllocaInfo &Info);
  79
  80   void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr,
  81                  uint64_t Size);
  82   void untagAlloca(AllocaInst *AI, Instruction *InsertBefore, uint64_t Size);
  83
  84   Instruction *
  85   insertBaseTaggedPointer(const MapVector<AllocaInst *, AllocaInfo> &Allocas,
  86                           const DominatorTree *DT);
  87   bool runOnFunction(Function &F) override;
  88
  89   StringRef getPassName() const override { return "AArch64 Stack Tagging"; }
  90
  91 private:
  92   Function *F;
  93   Function *SetTagFunc;
  94   const DataLayout *DL;
  95
  96   void getAnalysisUsage(AnalysisUsage &AU) const override {
  97     AU.setPreservesCFG();
  98   }
  99 };
 100
 101 } // end anonymous namespace
 102
 103 char AArch64StackTagging::ID = 0;
 104
 105 INITIALIZE_PASS_BEGIN(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",
 106                       false, false)
 107 INITIALIZE_PASS_END(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",
 108                     false, false)
 109
 110 FunctionPass *llvm::createAArch64StackTaggingPass() {
 111   return new AArch64StackTagging();
 112 }
 113
 114 bool AArch64StackTagging::isInterestingAlloca(const AllocaInst &AI) {
 115   // FIXME: support dynamic allocas
 116   bool IsInteresting =
 117       AI.getAllocatedType()->isSized() && AI.isStaticAlloca() &&
 118       // alloca() may be called with 0 size, ignore it.
 119       AI.getAllocationSizeInBits(*DL).getValue() > 0 &&
 120       // inalloca allocas are not treated as static, and we don't want
 121       // dynamic alloca instrumentation for them as well.
 122       !AI.isUsedWithInAlloca() &&
 123       // swifterror allocas are register promoted by ISel
 124       !AI.isSwiftError();
 125   return IsInteresting;
 126 }
 127
 128 void AArch64StackTagging::tagAlloca(AllocaInst *AI, Instruction *InsertBefore,
 129                                     Value *Ptr, uint64_t Size) {
 130   IRBuilder<> IRB(InsertBefore);
 131   IRB.CreateCall(SetTagFunc, {Ptr, ConstantInt::get(IRB.getInt64Ty(), Size)});
 132 }
 133
 134 void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
 135                                       uint64_t Size) {
 136   IRBuilder<> IRB(InsertBefore);
 137   IRB.CreateCall(SetTagFunc, {IRB.CreatePointerCast(AI, IRB.getInt8PtrTy()),
 138                               ConstantInt::get(IRB.getInt64Ty(), Size)});
 139 }
 140
 141 Instruction *AArch64StackTagging::insertBaseTaggedPointer(
 142     const MapVector<AllocaInst *, AllocaInfo> &Allocas,
 143     const DominatorTree *DT) {
 144   BasicBlock *PrologueBB = nullptr;
 145   // Try sinking IRG as deep as possible to avoid hurting shrink wrap.
 146   for (auto &I : Allocas) {
 147     const AllocaInfo &Info = I.second;
 148     AllocaInst *AI = Info.AI;
 149     if (Info.Tag < 0)
 150       continue;
 151     if (!PrologueBB) {
 152       PrologueBB = AI->getParent();
 153       continue;
 154     }
 155     PrologueBB = DT->findNearestCommonDominator(PrologueBB, AI->getParent());
 156   }
 157   assert(PrologueBB);
 158
 159   IRBuilder<> IRB(&PrologueBB->front());
 160   Function *IRG_SP =
 161       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_irg_sp);
 162   Instruction *Base =
 163       IRB.CreateCall(IRG_SP, {Constant::getNullValue(IRB.getInt64Ty())});
 164   Base->setName("basetag");
 165   return Base;
 166 }
 167
 168 void AArch64StackTagging::alignAndPadAlloca(AllocaInfo &Info) {
 169   unsigned NewAlignment = std::max(Info.AI->getAlignment(), kTagGranuleSize);
 170   Info.AI->setAlignment(NewAlignment);
 171
 172   uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
 173   uint64_t AlignedSize = alignTo(Size, kTagGranuleSize);
 174   if (Size == AlignedSize)
 175     return;
 176
 177   // Add padding to the alloca.
 178   Type *AllocatedType =
 179       Info.AI->isArrayAllocation()
 180           ? ArrayType::get(
 181                 Info.AI->getAllocatedType(),
 182                 dyn_cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue())
 183           : Info.AI->getAllocatedType();
 184   Type *PaddingType =
 185       ArrayType::get(Type::getInt8Ty(F->getContext()), AlignedSize - Size);
 186   Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType);
 187   auto *NewAI = new AllocaInst(
 188       TypeWithPadding, Info.AI->getType()->getAddressSpace(), nullptr, "", Info.AI);
 189   NewAI->takeName(Info.AI);
 190   NewAI->setAlignment(Info.AI->getAlignment());
 191   NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca());
 192   NewAI->setSwiftError(Info.AI->isSwiftError());
 193   NewAI->copyMetadata(*Info.AI);
 194
 195   auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI);
 196   Info.AI->replaceAllUsesWith(NewPtr);
 197   Info.AI->eraseFromParent();
 198   Info.AI = NewAI;
 199 }
 200
 201 // FIXME: check for MTE extension
 202 bool AArch64StackTagging::runOnFunction(Function &Fn) {
 203   if (!Fn.hasFnAttribute(Attribute::SanitizeMemTag))
 204     return false;
 205
 206   F = &Fn;
 207   DL = &Fn.getParent()->getDataLayout();
 208
 209   MapVector<AllocaInst *, AllocaInfo> Allocas; // need stable iteration order
 210   SmallVector<Instruction *, 8> RetVec;
 211   DenseMap<Value *, AllocaInst *> AllocaForValue;
 212   SmallVector<Instruction *, 4> UnrecognizedLifetimes;
 213
 214   for (auto &BB : *F) {
 215     for (BasicBlock::iterator IT = BB.begin(); IT != BB.end(); ++IT) {
 216       Instruction *I = &*IT;
 217       if (auto *AI = dyn_cast<AllocaInst>(I)) {
 218         Allocas[AI].AI = AI;
 219         continue;
 220       }
 221
 222       if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(I)) {
 223         if (auto *AI =
 224                 dyn_cast_or_null<AllocaInst>(DVI->getVariableLocation())) {
 225           Allocas[AI].DbgVariableIntrinsics.push_back(DVI);
 226         }
 227         continue;
 228       }
 229
 230       auto *II = dyn_cast<IntrinsicInst>(I);
 231       if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
 232                  II->getIntrinsicID() == Intrinsic::lifetime_end)) {
 233         AllocaInst *AI =
 234             llvm::findAllocaForValue(II->getArgOperand(1), AllocaForValue);
 235         if (!AI) {
 236           UnrecognizedLifetimes.push_back(I);
 237           continue;
 238         }
 239         if (II->getIntrinsicID() == Intrinsic::lifetime_start)
 240           Allocas[AI].LifetimeStart.push_back(II);
 241         else
 242           Allocas[AI].LifetimeEnd.push_back(II);
 243       }
 244
 245       if (isa<ReturnInst>(I) || isa<ResumeInst>(I) || isa<CleanupReturnInst>(I))
 246         RetVec.push_back(I);
 247     }
 248   }
 249
 250   if (Allocas.empty())
 251     return false;
 252
 253   int NextTag = 0;
 254   int NumInterestingAllocas = 0;
 255   for (auto &I : Allocas) {
 256     AllocaInfo &Info = I.second;
 257     assert(Info.AI);
 258
 259     if (!isInterestingAlloca(*Info.AI)) {
 260       Info.Tag = -1;
 261       continue;
 262     }
 263
 264     alignAndPadAlloca(Info);
 265     NumInterestingAllocas++;
 266     Info.Tag = NextTag;
 267     NextTag = (NextTag + 1) % 16;
 268   }
 269
 270   if (NumInterestingAllocas == 0)
 271     return true;
 272
 273   SetTagFunc =
 274       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 275
 276   // Compute DT only if the function has the attribute, there are more than 1
 277   // interesting allocas, and it is not available for free.
 278   Instruction *Base;
 279   if (NumInterestingAllocas > 1) {
 280     auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
 281     if (DTWP) {
 282       Base = insertBaseTaggedPointer(Allocas, &DTWP->getDomTree());
 283     } else {
 284       DominatorTree DT(*F);
 285       Base = insertBaseTaggedPointer(Allocas, &DT);
 286     }
 287   } else {
 288     Base = insertBaseTaggedPointer(Allocas, nullptr);
 289   }
 290
 291   for (auto &I : Allocas) {
 292     const AllocaInfo &Info = I.second;
 293     AllocaInst *AI = Info.AI;
 294     if (Info.Tag < 0)
 295       continue;
 296
 297     // Replace alloca with tagp(alloca).
 298     IRBuilder<> IRB(Info.AI->getNextNode());
 299     Function *TagP = Intrinsic::getDeclaration(
 300         F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()});
 301     Instruction *TagPCall =
 302         IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base,
 303                               ConstantInt::get(IRB.getInt64Ty(), Info.Tag)});
 304     if (Info.AI->hasName())
 305       TagPCall->setName(Info.AI->getName() + ".tag");
 306     Info.AI->replaceAllUsesWith(TagPCall);
 307     TagPCall->setOperand(0, Info.AI);
 308
 309     if (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 &&
 310         Info.LifetimeEnd.size() == 1) {
 311       IntrinsicInst *Start = Info.LifetimeStart[0];
 312       uint64_t Size =
 313           dyn_cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
 314       Size = alignTo(Size, kTagGranuleSize);
 315       tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size);
 316       untagAlloca(AI, Info.LifetimeEnd[0], Size);
 317     } else {
 318       uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
 319       Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
 320       tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);
 321       for (auto &RI : RetVec) {
 322         untagAlloca(AI, RI, Size);
 323       }
 324       // We may have inserted tag/untag outside of any lifetime interval.
 325       // Remove all lifetime intrinsics for this alloca.
 326       for (auto &II : Info.LifetimeStart)
 327         II->eraseFromParent();
 328       for (auto &II : Info.LifetimeEnd)
 329         II->eraseFromParent();
 330     }
 331
 332     // Fixup debug intrinsics to point to the new alloca.
 333     for (auto DVI : Info.DbgVariableIntrinsics)
 334       DVI->setArgOperand(
 335           0,
 336           MetadataAsValue::get(F->getContext(), LocalAsMetadata::get(Info.AI)));
 337   }
 338
 339   // If we have instrumented at least one alloca, all unrecognized lifetime
 340   // instrinsics have to go.
 341   for (auto &I : UnrecognizedLifetimes)
 342     I->eraseFromParent();
 343
 344   return true;
 345 }