Coverage Report

Created: 2019-07-24 05:18

/Users/buildslave/jenkins/workspace/clang-stage2-coverage-R/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
Line
Count
Source (jump to first uncovered line)
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
8.78k
  AArch64StackTagging() : FunctionPass(ID) {
74
8.78k
    initializeAArch64StackTaggingPass(*PassRegistry::getPassRegistry());
75
8.78k
  }
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
258k
  StringRef getPassName() const override { return "AArch64 Stack Tagging"; }
90
91
private:
92
  Function *F;
93
  Function *SetTagFunc;
94
  const DataLayout *DL;
95
96
8.75k
  void getAnalysisUsage(AnalysisUsage &AU) const override {
97
8.75k
    AU.setPreservesCFG();
98
8.75k
  }
99
};
100
101
} // end anonymous namespace
102
103
char AArch64StackTagging::ID = 0;
104
105
101k
INITIALIZE_PASS_BEGIN(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",
106
101k
                      false, false)
107
101k
INITIALIZE_PASS_END(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",
108
                    false, false)
109
110
8.78k
FunctionPass *llvm::createAArch64StackTaggingPass() {
111
8.78k
  return new AArch64StackTagging();
112
8.78k
}
113
114
12
bool AArch64StackTagging::isInterestingAlloca(const AllocaInst &AI) {
115
12
  // FIXME: support dynamic allocas
116
12
  bool IsInteresting =
117
12
      AI.getAllocatedType()->isSized() && AI.isStaticAlloca() &&
118
12
      // alloca() may be called with 0 size, ignore it.
119
12
      
AI.getAllocationSizeInBits(*DL).getValue() > 010
&&
120
12
      // inalloca allocas are not treated as static, and we don't want
121
12
      // dynamic alloca instrumentation for them as well.
122
12
      
!AI.isUsedWithInAlloca()10
&&
123
12
      // swifterror allocas are register promoted by ISel
124
12
      
!AI.isSwiftError()10
;
125
12
  return IsInteresting;
126
12
}
127
128
void AArch64StackTagging::tagAlloca(AllocaInst *AI, Instruction *InsertBefore,
129
10
                                    Value *Ptr, uint64_t Size) {
130
10
  IRBuilder<> IRB(InsertBefore);
131
10
  IRB.CreateCall(SetTagFunc, {Ptr, ConstantInt::get(IRB.getInt64Ty(), Size)});
132
10
}
133
134
void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
135
10
                                      uint64_t Size) {
136
10
  IRBuilder<> IRB(InsertBefore);
137
10
  IRB.CreateCall(SetTagFunc, {IRB.CreatePointerCast(AI, IRB.getInt8PtrTy()),
138
10
                              ConstantInt::get(IRB.getInt64Ty(), Size)});
139
10
}
140
141
Instruction *AArch64StackTagging::insertBaseTaggedPointer(
142
    const MapVector<AllocaInst *, AllocaInfo> &Allocas,
143
6
    const DominatorTree *DT) {
144
6
  BasicBlock *PrologueBB = nullptr;
145
6
  // Try sinking IRG as deep as possible to avoid hurting shrink wrap.
146
10
  for (auto &I : Allocas) {
147
10
    const AllocaInfo &Info = I.second;
148
10
    AllocaInst *AI = Info.AI;
149
10
    if (Info.Tag < 0)
150
0
      continue;
151
10
    if (!PrologueBB) {
152
6
      PrologueBB = AI->getParent();
153
6
      continue;
154
6
    }
155
4
    PrologueBB = DT->findNearestCommonDominator(PrologueBB, AI->getParent());
156
4
  }
157
6
  assert(PrologueBB);
158
6
159
6
  IRBuilder<> IRB(&PrologueBB->front());
160
6
  Function *IRG_SP =
161
6
      Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_irg_sp);
162
6
  Instruction *Base =
163
6
      IRB.CreateCall(IRG_SP, {Constant::getNullValue(IRB.getInt64Ty())});
164
6
  Base->setName("basetag");
165
6
  return Base;
166
6
}
167
168
10
void AArch64StackTagging::alignAndPadAlloca(AllocaInfo &Info) {
169
10
  unsigned NewAlignment = std::max(Info.AI->getAlignment(), kTagGranuleSize);
170
10
  Info.AI->setAlignment(NewAlignment);
171
10
172
10
  uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
173
10
  uint64_t AlignedSize = alignTo(Size, kTagGranuleSize);
174
10
  if (Size == AlignedSize)
175
0
    return;
176
10
177
10
  // Add padding to the alloca.
178
10
  Type *AllocatedType =
179
10
      Info.AI->isArrayAllocation()
180
10
          ? ArrayType::get(
181
1
                Info.AI->getAllocatedType(),
182
1
                dyn_cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue())
183
10
          : 
Info.AI->getAllocatedType()9
;
184
10
  Type *PaddingType =
185
10
      ArrayType::get(Type::getInt8Ty(F->getContext()), AlignedSize - Size);
186
10
  Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType);
187
10
  auto *NewAI = new AllocaInst(
188
10
      TypeWithPadding, Info.AI->getType()->getAddressSpace(), nullptr, "", Info.AI);
189
10
  NewAI->takeName(Info.AI);
190
10
  NewAI->setAlignment(Info.AI->getAlignment());
191
10
  NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca());
192
10
  NewAI->setSwiftError(Info.AI->isSwiftError());
193
10
  NewAI->copyMetadata(*Info.AI);
194
10
195
10
  auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI);
196
10
  Info.AI->replaceAllUsesWith(NewPtr);
197
10
  Info.AI->eraseFromParent();
198
10
  Info.AI = NewAI;
199
10
}
200
201
// FIXME: check for MTE extension
202
258k
bool AArch64StackTagging::runOnFunction(Function &Fn) {
203
258k
  if (!Fn.hasFnAttribute(Attribute::SanitizeMemTag))
204
258k
    return false;
205
7
206
7
  F = &Fn;
207
7
  DL = &Fn.getParent()->getDataLayout();
208
7
209
7
  MapVector<AllocaInst *, AllocaInfo> Allocas; // need stable iteration order
210
7
  SmallVector<Instruction *, 8> RetVec;
211
7
  DenseMap<Value *, AllocaInst *> AllocaForValue;
212
7
  SmallVector<Instruction *, 4> UnrecognizedLifetimes;
213
7
214
13
  for (auto &BB : *F) {
215
75
    for (BasicBlock::iterator IT = BB.begin(); IT != BB.end(); 
++IT62
) {
216
62
      Instruction *I = &*IT;
217
62
      if (auto *AI = dyn_cast<AllocaInst>(I)) {
218
12
        Allocas[AI].AI = AI;
219
12
        continue;
220
12
      }
221
50
222
50
      if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(I)) {
223
1
        if (auto *AI =
224
1
                dyn_cast_or_null<AllocaInst>(DVI->getVariableLocation())) {
225
1
          Allocas[AI].DbgVariableIntrinsics.push_back(DVI);
226
1
        }
227
1
        continue;
228
1
      }
229
49
230
49
      auto *II = dyn_cast<IntrinsicInst>(I);
231
49
      if (II && 
(12
II->getIntrinsicID() == Intrinsic::lifetime_start12
||
232
12
                 
II->getIntrinsicID() == Intrinsic::lifetime_end6
)) {
233
12
        AllocaInst *AI =
234
12
            llvm::findAllocaForValue(II->getArgOperand(1), AllocaForValue);
235
12
        if (!AI) {
236
2
          UnrecognizedLifetimes.push_back(I);
237
2
          continue;
238
2
        }
239
10
        if (II->getIntrinsicID() == Intrinsic::lifetime_start)
240
5
          Allocas[AI].LifetimeStart.push_back(II);
241
5
        else
242
5
          Allocas[AI].LifetimeEnd.push_back(II);
243
10
      }
244
49
245
49
      
if (47
isa<ReturnInst>(I)47
||
isa<ResumeInst>(I)40
||
isa<CleanupReturnInst>(I)40
)
246
7
        RetVec.push_back(I);
247
47
    }
248
13
  }
249
7
250
7
  if (Allocas.empty())
251
0
    return false;
252
7
253
7
  int NextTag = 0;
254
7
  int NumInterestingAllocas = 0;
255
12
  for (auto &I : Allocas) {
256
12
    AllocaInfo &Info = I.second;
257
12
    assert(Info.AI);
258
12
259
12
    if (!isInterestingAlloca(*Info.AI)) {
260
2
      Info.Tag = -1;
261
2
      continue;
262
2
    }
263
10
264
10
    alignAndPadAlloca(Info);
265
10
    NumInterestingAllocas++;
266
10
    Info.Tag = NextTag;
267
10
    NextTag = (NextTag + 1) % 16;
268
10
  }
269
7
270
7
  if (NumInterestingAllocas == 0)
271
1
    return true;
272
6
273
6
  SetTagFunc =
274
6
      Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
275
6
276
6
  // Compute DT only if the function has the attribute, there are more than 1
277
6
  // interesting allocas, and it is not available for free.
278
6
  Instruction *Base;
279
6
  if (NumInterestingAllocas > 1) {
280
2
    auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
281
2
    if (DTWP) {
282
0
      Base = insertBaseTaggedPointer(Allocas, &DTWP->getDomTree());
283
2
    } else {
284
2
      DominatorTree DT(*F);
285
2
      Base = insertBaseTaggedPointer(Allocas, &DT);
286
2
    }
287
4
  } else {
288
4
    Base = insertBaseTaggedPointer(Allocas, nullptr);
289
4
  }
290
6
291
10
  for (auto &I : Allocas) {
292
10
    const AllocaInfo &Info = I.second;
293
10
    AllocaInst *AI = Info.AI;
294
10
    if (Info.Tag < 0)
295
0
      continue;
296
10
297
10
    // Replace alloca with tagp(alloca).
298
10
    IRBuilder<> IRB(Info.AI->getNextNode());
299
10
    Function *TagP = Intrinsic::getDeclaration(
300
10
        F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()});
301
10
    Instruction *TagPCall =
302
10
        IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base,
303
10
                              ConstantInt::get(IRB.getInt64Ty(), Info.Tag)});
304
10
    if (Info.AI->hasName())
305
10
      TagPCall->setName(Info.AI->getName() + ".tag");
306
10
    Info.AI->replaceAllUsesWith(TagPCall);
307
10
    TagPCall->setOperand(0, Info.AI);
308
10
309
10
    if (UnrecognizedLifetimes.empty() && 
Info.LifetimeStart.size() == 17
&&
310
10
        
Info.LifetimeEnd.size() == 11
) {
311
1
      IntrinsicInst *Start = Info.LifetimeStart[0];
312
1
      uint64_t Size =
313
1
          dyn_cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
314
1
      Size = alignTo(Size, kTagGranuleSize);
315
1
      tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size);
316
1
      untagAlloca(AI, Info.LifetimeEnd[0], Size);
317
9
    } else {
318
9
      uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
319
9
      Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
320
9
      tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);
321
9
      for (auto &RI : RetVec) {
322
9
        untagAlloca(AI, RI, Size);
323
9
      }
324
9
      // We may have inserted tag/untag outside of any lifetime interval.
325
9
      // Remove all lifetime intrinsics for this alloca.
326
9
      for (auto &II : Info.LifetimeStart)
327
4
        II->eraseFromParent();
328
9
      for (auto &II : Info.LifetimeEnd)
329
4
        II->eraseFromParent();
330
9
    }
331
10
332
10
    // Fixup debug intrinsics to point to the new alloca.
333
10
    for (auto DVI : Info.DbgVariableIntrinsics)
334
1
      DVI->setArgOperand(
335
1
          0,
336
1
          MetadataAsValue::get(F->getContext(), LocalAsMetadata::get(Info.AI)));
337
10
  }
338
6
339
6
  // If we have instrumented at least one alloca, all unrecognized lifetime
340
6
  // instrinsics have to go.
341
6
  for (auto &I : UnrecognizedLifetimes)
342
2
    I->eraseFromParent();
343
6
344
6
  return true;
345
6
}