File: | llvm/lib/Target/AArch64/AArch64StackTagging.cpp |
Warning: | line 480, column 20 Called C++ object pointer is null |
Press '?' to see keyboard shortcuts
Keyboard shortcuts:
1 | //===- AArch64StackTagging.cpp - Stack tagging in IR --===// | |||
2 | // | |||
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | |||
4 | // See https://llvm.org/LICENSE.txt for license information. | |||
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |||
6 | // | |||
7 | //===----------------------------------------------------------------------===// | |||
8 | //===----------------------------------------------------------------------===// | |||
9 | ||||
10 | #include "AArch64.h" | |||
11 | #include "AArch64InstrInfo.h" | |||
12 | #include "AArch64Subtarget.h" | |||
13 | #include "AArch64TargetMachine.h" | |||
14 | #include "llvm/ADT/DenseMap.h" | |||
15 | #include "llvm/ADT/DepthFirstIterator.h" | |||
16 | #include "llvm/ADT/MapVector.h" | |||
17 | #include "llvm/ADT/None.h" | |||
18 | #include "llvm/ADT/Optional.h" | |||
19 | #include "llvm/ADT/SmallVector.h" | |||
20 | #include "llvm/ADT/Statistic.h" | |||
21 | #include "llvm/Analysis/AliasAnalysis.h" | |||
22 | #include "llvm/Analysis/CFG.h" | |||
23 | #include "llvm/Analysis/LoopInfo.h" | |||
24 | #include "llvm/Analysis/PostDominators.h" | |||
25 | #include "llvm/Analysis/ScalarEvolution.h" | |||
26 | #include "llvm/Analysis/ScalarEvolutionExpressions.h" | |||
27 | #include "llvm/Analysis/StackSafetyAnalysis.h" | |||
28 | #include "llvm/Analysis/ValueTracking.h" | |||
29 | #include "llvm/CodeGen/LiveRegUnits.h" | |||
30 | #include "llvm/CodeGen/MachineBasicBlock.h" | |||
31 | #include "llvm/CodeGen/MachineFunction.h" | |||
32 | #include "llvm/CodeGen/MachineFunctionPass.h" | |||
33 | #include "llvm/CodeGen/MachineInstr.h" | |||
34 | #include "llvm/CodeGen/MachineInstrBuilder.h" | |||
35 | #include "llvm/CodeGen/MachineLoopInfo.h" | |||
36 | #include "llvm/CodeGen/MachineOperand.h" | |||
37 | #include "llvm/CodeGen/MachineRegisterInfo.h" | |||
38 | #include "llvm/CodeGen/TargetPassConfig.h" | |||
39 | #include "llvm/CodeGen/TargetRegisterInfo.h" | |||
40 | #include "llvm/IR/DebugLoc.h" | |||
41 | #include "llvm/IR/Dominators.h" | |||
42 | #include "llvm/IR/Function.h" | |||
43 | #include "llvm/IR/GetElementPtrTypeIterator.h" | |||
44 | #include "llvm/IR/IRBuilder.h" | |||
45 | #include "llvm/IR/Instruction.h" | |||
46 | #include "llvm/IR/Instructions.h" | |||
47 | #include "llvm/IR/IntrinsicInst.h" | |||
48 | #include "llvm/IR/IntrinsicsAArch64.h" | |||
49 | #include "llvm/IR/Metadata.h" | |||
50 | #include "llvm/InitializePasses.h" | |||
51 | #include "llvm/Pass.h" | |||
52 | #include "llvm/Support/Casting.h" | |||
53 | #include "llvm/Support/Debug.h" | |||
54 | #include "llvm/Support/raw_ostream.h" | |||
55 | #include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h" | |||
56 | #include "llvm/Transforms/Utils/Local.h" | |||
57 | #include <cassert> | |||
58 | #include <iterator> | |||
59 | #include <utility> | |||
60 | ||||
61 | using namespace llvm; | |||
62 | ||||
63 | #define DEBUG_TYPE"aarch64-stack-tagging" "aarch64-stack-tagging" | |||
64 | ||||
65 | static cl::opt<bool> ClMergeInit( | |||
66 | "stack-tagging-merge-init", cl::Hidden, cl::init(true), cl::ZeroOrMore, | |||
67 | cl::desc("merge stack variable initializers with tagging when possible")); | |||
68 | ||||
69 | static cl::opt<bool> | |||
70 | ClUseStackSafety("stack-tagging-use-stack-safety", cl::Hidden, | |||
71 | cl::init(true), cl::ZeroOrMore, | |||
72 | cl::desc("Use Stack Safety analysis results")); | |||
73 | ||||
74 | static cl::opt<unsigned> ClScanLimit("stack-tagging-merge-init-scan-limit", | |||
75 | cl::init(40), cl::Hidden); | |||
76 | ||||
77 | static cl::opt<unsigned> | |||
78 | ClMergeInitSizeLimit("stack-tagging-merge-init-size-limit", cl::init(272), | |||
79 | cl::Hidden); | |||
80 | ||||
81 | static const Align kTagGranuleSize = Align(16); | |||
82 | ||||
83 | namespace { | |||
84 | ||||
85 | class InitializerBuilder { | |||
86 | uint64_t Size; | |||
87 | const DataLayout *DL; | |||
88 | Value *BasePtr; | |||
89 | Function *SetTagFn; | |||
90 | Function *SetTagZeroFn; | |||
91 | Function *StgpFn; | |||
92 | ||||
93 | // List of initializers sorted by start offset. | |||
94 | struct Range { | |||
95 | uint64_t Start, End; | |||
96 | Instruction *Inst; | |||
97 | }; | |||
98 | SmallVector<Range, 4> Ranges; | |||
99 | // 8-aligned offset => 8-byte initializer | |||
100 | // Missing keys are zero initialized. | |||
101 | std::map<uint64_t, Value *> Out; | |||
102 | ||||
103 | public: | |||
104 | InitializerBuilder(uint64_t Size, const DataLayout *DL, Value *BasePtr, | |||
105 | Function *SetTagFn, Function *SetTagZeroFn, | |||
106 | Function *StgpFn) | |||
107 | : Size(Size), DL(DL), BasePtr(BasePtr), SetTagFn(SetTagFn), | |||
108 | SetTagZeroFn(SetTagZeroFn), StgpFn(StgpFn) {} | |||
109 | ||||
110 | bool addRange(uint64_t Start, uint64_t End, Instruction *Inst) { | |||
111 | auto I = | |||
112 | llvm::lower_bound(Ranges, Start, [](const Range &LHS, uint64_t RHS) { | |||
113 | return LHS.End <= RHS; | |||
114 | }); | |||
115 | if (I != Ranges.end() && End > I->Start) { | |||
116 | // Overlap - bail. | |||
117 | return false; | |||
118 | } | |||
119 | Ranges.insert(I, {Start, End, Inst}); | |||
120 | return true; | |||
121 | } | |||
122 | ||||
123 | bool addStore(uint64_t Offset, StoreInst *SI, const DataLayout *DL) { | |||
124 | int64_t StoreSize = DL->getTypeStoreSize(SI->getOperand(0)->getType()); | |||
125 | if (!addRange(Offset, Offset + StoreSize, SI)) | |||
126 | return false; | |||
127 | IRBuilder<> IRB(SI); | |||
128 | applyStore(IRB, Offset, Offset + StoreSize, SI->getOperand(0)); | |||
129 | return true; | |||
130 | } | |||
131 | ||||
132 | bool addMemSet(uint64_t Offset, MemSetInst *MSI) { | |||
133 | uint64_t StoreSize = cast<ConstantInt>(MSI->getLength())->getZExtValue(); | |||
134 | if (!addRange(Offset, Offset + StoreSize, MSI)) | |||
135 | return false; | |||
136 | IRBuilder<> IRB(MSI); | |||
137 | applyMemSet(IRB, Offset, Offset + StoreSize, | |||
138 | cast<ConstantInt>(MSI->getValue())); | |||
139 | return true; | |||
140 | } | |||
141 | ||||
142 | void applyMemSet(IRBuilder<> &IRB, int64_t Start, int64_t End, | |||
143 | ConstantInt *V) { | |||
144 | // Out[] does not distinguish between zero and undef, and we already know | |||
145 | // that this memset does not overlap with any other initializer. Nothing to | |||
146 | // do for memset(0). | |||
147 | if (V->isZero()) | |||
148 | return; | |||
149 | for (int64_t Offset = Start - Start % 8; Offset < End; Offset += 8) { | |||
150 | uint64_t Cst = 0x0101010101010101UL; | |||
151 | int LowBits = Offset < Start ? (Start - Offset) * 8 : 0; | |||
152 | if (LowBits) | |||
153 | Cst = (Cst >> LowBits) << LowBits; | |||
154 | int HighBits = End - Offset < 8 ? (8 - (End - Offset)) * 8 : 0; | |||
155 | if (HighBits) | |||
156 | Cst = (Cst << HighBits) >> HighBits; | |||
157 | ConstantInt *C = | |||
158 | ConstantInt::get(IRB.getInt64Ty(), Cst * V->getZExtValue()); | |||
159 | ||||
160 | Value *&CurrentV = Out[Offset]; | |||
161 | if (!CurrentV) { | |||
162 | CurrentV = C; | |||
163 | } else { | |||
164 | CurrentV = IRB.CreateOr(CurrentV, C); | |||
165 | } | |||
166 | } | |||
167 | } | |||
168 | ||||
169 | // Take a 64-bit slice of the value starting at the given offset (in bytes). | |||
170 | // Offset can be negative. Pad with zeroes on both sides when necessary. | |||
171 | Value *sliceValue(IRBuilder<> &IRB, Value *V, int64_t Offset) { | |||
172 | if (Offset > 0) { | |||
173 | V = IRB.CreateLShr(V, Offset * 8); | |||
174 | V = IRB.CreateZExtOrTrunc(V, IRB.getInt64Ty()); | |||
175 | } else if (Offset < 0) { | |||
176 | V = IRB.CreateZExtOrTrunc(V, IRB.getInt64Ty()); | |||
177 | V = IRB.CreateShl(V, -Offset * 8); | |||
178 | } else { | |||
179 | V = IRB.CreateZExtOrTrunc(V, IRB.getInt64Ty()); | |||
180 | } | |||
181 | return V; | |||
182 | } | |||
183 | ||||
184 | void applyStore(IRBuilder<> &IRB, int64_t Start, int64_t End, | |||
185 | Value *StoredValue) { | |||
186 | StoredValue = flatten(IRB, StoredValue); | |||
187 | for (int64_t Offset = Start - Start % 8; Offset < End; Offset += 8) { | |||
188 | Value *V = sliceValue(IRB, StoredValue, Offset - Start); | |||
189 | Value *&CurrentV = Out[Offset]; | |||
190 | if (!CurrentV) { | |||
191 | CurrentV = V; | |||
192 | } else { | |||
193 | CurrentV = IRB.CreateOr(CurrentV, V); | |||
194 | } | |||
195 | } | |||
196 | } | |||
197 | ||||
198 | void generate(IRBuilder<> &IRB) { | |||
199 | LLVM_DEBUG(dbgs() << "Combined initializer\n")do { } while (false); | |||
200 | // No initializers => the entire allocation is undef. | |||
201 | if (Ranges.empty()) { | |||
202 | emitUndef(IRB, 0, Size); | |||
203 | return; | |||
204 | } | |||
205 | ||||
206 | // Look through 8-byte initializer list 16 bytes at a time; | |||
207 | // If one of the two 8-byte halfs is non-zero non-undef, emit STGP. | |||
208 | // Otherwise, emit zeroes up to next available item. | |||
209 | uint64_t LastOffset = 0; | |||
210 | for (uint64_t Offset = 0; Offset < Size; Offset += 16) { | |||
211 | auto I1 = Out.find(Offset); | |||
212 | auto I2 = Out.find(Offset + 8); | |||
213 | if (I1 == Out.end() && I2 == Out.end()) | |||
214 | continue; | |||
215 | ||||
216 | if (Offset > LastOffset) | |||
217 | emitZeroes(IRB, LastOffset, Offset - LastOffset); | |||
218 | ||||
219 | Value *Store1 = I1 == Out.end() ? Constant::getNullValue(IRB.getInt64Ty()) | |||
220 | : I1->second; | |||
221 | Value *Store2 = I2 == Out.end() ? Constant::getNullValue(IRB.getInt64Ty()) | |||
222 | : I2->second; | |||
223 | emitPair(IRB, Offset, Store1, Store2); | |||
224 | LastOffset = Offset + 16; | |||
225 | } | |||
226 | ||||
227 | // memset(0) does not update Out[], therefore the tail can be either undef | |||
228 | // or zero. | |||
229 | if (LastOffset < Size) | |||
230 | emitZeroes(IRB, LastOffset, Size - LastOffset); | |||
231 | ||||
232 | for (const auto &R : Ranges) { | |||
233 | R.Inst->eraseFromParent(); | |||
234 | } | |||
235 | } | |||
236 | ||||
237 | void emitZeroes(IRBuilder<> &IRB, uint64_t Offset, uint64_t Size) { | |||
238 | LLVM_DEBUG(dbgs() << " [" << Offset << ", " << Offset + Sizedo { } while (false) | |||
239 | << ") zero\n")do { } while (false); | |||
240 | Value *Ptr = BasePtr; | |||
241 | if (Offset) | |||
242 | Ptr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), Ptr, Offset); | |||
243 | IRB.CreateCall(SetTagZeroFn, | |||
244 | {Ptr, ConstantInt::get(IRB.getInt64Ty(), Size)}); | |||
245 | } | |||
246 | ||||
247 | void emitUndef(IRBuilder<> &IRB, uint64_t Offset, uint64_t Size) { | |||
248 | LLVM_DEBUG(dbgs() << " [" << Offset << ", " << Offset + Sizedo { } while (false) | |||
249 | << ") undef\n")do { } while (false); | |||
250 | Value *Ptr = BasePtr; | |||
251 | if (Offset) | |||
252 | Ptr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), Ptr, Offset); | |||
253 | IRB.CreateCall(SetTagFn, {Ptr, ConstantInt::get(IRB.getInt64Ty(), Size)}); | |||
254 | } | |||
255 | ||||
256 | void emitPair(IRBuilder<> &IRB, uint64_t Offset, Value *A, Value *B) { | |||
257 | LLVM_DEBUG(dbgs() << " [" << Offset << ", " << Offset + 16 << "):\n")do { } while (false); | |||
258 | LLVM_DEBUG(dbgs() << " " << *A << "\n " << *B << "\n")do { } while (false); | |||
259 | Value *Ptr = BasePtr; | |||
260 | if (Offset) | |||
261 | Ptr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), Ptr, Offset); | |||
262 | IRB.CreateCall(StgpFn, {Ptr, A, B}); | |||
263 | } | |||
264 | ||||
265 | Value *flatten(IRBuilder<> &IRB, Value *V) { | |||
266 | if (V->getType()->isIntegerTy()) | |||
267 | return V; | |||
268 | // vector of pointers -> vector of ints | |||
269 | if (VectorType *VecTy = dyn_cast<VectorType>(V->getType())) { | |||
270 | LLVMContext &Ctx = IRB.getContext(); | |||
271 | Type *EltTy = VecTy->getElementType(); | |||
272 | if (EltTy->isPointerTy()) { | |||
273 | uint32_t EltSize = DL->getTypeSizeInBits(EltTy); | |||
274 | auto *NewTy = FixedVectorType::get( | |||
275 | IntegerType::get(Ctx, EltSize), | |||
276 | cast<FixedVectorType>(VecTy)->getNumElements()); | |||
277 | V = IRB.CreatePointerCast(V, NewTy); | |||
278 | } | |||
279 | } | |||
280 | return IRB.CreateBitOrPointerCast( | |||
281 | V, IRB.getIntNTy(DL->getTypeStoreSize(V->getType()) * 8)); | |||
282 | } | |||
283 | }; | |||
284 | ||||
285 | class AArch64StackTagging : public FunctionPass { | |||
286 | struct AllocaInfo { | |||
287 | AllocaInst *AI; | |||
288 | TrackingVH<Instruction> OldAI; // Track through RAUW to replace debug uses. | |||
289 | SmallVector<IntrinsicInst *, 2> LifetimeStart; | |||
290 | SmallVector<IntrinsicInst *, 2> LifetimeEnd; | |||
291 | SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics; | |||
292 | int Tag; // -1 for non-tagged allocations | |||
293 | }; | |||
294 | ||||
295 | const bool MergeInit; | |||
296 | const bool UseStackSafety; | |||
297 | ||||
298 | public: | |||
299 | static char ID; // Pass ID, replacement for typeid | |||
300 | ||||
301 | AArch64StackTagging(bool IsOptNone = false) | |||
302 | : FunctionPass(ID), | |||
303 | MergeInit(ClMergeInit.getNumOccurrences() ? ClMergeInit : !IsOptNone), | |||
304 | UseStackSafety(ClUseStackSafety.getNumOccurrences() ? ClUseStackSafety | |||
305 | : !IsOptNone) { | |||
306 | initializeAArch64StackTaggingPass(*PassRegistry::getPassRegistry()); | |||
307 | } | |||
308 | ||||
309 | bool isInterestingAlloca(const AllocaInst &AI); | |||
310 | void alignAndPadAlloca(AllocaInfo &Info); | |||
311 | ||||
312 | void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr, | |||
313 | uint64_t Size); | |||
314 | void untagAlloca(AllocaInst *AI, Instruction *InsertBefore, uint64_t Size); | |||
315 | ||||
316 | Instruction *collectInitializers(Instruction *StartInst, Value *StartPtr, | |||
317 | uint64_t Size, InitializerBuilder &IB); | |||
318 | ||||
319 | Instruction * | |||
320 | insertBaseTaggedPointer(const MapVector<AllocaInst *, AllocaInfo> &Allocas, | |||
321 | const DominatorTree *DT); | |||
322 | bool runOnFunction(Function &F) override; | |||
323 | ||||
324 | StringRef getPassName() const override { return "AArch64 Stack Tagging"; } | |||
325 | ||||
326 | private: | |||
327 | Function *F = nullptr; | |||
328 | Function *SetTagFunc = nullptr; | |||
329 | const DataLayout *DL = nullptr; | |||
330 | AAResults *AA = nullptr; | |||
331 | const StackSafetyGlobalInfo *SSI = nullptr; | |||
332 | ||||
333 | void getAnalysisUsage(AnalysisUsage &AU) const override { | |||
334 | AU.setPreservesCFG(); | |||
335 | if (UseStackSafety) | |||
336 | AU.addRequired<StackSafetyGlobalInfoWrapperPass>(); | |||
337 | if (MergeInit) | |||
338 | AU.addRequired<AAResultsWrapperPass>(); | |||
339 | } | |||
340 | }; | |||
341 | ||||
342 | } // end anonymous namespace | |||
343 | ||||
344 | char AArch64StackTagging::ID = 0; | |||
345 | ||||
346 | INITIALIZE_PASS_BEGIN(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",static void *initializeAArch64StackTaggingPassOnce(PassRegistry &Registry) { | |||
347 | false, false)static void *initializeAArch64StackTaggingPassOnce(PassRegistry &Registry) { | |||
348 | INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)initializeAAResultsWrapperPassPass(Registry); | |||
349 | INITIALIZE_PASS_DEPENDENCY(StackSafetyGlobalInfoWrapperPass)initializeStackSafetyGlobalInfoWrapperPassPass(Registry); | |||
350 | INITIALIZE_PASS_END(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",PassInfo *PI = new PassInfo( "AArch64 Stack Tagging", "aarch64-stack-tagging" , &AArch64StackTagging::ID, PassInfo::NormalCtor_t(callDefaultCtor <AArch64StackTagging>), false, false); Registry.registerPass (*PI, true); return PI; } static llvm::once_flag InitializeAArch64StackTaggingPassFlag ; void llvm::initializeAArch64StackTaggingPass(PassRegistry & Registry) { llvm::call_once(InitializeAArch64StackTaggingPassFlag , initializeAArch64StackTaggingPassOnce, std::ref(Registry)); } | |||
351 | false, false)PassInfo *PI = new PassInfo( "AArch64 Stack Tagging", "aarch64-stack-tagging" , &AArch64StackTagging::ID, PassInfo::NormalCtor_t(callDefaultCtor <AArch64StackTagging>), false, false); Registry.registerPass (*PI, true); return PI; } static llvm::once_flag InitializeAArch64StackTaggingPassFlag ; void llvm::initializeAArch64StackTaggingPass(PassRegistry & Registry) { llvm::call_once(InitializeAArch64StackTaggingPassFlag , initializeAArch64StackTaggingPassOnce, std::ref(Registry)); } | |||
352 | ||||
353 | FunctionPass *llvm::createAArch64StackTaggingPass(bool IsOptNone) { | |||
354 | return new AArch64StackTagging(IsOptNone); | |||
355 | } | |||
356 | ||||
357 | Instruction *AArch64StackTagging::collectInitializers(Instruction *StartInst, | |||
358 | Value *StartPtr, | |||
359 | uint64_t Size, | |||
360 | InitializerBuilder &IB) { | |||
361 | MemoryLocation AllocaLoc{StartPtr, Size}; | |||
362 | Instruction *LastInst = StartInst; | |||
363 | BasicBlock::iterator BI(StartInst); | |||
364 | ||||
365 | unsigned Count = 0; | |||
366 | for (; Count < ClScanLimit && !BI->isTerminator(); ++BI) { | |||
367 | if (!isa<DbgInfoIntrinsic>(*BI)) | |||
368 | ++Count; | |||
369 | ||||
370 | if (isNoModRef(AA->getModRefInfo(&*BI, AllocaLoc))) | |||
371 | continue; | |||
372 | ||||
373 | if (!isa<StoreInst>(BI) && !isa<MemSetInst>(BI)) { | |||
374 | // If the instruction is readnone, ignore it, otherwise bail out. We | |||
375 | // don't even allow readonly here because we don't want something like: | |||
376 | // A[1] = 2; strlen(A); A[2] = 2; -> memcpy(A, ...); strlen(A). | |||
377 | if (BI->mayWriteToMemory() || BI->mayReadFromMemory()) | |||
378 | break; | |||
379 | continue; | |||
380 | } | |||
381 | ||||
382 | if (StoreInst *NextStore = dyn_cast<StoreInst>(BI)) { | |||
383 | if (!NextStore->isSimple()) | |||
384 | break; | |||
385 | ||||
386 | // Check to see if this store is to a constant offset from the start ptr. | |||
387 | Optional<int64_t> Offset = | |||
388 | isPointerOffset(StartPtr, NextStore->getPointerOperand(), *DL); | |||
389 | if (!Offset) | |||
390 | break; | |||
391 | ||||
392 | if (!IB.addStore(*Offset, NextStore, DL)) | |||
393 | break; | |||
394 | LastInst = NextStore; | |||
395 | } else { | |||
396 | MemSetInst *MSI = cast<MemSetInst>(BI); | |||
397 | ||||
398 | if (MSI->isVolatile() || !isa<ConstantInt>(MSI->getLength())) | |||
399 | break; | |||
400 | ||||
401 | if (!isa<ConstantInt>(MSI->getValue())) | |||
402 | break; | |||
403 | ||||
404 | // Check to see if this store is to a constant offset from the start ptr. | |||
405 | Optional<int64_t> Offset = isPointerOffset(StartPtr, MSI->getDest(), *DL); | |||
406 | if (!Offset) | |||
407 | break; | |||
408 | ||||
409 | if (!IB.addMemSet(*Offset, MSI)) | |||
410 | break; | |||
411 | LastInst = MSI; | |||
412 | } | |||
413 | } | |||
414 | return LastInst; | |||
415 | } | |||
416 | ||||
417 | bool AArch64StackTagging::isInterestingAlloca(const AllocaInst &AI) { | |||
418 | // FIXME: support dynamic allocas | |||
419 | bool IsInteresting = | |||
420 | AI.getAllocatedType()->isSized() && AI.isStaticAlloca() && | |||
421 | // alloca() may be called with 0 size, ignore it. | |||
422 | AI.getAllocationSizeInBits(*DL).getValue() > 0 && | |||
423 | // inalloca allocas are not treated as static, and we don't want | |||
424 | // dynamic alloca instrumentation for them as well. | |||
425 | !AI.isUsedWithInAlloca() && | |||
426 | // swifterror allocas are register promoted by ISel | |||
427 | !AI.isSwiftError() && | |||
428 | // safe allocas are not interesting | |||
429 | !(SSI && SSI->isSafe(AI)); | |||
430 | return IsInteresting; | |||
431 | } | |||
432 | ||||
433 | void AArch64StackTagging::tagAlloca(AllocaInst *AI, Instruction *InsertBefore, | |||
434 | Value *Ptr, uint64_t Size) { | |||
435 | auto SetTagZeroFunc = | |||
436 | Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag_zero); | |||
437 | auto StgpFunc = | |||
438 | Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_stgp); | |||
439 | ||||
440 | InitializerBuilder IB(Size, DL, Ptr, SetTagFunc, SetTagZeroFunc, StgpFunc); | |||
441 | bool LittleEndian = | |||
442 | Triple(AI->getModule()->getTargetTriple()).isLittleEndian(); | |||
443 | // Current implementation of initializer merging assumes little endianness. | |||
444 | if (MergeInit && !F->hasOptNone() && LittleEndian && | |||
445 | Size < ClMergeInitSizeLimit) { | |||
446 | LLVM_DEBUG(dbgs() << "collecting initializers for " << *AIdo { } while (false) | |||
447 | << ", size = " << Size << "\n")do { } while (false); | |||
448 | InsertBefore = collectInitializers(InsertBefore, Ptr, Size, IB); | |||
449 | } | |||
450 | ||||
451 | IRBuilder<> IRB(InsertBefore); | |||
452 | IB.generate(IRB); | |||
453 | } | |||
454 | ||||
455 | void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore, | |||
456 | uint64_t Size) { | |||
457 | IRBuilder<> IRB(InsertBefore); | |||
458 | IRB.CreateCall(SetTagFunc, {IRB.CreatePointerCast(AI, IRB.getInt8PtrTy()), | |||
459 | ConstantInt::get(IRB.getInt64Ty(), Size)}); | |||
460 | } | |||
461 | ||||
462 | Instruction *AArch64StackTagging::insertBaseTaggedPointer( | |||
463 | const MapVector<AllocaInst *, AllocaInfo> &Allocas, | |||
464 | const DominatorTree *DT) { | |||
465 | BasicBlock *PrologueBB = nullptr; | |||
466 | // Try sinking IRG as deep as possible to avoid hurting shrink wrap. | |||
467 | for (auto &I : Allocas) { | |||
468 | const AllocaInfo &Info = I.second; | |||
469 | AllocaInst *AI = Info.AI; | |||
470 | if (Info.Tag < 0) | |||
471 | continue; | |||
472 | if (!PrologueBB) { | |||
473 | PrologueBB = AI->getParent(); | |||
474 | continue; | |||
475 | } | |||
476 | PrologueBB = DT->findNearestCommonDominator(PrologueBB, AI->getParent()); | |||
477 | } | |||
478 | assert(PrologueBB)(static_cast<void> (0)); | |||
479 | ||||
480 | IRBuilder<> IRB(&PrologueBB->front()); | |||
| ||||
481 | Function *IRG_SP = | |||
482 | Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_irg_sp); | |||
483 | Instruction *Base = | |||
484 | IRB.CreateCall(IRG_SP, {Constant::getNullValue(IRB.getInt64Ty())}); | |||
485 | Base->setName("basetag"); | |||
486 | return Base; | |||
487 | } | |||
488 | ||||
489 | void AArch64StackTagging::alignAndPadAlloca(AllocaInfo &Info) { | |||
490 | const Align NewAlignment = | |||
491 | max(MaybeAlign(Info.AI->getAlignment()), kTagGranuleSize); | |||
492 | Info.AI->setAlignment(NewAlignment); | |||
493 | ||||
494 | uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8; | |||
495 | uint64_t AlignedSize = alignTo(Size, kTagGranuleSize); | |||
496 | if (Size == AlignedSize) | |||
497 | return; | |||
498 | ||||
499 | // Add padding to the alloca. | |||
500 | Type *AllocatedType = | |||
501 | Info.AI->isArrayAllocation() | |||
502 | ? ArrayType::get( | |||
503 | Info.AI->getAllocatedType(), | |||
504 | cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue()) | |||
505 | : Info.AI->getAllocatedType(); | |||
506 | Type *PaddingType = | |||
507 | ArrayType::get(Type::getInt8Ty(F->getContext()), AlignedSize - Size); | |||
508 | Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType); | |||
509 | auto *NewAI = new AllocaInst( | |||
510 | TypeWithPadding, Info.AI->getType()->getAddressSpace(), nullptr, "", Info.AI); | |||
511 | NewAI->takeName(Info.AI); | |||
512 | NewAI->setAlignment(Info.AI->getAlign()); | |||
513 | NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca()); | |||
514 | NewAI->setSwiftError(Info.AI->isSwiftError()); | |||
515 | NewAI->copyMetadata(*Info.AI); | |||
516 | ||||
517 | auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI); | |||
518 | Info.AI->replaceAllUsesWith(NewPtr); | |||
519 | Info.AI->eraseFromParent(); | |||
520 | Info.AI = NewAI; | |||
521 | } | |||
522 | ||||
523 | // FIXME: check for MTE extension | |||
524 | bool AArch64StackTagging::runOnFunction(Function &Fn) { | |||
525 | if (!Fn.hasFnAttribute(Attribute::SanitizeMemTag)) | |||
| ||||
526 | return false; | |||
527 | ||||
528 | if (UseStackSafety) | |||
529 | SSI = &getAnalysis<StackSafetyGlobalInfoWrapperPass>().getResult(); | |||
530 | F = &Fn; | |||
531 | DL = &Fn.getParent()->getDataLayout(); | |||
532 | if (MergeInit) | |||
533 | AA = &getAnalysis<AAResultsWrapperPass>().getAAResults(); | |||
534 | ||||
535 | MapVector<AllocaInst *, AllocaInfo> Allocas; // need stable iteration order | |||
536 | SmallVector<Instruction *, 8> RetVec; | |||
537 | SmallVector<Instruction *, 4> UnrecognizedLifetimes; | |||
538 | ||||
539 | for (auto &BB : *F) { | |||
540 | for (BasicBlock::iterator IT = BB.begin(); IT != BB.end(); ++IT) { | |||
541 | Instruction *I = &*IT; | |||
542 | if (auto *AI = dyn_cast<AllocaInst>(I)) { | |||
543 | Allocas[AI].AI = AI; | |||
544 | Allocas[AI].OldAI = AI; | |||
545 | continue; | |||
546 | } | |||
547 | ||||
548 | if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(I)) { | |||
549 | for (Value *V : DVI->location_ops()) | |||
550 | if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) | |||
551 | if (Allocas[AI].DbgVariableIntrinsics.empty() || | |||
552 | Allocas[AI].DbgVariableIntrinsics.back() != DVI) | |||
553 | Allocas[AI].DbgVariableIntrinsics.push_back(DVI); | |||
554 | continue; | |||
555 | } | |||
556 | ||||
557 | auto *II = dyn_cast<IntrinsicInst>(I); | |||
558 | if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start || | |||
559 | II->getIntrinsicID() == Intrinsic::lifetime_end)) { | |||
560 | AllocaInst *AI = findAllocaForValue(II->getArgOperand(1)); | |||
561 | if (!AI) { | |||
562 | UnrecognizedLifetimes.push_back(I); | |||
563 | continue; | |||
564 | } | |||
565 | if (II->getIntrinsicID() == Intrinsic::lifetime_start) | |||
566 | Allocas[AI].LifetimeStart.push_back(II); | |||
567 | else | |||
568 | Allocas[AI].LifetimeEnd.push_back(II); | |||
569 | } | |||
570 | ||||
571 | if (isa<ReturnInst>(I) || isa<ResumeInst>(I) || isa<CleanupReturnInst>(I)) | |||
572 | RetVec.push_back(I); | |||
573 | } | |||
574 | } | |||
575 | ||||
576 | if (Allocas.empty()) | |||
577 | return false; | |||
578 | ||||
579 | int NextTag = 0; | |||
580 | int NumInterestingAllocas = 0; | |||
581 | for (auto &I : Allocas) { | |||
582 | AllocaInfo &Info = I.second; | |||
583 | assert(Info.AI)(static_cast<void> (0)); | |||
584 | ||||
585 | if (!isInterestingAlloca(*Info.AI)) { | |||
586 | Info.Tag = -1; | |||
587 | continue; | |||
588 | } | |||
589 | ||||
590 | alignAndPadAlloca(Info); | |||
591 | NumInterestingAllocas++; | |||
592 | Info.Tag = NextTag; | |||
593 | NextTag = (NextTag + 1) % 16; | |||
594 | } | |||
595 | ||||
596 | if (NumInterestingAllocas
| |||
597 | return true; | |||
598 | ||||
599 | std::unique_ptr<DominatorTree> DeleteDT; | |||
600 | DominatorTree *DT = nullptr; | |||
601 | if (auto *P
| |||
602 | DT = &P->getDomTree(); | |||
603 | ||||
604 | if (DT == nullptr && (NumInterestingAllocas
| |||
605 | !F->hasFnAttribute(Attribute::OptimizeNone))) { | |||
606 | DeleteDT = std::make_unique<DominatorTree>(*F); | |||
607 | DT = DeleteDT.get(); | |||
608 | } | |||
609 | ||||
610 | std::unique_ptr<PostDominatorTree> DeletePDT; | |||
611 | PostDominatorTree *PDT = nullptr; | |||
612 | if (auto *P
| |||
613 | PDT = &P->getPostDomTree(); | |||
614 | ||||
615 | if (PDT == nullptr && !F->hasFnAttribute(Attribute::OptimizeNone)) { | |||
616 | DeletePDT = std::make_unique<PostDominatorTree>(*F); | |||
617 | PDT = DeletePDT.get(); | |||
618 | } | |||
619 | ||||
620 | SetTagFunc = | |||
621 | Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag); | |||
622 | ||||
623 | Instruction *Base = insertBaseTaggedPointer(Allocas, DT); | |||
624 | ||||
625 | for (auto &I : Allocas) { | |||
626 | const AllocaInfo &Info = I.second; | |||
627 | AllocaInst *AI = Info.AI; | |||
628 | if (Info.Tag < 0) | |||
629 | continue; | |||
630 | ||||
631 | // Replace alloca with tagp(alloca). | |||
632 | IRBuilder<> IRB(Info.AI->getNextNode()); | |||
633 | Function *TagP = Intrinsic::getDeclaration( | |||
634 | F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()}); | |||
635 | Instruction *TagPCall = | |||
636 | IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base, | |||
637 | ConstantInt::get(IRB.getInt64Ty(), Info.Tag)}); | |||
638 | if (Info.AI->hasName()) | |||
639 | TagPCall->setName(Info.AI->getName() + ".tag"); | |||
640 | Info.AI->replaceAllUsesWith(TagPCall); | |||
641 | TagPCall->setOperand(0, Info.AI); | |||
642 | ||||
643 | if (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 && | |||
644 | Info.LifetimeEnd.size() == 1) { | |||
645 | IntrinsicInst *Start = Info.LifetimeStart[0]; | |||
646 | IntrinsicInst *End = Info.LifetimeEnd[0]; | |||
647 | uint64_t Size = | |||
648 | cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue(); | |||
649 | Size = alignTo(Size, kTagGranuleSize); | |||
650 | tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size); | |||
651 | ||||
652 | auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); }; | |||
653 | if (!DT || !PDT || | |||
654 | !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec, | |||
655 | TagEnd)) | |||
656 | End->eraseFromParent(); | |||
657 | } else { | |||
658 | uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8; | |||
659 | Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy()); | |||
660 | tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size); | |||
661 | for (auto &RI : RetVec) { | |||
662 | untagAlloca(AI, RI, Size); | |||
663 | } | |||
664 | // We may have inserted tag/untag outside of any lifetime interval. | |||
665 | // Remove all lifetime intrinsics for this alloca. | |||
666 | for (auto &II : Info.LifetimeStart) | |||
667 | II->eraseFromParent(); | |||
668 | for (auto &II : Info.LifetimeEnd) | |||
669 | II->eraseFromParent(); | |||
670 | } | |||
671 | ||||
672 | // Fixup debug intrinsics to point to the new alloca. | |||
673 | for (auto DVI : Info.DbgVariableIntrinsics) | |||
674 | DVI->replaceVariableLocationOp(Info.OldAI, Info.AI); | |||
675 | } | |||
676 | ||||
677 | // If we have instrumented at least one alloca, all unrecognized lifetime | |||
678 | // instrinsics have to go. | |||
679 | for (auto &I : UnrecognizedLifetimes) | |||
680 | I->eraseFromParent(); | |||
681 | ||||
682 | return true; | |||
683 | } |