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