LLVM  10.0.0svn
AArch64StackTagging.cpp
Go to the documentation of this file.
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"
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"
37 #include "llvm/IR/DebugLoc.h"
38 #include "llvm/IR/Dominators.h"
39 #include "llvm/IR/Function.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"
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;
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) {
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 
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(
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
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
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());
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 }
A parsed version of the target data layout string in and methods for querying it. ...
Definition: DataLayout.h:111
static PassRegistry * getPassRegistry()
getPassRegistry - Access the global registry object, which is automatically initialized at applicatio...
GCNRegPressure max(const GCNRegPressure &P1, const GCNRegPressure &P2)
NodeTy * getNextNode()
Get the next node, or nullptr for the list tail.
Definition: ilist_node.h:288
This class represents lattice values for constants.
Definition: AllocatorList.h:23
bool isSized(SmallPtrSetImpl< Type *> *Visited=nullptr) const
Return true if it makes sense to take the size of this type.
Definition: Type.h:264
This file contains the declarations for metadata subclasses.
bool isSwiftError() const
Return true if this alloca is used as a swifterror argument to a call.
Definition: Instructions.h:135
bool hasFnAttribute(Attribute::AttrKind Kind) const
Return true if the function has the attribute.
Definition: Function.h:323
NodeT * findNearestCommonDominator(NodeT *A, NodeT *B) const
findNearestCommonDominator - Find nearest common dominator basic block for basic block A and B...
This class implements a map that also provides access to all stored values in a deterministic order...
Definition: MapVector.h:37
F(f)
#define DEBUG_TYPE
static Constant * getNullValue(Type *Ty)
Constructor to create a &#39;0&#39; constant of arbitrary type.
Definition: Constants.cpp:274
AArch64 Stack Tagging
Value * getArgOperand(unsigned i) const
Definition: InstrTypes.h:1241
const DataLayout & getDataLayout() const
Get the data layout for the module&#39;s target platform.
Definition: Module.cpp:369
IntegerType * getInt64Ty()
Fetch the type representing a 64-bit integer.
Definition: IRBuilder.h:388
This provides a uniform API for creating instructions and inserting them into a basic block: either a...
Definition: IRBuilder.h:779
bool empty() const
Definition: MapVector.h:79
void setName(const Twine &Name)
Change the name of the value.
Definition: Value.cpp:285
static StructType * get(LLVMContext &Context, ArrayRef< Type *> Elements, bool isPacked=false)
This static method is the primary way to create a literal StructType.
Definition: Type.cpp:341
const T & getValue() const LLVM_LVALUE_FUNCTION
Definition: Optional.h:255
bool isUsedWithInAlloca() const
Return true if this alloca is used as an inalloca argument to a call.
Definition: Instructions.h:124
This class represents a no-op cast from one type to another.
void takeName(Value *V)
Transfer the name from V to this value.
Definition: Value.cpp:291
Concrete subclass of DominatorTreeBase that is used to compute a normal dominator tree...
Definition: Dominators.h:144
Function * getDeclaration(Module *M, ID id, ArrayRef< Type *> Tys=None)
Create or insert an LLVM Function declaration for an intrinsic, and return it.
Definition: Function.cpp:1043
Analysis containing CSE Info
Definition: CSEInfo.cpp:20
static MetadataAsValue * get(LLVMContext &Context, Metadata *MD)
Definition: Metadata.cpp:105
static bool runOnFunction(Function &F, bool PostInlining)
LLVM Basic Block Representation.
Definition: BasicBlock.h:57
The instances of the Type class are immutable: once they are created, they are never changed...
Definition: Type.h:45
const Instruction & front() const
Definition: BasicBlock.h:280
Represent the analysis usage information of a pass.
AllocaInst * findAllocaForValue(Value *V, DenseMap< Value *, AllocaInst *> &AllocaForValue)
Finds alloca where the value comes from.
FunctionPass class - This class is used to implement most global optimizations.
Definition: Pass.h:284
A set of register units.
FunctionPass * createAArch64StackTaggingPass()
static LocalAsMetadata * get(Value *Local)
Definition: Metadata.h:435
PointerType * getInt8PtrTy(unsigned AddrSpace=0)
Fetch the type representing a pointer to an 8-bit integer value.
Definition: IRBuilder.h:421
INITIALIZE_PASS_END(RegBankSelect, DEBUG_TYPE, "Assign register bank of generic virtual registers", false, false) RegBankSelect
Type * getAllocatedType() const
Return the type that is being allocated by the instruction.
Definition: Instructions.h:105
Optional< uint64_t > getAllocationSizeInBits(const DataLayout &DL) const
Get allocation size in bits.
Iterator for intrusive lists based on ilist_node.
This is the shared class of boolean and integer constants.
Definition: Constants.h:83
This is a &#39;vector&#39; (really, a variable-sized array), optimized for the case when the array is small...
Definition: SmallVector.h:837
static Constant * get(Type *Ty, uint64_t V, bool isSigned=false)
If Ty is a vector type, return a Constant with a splat of the given value.
Definition: Constants.cpp:640
void setPreservesCFG()
This function should be called by the pass, iff they do not:
Definition: Pass.cpp:301
void initializeAArch64StackTaggingPass(PassRegistry &)
static constexpr unsigned kTagGranuleSize
Value * CreatePointerCast(Value *V, Type *DestTy, const Twine &Name="")
Definition: IRBuilder.h:2004
static cl::opt< ITMode > IT(cl::desc("IT block support"), cl::Hidden, cl::init(DefaultIT), cl::ZeroOrMore, cl::values(clEnumValN(DefaultIT, "arm-default-it", "Generate IT block based on arch"), clEnumValN(RestrictedIT, "arm-restrict-it", "Disallow deprecated IT based on ARMv8"), clEnumValN(NoRestrictedIT, "arm-no-restrict-it", "Allow IT blocks based on ARMv7")))
uint64_t alignTo(uint64_t Size, Align A)
Returns a multiple of A needed to store Size bytes.
Definition: Alignment.h:126
INITIALIZE_PASS_BEGIN(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging", false, false) INITIALIZE_PASS_END(AArch64StackTagging
LLVM_NODISCARD bool empty() const
Definition: SmallVector.h:55
#define I(x, y, z)
Definition: MD5.cpp:58
static ArrayType * get(Type *ElementType, uint64_t NumElements)
This static method is the primary way to construct an ArrayType.
Definition: Type.cpp:582
LLVM_NODISCARD std::enable_if<!is_simple_type< Y >::value, typename cast_retty< X, const Y >::ret_type >::type dyn_cast(const Y &Val)
Definition: Casting.h:332
uint32_t Size
Definition: Profile.cpp:46
CallInst * CreateCall(FunctionType *FTy, Value *Callee, ArrayRef< Value *> Args=None, const Twine &Name="", MDNode *FPMathTag=nullptr)
Definition: IRBuilder.h:2223
assert(ImpDefSCC.getReg()==AMDGPU::SCC &&ImpDefSCC.isDef())
Module * getParent()
Get the module that this global value is contained inside of...
Definition: GlobalValue.h:575
LLVM Value Representation.
Definition: Value.h:73
StringRef - Represent a constant reference to a string, i.e.
Definition: StringRef.h:48
bool isStaticAlloca() const
Return true if this alloca is in the entry block of the function and is a constant size...
static IntegerType * getInt8Ty(LLVMContext &C)
Definition: Type.cpp:173
A wrapper class for inspecting calls to intrinsic functions.
Definition: IntrinsicInst.h:43
const BasicBlock * getParent() const
Definition: Instruction.h:66
an instruction to allocate memory on the stack
Definition: Instructions.h:59