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
AArch64StackTagging()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
getPassName() const89 StringRef getPassName() const override { return "AArch64 Stack Tagging"; }
90
91 private:
92 Function *F;
93 Function *SetTagFunc;
94 const DataLayout *DL;
95
getAnalysisUsage(AnalysisUsage & AU) const96 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
createAArch64StackTaggingPass()110 FunctionPass *llvm::createAArch64StackTaggingPass() {
111 return new AArch64StackTagging();
112 }
113
isInterestingAlloca(const AllocaInst & AI)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
tagAlloca(AllocaInst * AI,Instruction * InsertBefore,Value * Ptr,uint64_t Size)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
untagAlloca(AllocaInst * AI,Instruction * InsertBefore,uint64_t Size)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
insertBaseTaggedPointer(const MapVector<AllocaInst *,AllocaInfo> & Allocas,const DominatorTree * DT)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
alignAndPadAlloca(AllocaInfo & Info)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
runOnFunction(Function & Fn)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 }
346