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