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/StringExtras.h"
14#include "llvm/ADT/StringRef.h"
15#include "llvm/Analysis/OptimizationRemarkEmitter.h"
16#include "llvm/Analysis/ProfileSummaryInfo.h"
17#include "llvm/IR/Constants.h"
18#include "llvm/IR/DiagnosticInfo.h"
19#include "llvm/IR/InstIterator.h"
20#include "llvm/IR/Instructions.h"
21#include "llvm/IR/IntrinsicInst.h"
22#include "llvm/IR/Intrinsics.h"
23#include "llvm/IR/Metadata.h"
24#include "llvm/IR/Module.h"
25#include "llvm/Support/Debug.h"
26#include "llvm/Support/RandomNumberGenerator.h"
27#include <memory>
28#include <optional>
29#include <random>
30
31using namespace llvm;
32
33#define DEBUG_TYPE "lower-allow-check"
34
35static cl::opt<int>
36 HotPercentileCutoff("lower-allow-check-percentile-cutoff-hot",
37 cl::desc("Hot percentile cutoff."));
38
39static cl::opt<float>
40 RandomRate("lower-allow-check-random-rate",
41 cl::desc("Probability value in the range [0.0, 1.0] of "
42 "unconditional pseudo-random checks."));
43
44STATISTIC(NumChecksTotal, "Number of checks");
45STATISTIC(NumChecksRemoved, "Number of removed checks");
46
47struct RemarkInfo {
48 ore::NV Kind;
49 ore::NV F;
50 ore::NV BB;
51 explicit RemarkInfo(IntrinsicInst *II)
52 : Kind("Kind", II->getArgOperand(i: 0)),
53 F("Function", II->getParent()->getParent()),
54 BB("Block", II->getParent()->getName()) {}
55};
56
57static void emitRemark(IntrinsicInst *II, OptimizationRemarkEmitter &ORE,
58 bool Removed) {
59 if (Removed) {
60 ORE.emit(RemarkBuilder: [&]() {
61 RemarkInfo Info(II);
62 return OptimizationRemark(DEBUG_TYPE, "Removed", II)
63 << "Removed check: Kind=" << Info.Kind << " F=" << Info.F
64 << " BB=" << Info.BB;
65 });
66 } else {
67 ORE.emit(RemarkBuilder: [&]() {
68 RemarkInfo Info(II);
69 return OptimizationRemarkMissed(DEBUG_TYPE, "Allowed", II)
70 << "Allowed check: Kind=" << Info.Kind << " F=" << Info.F
71 << " BB=" << Info.BB;
72 });
73 }
74}
75
76static bool lowerAllowChecks(Function &F, FunctionAnalysisManager &AM,
77 const LowerAllowCheckPass::Options &Opts) {
78 // Lazy analysis getters.
79 auto GetBFI = [&AM, &F, BFI = (BlockFrequencyInfo *)nullptr]() mutable
80 -> const BlockFrequencyInfo & {
81 if (!BFI)
82 BFI = &AM.getResult<BlockFrequencyAnalysis>(IR&: F);
83 return *BFI;
84 };
85 auto GetPSI = [&AM, &F, PSI = std::optional<ProfileSummaryInfo *>()]() mutable
86 -> const ProfileSummaryInfo * {
87 if (!PSI.has_value()) {
88 auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(IR&: F);
89 PSI = MAMProxy.getCachedResult<ProfileSummaryAnalysis>(IR&: *F.getParent());
90 }
91 return *PSI;
92 };
93 auto GetORE = [&AM, &F, ORE = (OptimizationRemarkEmitter *)nullptr]() mutable
94 -> OptimizationRemarkEmitter & {
95 if (!ORE)
96 ORE = &AM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: F);
97 return *ORE;
98 };
99
100 // List of intrinsics and the constant value they should be lowered to.
101 SmallVector<std::pair<IntrinsicInst *, bool>, 16> ReplaceWithValue;
102 std::unique_ptr<RandomNumberGenerator> Rng;
103
104 auto GetRng = [&]() -> RandomNumberGenerator & {
105 if (!Rng)
106 Rng = F.getParent()->createRNG(Name: F.getName());
107 return *Rng;
108 };
109
110 auto GetCutoff = [&](const IntrinsicInst *II) -> unsigned {
111 if (HotPercentileCutoff.getNumOccurrences())
112 return HotPercentileCutoff;
113 else if (II->getIntrinsicID() == Intrinsic::allow_ubsan_check) {
114 auto *Kind = cast<ConstantInt>(Val: II->getArgOperand(i: 0));
115 if (Kind->getZExtValue() < Opts.cutoffs.size())
116 return Opts.cutoffs[Kind->getZExtValue()];
117 } else if (II->getIntrinsicID() == Intrinsic::allow_runtime_check) {
118 return Opts.runtime_check;
119 }
120
121 return 0;
122 };
123
124 auto ShouldRemoveHot = [&](const BasicBlock &BB, unsigned int cutoff) {
125 if (cutoff == 1000000)
126 return true;
127 const ProfileSummaryInfo *PSI = GetPSI();
128 return PSI && PSI->isHotCountNthPercentile(
129 PercentileCutoff: cutoff, C: GetBFI().getBlockProfileCount(BB: &BB).value_or(u: 0));
130 };
131
132 auto ShouldRemoveRandom = [&]() {
133 return RandomRate.getNumOccurrences() &&
134 !std::bernoulli_distribution(RandomRate)(GetRng());
135 };
136
137 auto ShouldRemove = [&](const IntrinsicInst *II) {
138 unsigned int cutoff = GetCutoff(II);
139 return ShouldRemoveRandom() || ShouldRemoveHot(*(II->getParent()), cutoff);
140 };
141
142 for (Instruction &I : instructions(F)) {
143 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: &I);
144 if (!II)
145 continue;
146 auto ID = II->getIntrinsicID();
147 switch (ID) {
148 case Intrinsic::allow_ubsan_check:
149 case Intrinsic::allow_runtime_check: {
150 bool ToRemove = ShouldRemove(II);
151
152 ReplaceWithValue.push_back(Elt: {
153 II,
154 !ToRemove,
155 });
156 emitRemark(II, ORE&: GetORE(), Removed: ToRemove);
157 break;
158 }
159 case Intrinsic::allow_sanitize_address:
160 ReplaceWithValue.push_back(
161 Elt: {II, F.hasFnAttribute(Kind: Attribute::SanitizeAddress)});
162 break;
163 case Intrinsic::allow_sanitize_thread:
164 ReplaceWithValue.push_back(
165 Elt: {II, F.hasFnAttribute(Kind: Attribute::SanitizeThread)});
166 break;
167 case Intrinsic::allow_sanitize_memory:
168 ReplaceWithValue.push_back(
169 Elt: {II, F.hasFnAttribute(Kind: Attribute::SanitizeMemory)});
170 break;
171 case Intrinsic::allow_sanitize_hwaddress:
172 ReplaceWithValue.push_back(
173 Elt: {II, F.hasFnAttribute(Kind: Attribute::SanitizeHWAddress)});
174 break;
175 default:
176 break;
177 }
178 }
179
180 for (auto [I, V] : ReplaceWithValue) {
181 ++NumChecksTotal;
182 if (!V) // If the final value is false, the check is considered removed.
183 ++NumChecksRemoved;
184 I->replaceAllUsesWith(V: ConstantInt::getBool(Ty: I->getType(), V));
185 I->eraseFromParent();
186 }
187
188 return !ReplaceWithValue.empty();
189}
190
191PreservedAnalyses LowerAllowCheckPass::run(Function &F,
192 FunctionAnalysisManager &AM) {
193 if (F.isDeclaration())
194 return PreservedAnalyses::all();
195
196 return lowerAllowChecks(F, AM, Opts)
197 // We do not change the CFG, we only replace the intrinsics with
198 // true or false.
199 ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()
200 : PreservedAnalyses::all();
201}
202
203bool LowerAllowCheckPass::IsRequested() {
204 return RandomRate.getNumOccurrences() ||
205 HotPercentileCutoff.getNumOccurrences();
206}
207
208void LowerAllowCheckPass::printPipeline(
209 raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
210 static_cast<PassInfoMixin<LowerAllowCheckPass> *>(this)->printPipeline(
211 OS, MapClassName2PassName);
212 OS << "<";
213
214 // Format is <cutoffs[0,1,2]=70000;cutoffs[5,6,8]=90000>
215 // but it's equally valid to specify
216 // cutoffs[0]=70000;cutoffs[1]=70000;cutoffs[2]=70000;cutoffs[5]=90000;...
217 // and that's what we do here. It is verbose but valid and easy to verify
218 // correctness.
219 // TODO: print shorter output by combining adjacent runs, etc.
220 int i = 0;
221 ListSeparator LS(";");
222 for (unsigned int cutoff : Opts.cutoffs) {
223 if (cutoff > 0)
224 OS << LS << "cutoffs[" << i << "]=" << cutoff;
225 i++;
226 }
227 if (Opts.runtime_check)
228 OS << LS << "runtime_check=" << Opts.runtime_check;
229
230 OS << '>';
231}
232