xref: /llvm-project/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp (revision 67efbd0bf1b2df8a479e09eb2be7db4c3c892f2c)
1 //===- LowerAllowCheckPass.cpp ----------------------------------*- C++ -*-===//
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 #include "llvm/Transforms/Instrumentation/LowerAllowCheckPass.h"
10 
11 #include "llvm/ADT/SmallVector.h"
12 #include "llvm/ADT/Statistic.h"
13 #include "llvm/ADT/StringRef.h"
14 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
15 #include "llvm/Analysis/ProfileSummaryInfo.h"
16 #include "llvm/IR/Constants.h"
17 #include "llvm/IR/DiagnosticInfo.h"
18 #include "llvm/IR/Instructions.h"
19 #include "llvm/IR/IntrinsicInst.h"
20 #include "llvm/IR/Intrinsics.h"
21 #include "llvm/IR/Metadata.h"
22 #include "llvm/IR/Module.h"
23 #include "llvm/Support/RandomNumberGenerator.h"
24 #include <memory>
25 #include <random>
26 
27 using namespace llvm;
28 
29 #define DEBUG_TYPE "lower-allow-check"
30 
31 static cl::opt<int>
32     HotPercentileCutoff("lower-allow-check-percentile-cutoff-hot",
33                         cl::desc("Hot percentile cutoff."));
34 
35 static cl::opt<float>
36     RandomRate("lower-allow-check-random-rate",
37                cl::desc("Probability value in the range [0.0, 1.0] of "
38                         "unconditional pseudo-random checks."));
39 
40 STATISTIC(NumChecksTotal, "Number of checks");
41 STATISTIC(NumChecksRemoved, "Number of removed checks");
42 
43 struct RemarkInfo {
44   ore::NV Kind;
45   ore::NV F;
46   ore::NV BB;
47   explicit RemarkInfo(IntrinsicInst *II)
48       : Kind("Kind", II->getArgOperand(0)),
49         F("Function", II->getParent()->getParent()),
50         BB("Block", II->getParent()->getName()) {}
51 };
52 
53 static void emitRemark(IntrinsicInst *II, OptimizationRemarkEmitter &ORE,
54                        bool Removed) {
55   if (Removed) {
56     ORE.emit([&]() {
57       RemarkInfo Info(II);
58       return OptimizationRemark(DEBUG_TYPE, "Removed", II)
59              << "Removed check: Kind=" << Info.Kind << " F=" << Info.F
60              << " BB=" << Info.BB;
61     });
62   } else {
63     ORE.emit([&]() {
64       RemarkInfo Info(II);
65       return OptimizationRemarkMissed(DEBUG_TYPE, "Allowed", II)
66              << "Allowed check: Kind=" << Info.Kind << " F=" << Info.F
67              << " BB=" << Info.BB;
68     });
69   }
70 }
71 
72 static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
73                              const ProfileSummaryInfo *PSI,
74                              OptimizationRemarkEmitter &ORE) {
75   SmallVector<std::pair<IntrinsicInst *, bool>, 16> ReplaceWithValue;
76   std::unique_ptr<RandomNumberGenerator> Rng;
77 
78   auto GetRng = [&]() -> RandomNumberGenerator & {
79     if (!Rng)
80       Rng = F.getParent()->createRNG(F.getName());
81     return *Rng;
82   };
83 
84   auto ShouldRemoveHot = [&](const BasicBlock &BB) {
85     return HotPercentileCutoff.getNumOccurrences() && PSI &&
86            PSI->isHotCountNthPercentile(
87                HotPercentileCutoff, BFI.getBlockProfileCount(&BB).value_or(0));
88   };
89 
90   auto ShouldRemoveRandom = [&]() {
91     return RandomRate.getNumOccurrences() &&
92            !std::bernoulli_distribution(RandomRate)(GetRng());
93   };
94 
95   auto ShouldRemove = [&](const BasicBlock &BB) {
96     return ShouldRemoveRandom() || ShouldRemoveHot(BB);
97   };
98 
99   for (BasicBlock &BB : F) {
100     for (Instruction &I : BB) {
101       IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
102       if (!II)
103         continue;
104       auto ID = II->getIntrinsicID();
105       switch (ID) {
106       case Intrinsic::allow_ubsan_check:
107       case Intrinsic::allow_runtime_check: {
108         ++NumChecksTotal;
109 
110         bool ToRemove = ShouldRemove(BB);
111         ReplaceWithValue.push_back({
112             II,
113             ToRemove,
114         });
115         if (ToRemove)
116           ++NumChecksRemoved;
117         emitRemark(II, ORE, ToRemove);
118         break;
119       }
120       default:
121         break;
122       }
123     }
124   }
125 
126   for (auto [I, V] : ReplaceWithValue) {
127     I->replaceAllUsesWith(ConstantInt::getBool(I->getType(), !V));
128     I->eraseFromParent();
129   }
130 
131   return !ReplaceWithValue.empty();
132 }
133 
134 PreservedAnalyses LowerAllowCheckPass::run(Function &F,
135                                            FunctionAnalysisManager &AM) {
136   if (F.isDeclaration())
137     return PreservedAnalyses::all();
138   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
139   ProfileSummaryInfo *PSI =
140       MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
141   BlockFrequencyInfo &BFI = AM.getResult<BlockFrequencyAnalysis>(F);
142   OptimizationRemarkEmitter &ORE =
143       AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
144 
145   return removeUbsanTraps(F, BFI, PSI, ORE) ? PreservedAnalyses::none()
146                                             : PreservedAnalyses::all();
147 }
148 
149 bool LowerAllowCheckPass::IsRequested() {
150   return RandomRate.getNumOccurrences() ||
151          HotPercentileCutoff.getNumOccurrences();
152 }
153