1//===- AssumptionCache.cpp - Cache finding @llvm.assume calls -------------===//
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// This file contains a pass that keeps track of @llvm.assume intrinsics in
10// the functions of a module.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/Analysis/AssumptionCache.h"
15#include "llvm/ADT/DenseSet.h"
16#include "llvm/ADT/STLExtras.h"
17#include "llvm/ADT/SmallPtrSet.h"
18#include "llvm/ADT/SmallVector.h"
19#include "llvm/ADT/StringExtras.h"
20#include "llvm/Analysis/AssumeBundleQueries.h"
21#include "llvm/Analysis/TargetTransformInfo.h"
22#include "llvm/Analysis/ValueTracking.h"
23#include "llvm/IR/BasicBlock.h"
24#include "llvm/IR/Function.h"
25#include "llvm/IR/InstrTypes.h"
26#include "llvm/IR/Instruction.h"
27#include "llvm/IR/Instructions.h"
28#include "llvm/IR/PassManager.h"
29#include "llvm/IR/PatternMatch.h"
30#include "llvm/InitializePasses.h"
31#include "llvm/Pass.h"
32#include "llvm/Support/Casting.h"
33#include "llvm/Support/CommandLine.h"
34#include "llvm/Support/ErrorHandling.h"
35#include "llvm/Support/raw_ostream.h"
36#include <cassert>
37
38using namespace llvm;
39using namespace llvm::PatternMatch;
40
41static cl::opt<bool>
42 VerifyAssumptionCache("verify-assumption-cache", cl::Hidden,
43 cl::desc("Enable verification of assumption cache"),
44 cl::init(Val: false));
45
46SmallVector<AssumptionCache::ResultElem, 1> &
47AssumptionCache::getOrInsertAffectedValues(Value *V) {
48 // Try using find_as first to avoid creating extra value handles just for the
49 // purpose of doing the lookup.
50 auto AVI = AffectedValues.find_as(Val: V);
51 if (AVI != AffectedValues.end())
52 return AVI->second;
53
54 return AffectedValues[AffectedValueCallbackVH(V, this)];
55}
56
57void AssumptionCache::findValuesAffectedByOperandBundle(
58 OperandBundleUse Bundle, function_ref<void(Value *)> InsertAffected) {
59 auto AddAffectedVal = [&](Value *V) {
60 if (isa<Argument, GlobalValue, Instruction>(Val: V))
61 InsertAffected(V);
62 };
63
64 if (Bundle.getTagName() == "separate_storage") {
65 assert(Bundle.Inputs.size() == 2 && "separate_storage must have two args");
66 AddAffectedVal(getUnderlyingObject(V: Bundle.Inputs[0]));
67 AddAffectedVal(getUnderlyingObject(V: Bundle.Inputs[1]));
68 } else if (Bundle.Inputs.size() > ABA_WasOn &&
69 Bundle.getTagName() != IgnoreBundleTag)
70 AddAffectedVal(Bundle.Inputs[ABA_WasOn]);
71}
72
73static void
74findAffectedValues(CallBase *CI, TargetTransformInfo *TTI,
75 SmallVectorImpl<AssumptionCache::ResultElem> &Affected) {
76 // Note: This code must be kept in-sync with the code in
77 // computeKnownBitsFromAssume in ValueTracking.
78
79 auto InsertAffected = [&Affected](Value *V) {
80 Affected.push_back(Elt: {.Assume: V, .Index: AssumptionCache::ExprResultIdx});
81 };
82
83 auto AddAffectedVal = [&Affected](Value *V, unsigned Idx) {
84 if (isa<Argument>(Val: V) || isa<GlobalValue>(Val: V) || isa<Instruction>(Val: V)) {
85 Affected.push_back(Elt: {.Assume: V, .Index: Idx});
86 }
87 };
88
89 for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++)
90 AssumptionCache::findValuesAffectedByOperandBundle(
91 Bundle: CI->getOperandBundleAt(Index: Idx),
92 InsertAffected: [&](Value *V) { Affected.push_back(Elt: {.Assume: V, .Index: Idx}); });
93
94 Value *Cond = CI->getArgOperand(i: 0);
95 findValuesAffectedByCondition(Cond, /*IsAssume=*/true, InsertAffected);
96
97 if (TTI) {
98 const Value *Ptr;
99 unsigned AS;
100 std::tie(args&: Ptr, args&: AS) = TTI->getPredicatedAddrSpace(V: Cond);
101 if (Ptr)
102 AddAffectedVal(const_cast<Value *>(Ptr->stripInBoundsOffsets()),
103 AssumptionCache::ExprResultIdx);
104 }
105}
106
107void AssumptionCache::updateAffectedValues(AssumeInst *CI) {
108 SmallVector<AssumptionCache::ResultElem, 16> Affected;
109 findAffectedValues(CI, TTI, Affected);
110
111 for (auto &AV : Affected) {
112 auto &AVV = getOrInsertAffectedValues(V: AV.Assume);
113 if (llvm::none_of(Range&: AVV, P: [&](ResultElem &Elem) {
114 return Elem.Assume == CI && Elem.Index == AV.Index;
115 }))
116 AVV.push_back(Elt: {.Assume: CI, .Index: AV.Index});
117 }
118}
119
120void AssumptionCache::removeAffectedValues(AssumeInst *CI) {
121 SmallVector<AssumptionCache::ResultElem, 16> Affected;
122 findAffectedValues(CI, TTI, Affected);
123
124 // If a value appears more than once in an AssumeInst e.g., 'ptr %arg1' in:
125 // call void @llvm.assume(i1 true)
126 // [ "dereferenceable"(ptr %arg1, i64 1),
127 // "align"(ptr %arg1, i64 8) ]
128 // it will appear multiple times in Affected, but we may (depending on
129 // how the results in AffectedValues.find_as(AV.Assume) are ordered)
130 // nullify multiple instances of Elem.Assume during one iteration of the
131 // 'for (auto &AV : Affected)' loop below. The next iteration of that for
132 // loop may then find only a match to a different AssumeInst, resulting in
133 // an assertion failure. Avoid this by counting the number of expected
134 // matches.
135#ifndef NDEBUG
136 SmallDenseSet<std::pair<Value *, unsigned>, 16> Seen;
137 DenseMap<Value *, int> ExpectedMatches;
138 for (auto &AV : Affected)
139 if (Seen.insert({AV.Assume, AV.Index}).second &&
140 AffectedValues.find_as(AV.Assume) != AffectedValues.end())
141 ExpectedMatches[AV.Assume]++;
142#endif
143
144 for (auto &AV : Affected) {
145 auto AVI = AffectedValues.find_as(Val: AV.Assume);
146 if (AVI == AffectedValues.end())
147 continue;
148 bool Found = false;
149 bool HasNonnull = false;
150 for (ResultElem &Elem : AVI->second) {
151 if (Elem.Assume == CI) {
152 Found = true;
153 Elem.Assume = nullptr;
154
155#ifndef NDEBUG
156 ExpectedMatches[AV.Assume]--;
157#endif
158 assert(ExpectedMatches[AV.Assume] >= 0);
159 // After ExpectedMatches[AV.Assume] == 0, we still need to iterate
160 // through this loop to determine the value of HasNonnull, to avoid
161 // prematurely calling AffectedValues.erase(AVI).
162 }
163 HasNonnull |= !!Elem.Assume;
164 if (HasNonnull && Found)
165 break;
166 }
167
168 assert(ExpectedMatches[AV.Assume] == 0 ||
169 Found && "already unregistered or incorrect cache state");
170
171 if (!HasNonnull)
172 AffectedValues.erase(I: AVI);
173 }
174
175 assert(
176 none_of(Affected, [&](auto &AV) { return ExpectedMatches[AV.Assume]; }) &&
177 "already unregistered or incorrect cache state");
178}
179
180void AssumptionCache::unregisterAssumption(AssumeInst *CI) {
181 removeAffectedValues(CI);
182 llvm::erase(C&: AssumeHandles, V: CI);
183}
184
185void AssumptionCache::replaceAssumption(WeakVH &Handle, AssumeInst *New) {
186 removeAffectedValues(CI: cast<AssumeInst>(Val&: Handle));
187 Handle = New;
188 updateAffectedValues(CI: New);
189}
190
191void AssumptionCache::AffectedValueCallbackVH::deleted() {
192 AC->AffectedValues.erase(Val: getValPtr());
193 // 'this' now dangles!
194}
195
196void AssumptionCache::transferAffectedValuesInCache(Value *OV, Value *NV) {
197 auto &NAVV = getOrInsertAffectedValues(V: NV);
198 auto AVI = AffectedValues.find(Val: OV);
199 if (AVI == AffectedValues.end())
200 return;
201
202 for (auto &A : AVI->second)
203 if (!llvm::is_contained(Range&: NAVV, Element: A))
204 NAVV.push_back(Elt: A);
205 AffectedValues.erase(Val: OV);
206}
207
208void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) {
209 if (!isa<Instruction>(Val: NV) && !isa<Argument>(Val: NV))
210 return;
211
212 // Any assumptions that affected this value now affect the new value.
213
214 AC->transferAffectedValuesInCache(OV: getValPtr(), NV);
215 // 'this' now might dangle! If the AffectedValues map was resized to add an
216 // entry for NV then this object might have been destroyed in favor of some
217 // copy in the grown map.
218}
219
220void AssumptionCache::scanFunction() {
221 assert(!Scanned && "Tried to scan the function twice!");
222 assert(AssumeHandles.empty() && "Already have assumes when scanning!");
223
224 // Go through all instructions in all blocks, add all calls to @llvm.assume
225 // to this cache.
226 for (BasicBlock &B : F)
227 for (Instruction &I : B)
228 if (isa<AssumeInst>(Val: &I))
229 AssumeHandles.push_back(Elt: &I);
230
231 // Mark the scan as complete.
232 Scanned = true;
233
234 // Update affected values.
235 for (auto &A : AssumeHandles)
236 updateAffectedValues(CI: cast<AssumeInst>(Val&: A));
237}
238
239void AssumptionCache::registerAssumption(AssumeInst *CI) {
240 // If we haven't scanned the function yet, just drop this assumption. It will
241 // be found when we scan later.
242 if (!Scanned)
243 return;
244
245 AssumeHandles.push_back(Elt: CI);
246
247#ifndef NDEBUG
248 assert(CI->getParent() &&
249 "Cannot register @llvm.assume call not in a basic block");
250 assert(&F == CI->getParent()->getParent() &&
251 "Cannot register @llvm.assume call not in this function");
252
253 // We expect the number of assumptions to be small, so in an asserts build
254 // check that we don't accumulate duplicates and that all assumptions point
255 // to the same function.
256 SmallPtrSet<Value *, 16> AssumptionSet;
257 for (auto &VH : AssumeHandles) {
258 if (!VH)
259 continue;
260
261 assert(&F == cast<Instruction>(VH)->getParent()->getParent() &&
262 "Cached assumption not inside this function!");
263 assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) &&
264 "Cached something other than a call to @llvm.assume!");
265 assert(AssumptionSet.insert(VH).second &&
266 "Cache contains multiple copies of a call!");
267 }
268#endif
269
270 updateAffectedValues(CI);
271}
272
273AssumptionCache AssumptionAnalysis::run(Function &F,
274 FunctionAnalysisManager &FAM) {
275 auto &TTI = FAM.getResult<TargetIRAnalysis>(IR&: F);
276 return AssumptionCache(F, &TTI);
277}
278
279AnalysisKey AssumptionAnalysis::Key;
280
281PreservedAnalyses AssumptionPrinterPass::run(Function &F,
282 FunctionAnalysisManager &AM) {
283 AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(IR&: F);
284
285 OS << "Cached assumptions for function: " << F.getName() << "\n";
286 for (auto &VH : AC.assumptions()) {
287 if (!VH)
288 continue;
289
290 auto *Assume = cast<CallInst>(Val&: VH);
291 if (!Assume->hasOperandBundles()) {
292 OS << " " << *Assume->getArgOperand(i: 0) << "\n";
293 continue;
294 }
295
296 assert(match(Assume->getArgOperand(0), m_One()) &&
297 "assume must have trivial cond");
298 OS << " [ ";
299 ListSeparator LS;
300 for (const OperandBundleUse &BU : Assume->operand_bundles()) {
301 OS << LS << '"' << BU.getTagName() << "\"(";
302 interleaveComma(c: BU.Inputs, os&: OS,
303 each_fn: [&](const Use &Input) { Input->printAsOperand(O&: OS); });
304 OS << ')';
305 }
306 OS << " ]\n";
307 }
308
309 return PreservedAnalyses::all();
310}
311
312void AssumptionCacheTracker::FunctionCallbackVH::deleted() {
313 auto I = ACT->AssumptionCaches.find_as(Val: cast<Function>(Val: getValPtr()));
314 if (I != ACT->AssumptionCaches.end())
315 ACT->AssumptionCaches.erase(I);
316 // 'this' now dangles!
317}
318
319AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) {
320 // We probe the function map twice to try and avoid creating a value handle
321 // around the function in common cases. This makes insertion a bit slower,
322 // but if we have to insert we're going to scan the whole function so that
323 // shouldn't matter.
324 auto I = AssumptionCaches.find_as(Val: &F);
325 if (I != AssumptionCaches.end())
326 return *I->second;
327
328 auto *TTIWP = getAnalysisIfAvailable<TargetTransformInfoWrapperPass>();
329 auto *TTI = TTIWP ? &TTIWP->getTTI(F) : nullptr;
330
331 // Ok, build a new cache by scanning the function, insert it and the value
332 // handle into our map, and return the newly populated cache.
333 auto IP = AssumptionCaches.insert(KV: std::make_pair(
334 x: FunctionCallbackVH(&F, this), y: std::make_unique<AssumptionCache>(args&: F, args&: TTI)));
335 assert(IP.second && "Scanning function already in the map?");
336 return *IP.first->second;
337}
338
339AssumptionCache *AssumptionCacheTracker::lookupAssumptionCache(Function &F) {
340 auto I = AssumptionCaches.find_as(Val: &F);
341 if (I != AssumptionCaches.end())
342 return I->second.get();
343 return nullptr;
344}
345
346void AssumptionCacheTracker::verifyAnalysis() const {
347 // FIXME: In the long term the verifier should not be controllable with a
348 // flag. We should either fix all passes to correctly update the assumption
349 // cache and enable the verifier unconditionally or somehow arrange for the
350 // assumption list to be updated automatically by passes.
351 if (!VerifyAssumptionCache)
352 return;
353
354 SmallPtrSet<const CallInst *, 4> AssumptionSet;
355 for (const auto &I : AssumptionCaches) {
356 for (auto &VH : I.second->assumptions())
357 if (VH)
358 AssumptionSet.insert(Ptr: cast<CallInst>(Val&: VH));
359
360 for (const BasicBlock &B : cast<Function>(Val&: *I.first))
361 for (const Instruction &II : B)
362 if (match(V: &II, P: m_Intrinsic<Intrinsic::assume>()) &&
363 !AssumptionSet.count(Ptr: cast<CallInst>(Val: &II)))
364 report_fatal_error(reason: "Assumption in scanned function not in cache");
365 }
366}
367
368AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) {}
369
370AssumptionCacheTracker::~AssumptionCacheTracker() = default;
371
372char AssumptionCacheTracker::ID = 0;
373
374INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker",
375 "Assumption Cache Tracker", false, true)
376