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/Analysis/AssumeBundleQueries.h"
19#include "llvm/Analysis/TargetTransformInfo.h"
20#include "llvm/Analysis/ValueTracking.h"
21#include "llvm/IR/BasicBlock.h"
22#include "llvm/IR/Function.h"
23#include "llvm/IR/InstrTypes.h"
24#include "llvm/IR/Instruction.h"
25#include "llvm/IR/Instructions.h"
26#include "llvm/IR/PassManager.h"
27#include "llvm/IR/PatternMatch.h"
28#include "llvm/InitializePasses.h"
29#include "llvm/Pass.h"
30#include "llvm/Support/Casting.h"
31#include "llvm/Support/CommandLine.h"
32#include "llvm/Support/ErrorHandling.h"
33#include "llvm/Support/raw_ostream.h"
34#include <cassert>
35#include <utility>
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
45SmallVector<AssumptionCache::ResultElem, 1> &
46AssumptionCache::getOrInsertAffectedValues(Value *V) {
47 // Try using find_as first to avoid creating extra value handles just for the
48 // purpose of doing the lookup.
49 auto AVI = AffectedValues.find_as(Val: V);
50 if (AVI != AffectedValues.end())
51 return AVI->second;
52
53 auto AVIP = AffectedValues.insert(
54 KV: {AffectedValueCallbackVH(V, this), SmallVector<ResultElem, 1>()});
55 return AVIP.first->second;
56}
57
58static void
59findAffectedValues(CallBase *CI, TargetTransformInfo *TTI,
60 SmallVectorImpl<AssumptionCache::ResultElem> &Affected) {
61 // Note: This code must be kept in-sync with the code in
62 // computeKnownBitsFromAssume in ValueTracking.
63
64 auto AddAffected = [&Affected](Value *V, unsigned Idx =
65 AssumptionCache::ExprResultIdx) {
66 if (isa<Argument>(Val: V) || isa<GlobalValue>(Val: V)) {
67 Affected.push_back(Elt: {.Assume: V, .Index: Idx});
68 } else if (auto *I = dyn_cast<Instruction>(Val: V)) {
69 Affected.push_back(Elt: {.Assume: I, .Index: Idx});
70
71 // Peek through unary operators to find the source of the condition.
72 Value *Op;
73 if (match(V: I, P: m_PtrToInt(Op: m_Value(V&: Op)))) {
74 if (isa<Instruction>(Val: Op) || isa<Argument>(Val: Op))
75 Affected.push_back(Elt: {.Assume: Op, .Index: Idx});
76 }
77 }
78 };
79
80 for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++) {
81 OperandBundleUse Bundle = CI->getOperandBundleAt(Index: Idx);
82 if (Bundle.getTagName() == "separate_storage") {
83 assert(Bundle.Inputs.size() == 2 &&
84 "separate_storage must have two args");
85 AddAffected(getUnderlyingObject(V: Bundle.Inputs[0]), Idx);
86 AddAffected(getUnderlyingObject(V: Bundle.Inputs[1]), Idx);
87 } else if (Bundle.Inputs.size() > ABA_WasOn &&
88 Bundle.getTagName() != IgnoreBundleTag)
89 AddAffected(Bundle.Inputs[ABA_WasOn], Idx);
90 }
91
92 Value *Cond = CI->getArgOperand(i: 0), *A, *B;
93 AddAffected(Cond);
94 if (match(V: Cond, P: m_Not(V: m_Value(V&: A))))
95 AddAffected(A);
96
97 CmpInst::Predicate Pred;
98 if (match(V: Cond, P: m_Cmp(Pred, L: m_Value(V&: A), R: m_Value(V&: B)))) {
99 AddAffected(A);
100 AddAffected(B);
101
102 if (Pred == ICmpInst::ICMP_EQ) {
103 if (match(V: B, P: m_ConstantInt())) {
104 Value *X;
105 // (X & C) or (X | C) or (X ^ C).
106 // (X << C) or (X >>_s C) or (X >>_u C).
107 if (match(V: A, P: m_BitwiseLogic(L: m_Value(V&: X), R: m_ConstantInt())) ||
108 match(V: A, P: m_Shift(L: m_Value(V&: X), R: m_ConstantInt())))
109 AddAffected(X);
110 }
111 } else if (Pred == ICmpInst::ICMP_NE) {
112 Value *X;
113 // Handle (X & pow2 != 0).
114 if (match(V: A, P: m_And(L: m_Value(V&: X), R: m_Power2())) && match(V: B, P: m_Zero()))
115 AddAffected(X);
116 } else if (Pred == ICmpInst::ICMP_ULT) {
117 Value *X;
118 // Handle (A + C1) u< C2, which is the canonical form of A > C3 && A < C4,
119 // and recognized by LVI at least.
120 if (match(V: A, P: m_Add(L: m_Value(V&: X), R: m_ConstantInt())) &&
121 match(V: B, P: m_ConstantInt()))
122 AddAffected(X);
123 } else if (CmpInst::isFPPredicate(P: Pred)) {
124 // fcmp fneg(x), y
125 // fcmp fabs(x), y
126 // fcmp fneg(fabs(x)), y
127 if (match(V: A, P: m_FNeg(X: m_Value(V&: A))))
128 AddAffected(A);
129 if (match(V: A, P: m_FAbs(Op0: m_Value(V&: A))))
130 AddAffected(A);
131 }
132 } else if (match(Cond, m_Intrinsic<Intrinsic::is_fpclass>(m_Value(V&: A),
133 m_Value(V&: B)))) {
134 AddAffected(A);
135 }
136
137 if (TTI) {
138 const Value *Ptr;
139 unsigned AS;
140 std::tie(args&: Ptr, args&: AS) = TTI->getPredicatedAddrSpace(V: Cond);
141 if (Ptr)
142 AddAffected(const_cast<Value *>(Ptr->stripInBoundsOffsets()));
143 }
144}
145
146void AssumptionCache::updateAffectedValues(AssumeInst *CI) {
147 SmallVector<AssumptionCache::ResultElem, 16> Affected;
148 findAffectedValues(CI, TTI, Affected);
149
150 for (auto &AV : Affected) {
151 auto &AVV = getOrInsertAffectedValues(V: AV.Assume);
152 if (llvm::none_of(Range&: AVV, P: [&](ResultElem &Elem) {
153 return Elem.Assume == CI && Elem.Index == AV.Index;
154 }))
155 AVV.push_back(Elt: {.Assume: CI, .Index: AV.Index});
156 }
157}
158
159void AssumptionCache::unregisterAssumption(AssumeInst *CI) {
160 SmallVector<AssumptionCache::ResultElem, 16> Affected;
161 findAffectedValues(CI, TTI, Affected);
162
163 for (auto &AV : Affected) {
164 auto AVI = AffectedValues.find_as(Val: AV.Assume);
165 if (AVI == AffectedValues.end())
166 continue;
167 bool Found = false;
168 bool HasNonnull = false;
169 for (ResultElem &Elem : AVI->second) {
170 if (Elem.Assume == CI) {
171 Found = true;
172 Elem.Assume = nullptr;
173 }
174 HasNonnull |= !!Elem.Assume;
175 if (HasNonnull && Found)
176 break;
177 }
178 assert(Found && "already unregistered or incorrect cache state");
179 if (!HasNonnull)
180 AffectedValues.erase(I: AVI);
181 }
182
183 llvm::erase(C&: AssumeHandles, V: CI);
184}
185
186void AssumptionCache::AffectedValueCallbackVH::deleted() {
187 AC->AffectedValues.erase(Val: getValPtr());
188 // 'this' now dangles!
189}
190
191void AssumptionCache::transferAffectedValuesInCache(Value *OV, Value *NV) {
192 auto &NAVV = getOrInsertAffectedValues(V: NV);
193 auto AVI = AffectedValues.find(Val: OV);
194 if (AVI == AffectedValues.end())
195 return;
196
197 for (auto &A : AVI->second)
198 if (!llvm::is_contained(Range&: NAVV, Element: A))
199 NAVV.push_back(Elt: A);
200 AffectedValues.erase(Val: OV);
201}
202
203void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) {
204 if (!isa<Instruction>(Val: NV) && !isa<Argument>(Val: NV))
205 return;
206
207 // Any assumptions that affected this value now affect the new value.
208
209 AC->transferAffectedValuesInCache(OV: getValPtr(), NV);
210 // 'this' now might dangle! If the AffectedValues map was resized to add an
211 // entry for NV then this object might have been destroyed in favor of some
212 // copy in the grown map.
213}
214
215void AssumptionCache::scanFunction() {
216 assert(!Scanned && "Tried to scan the function twice!");
217 assert(AssumeHandles.empty() && "Already have assumes when scanning!");
218
219 // Go through all instructions in all blocks, add all calls to @llvm.assume
220 // to this cache.
221 for (BasicBlock &B : F)
222 for (Instruction &I : B)
223 if (isa<AssumeInst>(Val: &I))
224 AssumeHandles.push_back(Elt: {.Assume: &I, .Index: ExprResultIdx});
225
226 // Mark the scan as complete.
227 Scanned = true;
228
229 // Update affected values.
230 for (auto &A : AssumeHandles)
231 updateAffectedValues(CI: cast<AssumeInst>(Val&: A));
232}
233
234void AssumptionCache::registerAssumption(AssumeInst *CI) {
235 // If we haven't scanned the function yet, just drop this assumption. It will
236 // be found when we scan later.
237 if (!Scanned)
238 return;
239
240 AssumeHandles.push_back(Elt: {.Assume: CI, .Index: ExprResultIdx});
241
242#ifndef NDEBUG
243 assert(CI->getParent() &&
244 "Cannot register @llvm.assume call not in a basic block");
245 assert(&F == CI->getParent()->getParent() &&
246 "Cannot register @llvm.assume call not in this function");
247
248 // We expect the number of assumptions to be small, so in an asserts build
249 // check that we don't accumulate duplicates and that all assumptions point
250 // to the same function.
251 SmallPtrSet<Value *, 16> AssumptionSet;
252 for (auto &VH : AssumeHandles) {
253 if (!VH)
254 continue;
255
256 assert(&F == cast<Instruction>(VH)->getParent()->getParent() &&
257 "Cached assumption not inside this function!");
258 assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) &&
259 "Cached something other than a call to @llvm.assume!");
260 assert(AssumptionSet.insert(VH).second &&
261 "Cache contains multiple copies of a call!");
262 }
263#endif
264
265 updateAffectedValues(CI);
266}
267
268AssumptionCache AssumptionAnalysis::run(Function &F,
269 FunctionAnalysisManager &FAM) {
270 auto &TTI = FAM.getResult<TargetIRAnalysis>(IR&: F);
271 return AssumptionCache(F, &TTI);
272}
273
274AnalysisKey AssumptionAnalysis::Key;
275
276PreservedAnalyses AssumptionPrinterPass::run(Function &F,
277 FunctionAnalysisManager &AM) {
278 AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(IR&: F);
279
280 OS << "Cached assumptions for function: " << F.getName() << "\n";
281 for (auto &VH : AC.assumptions())
282 if (VH)
283 OS << " " << *cast<CallInst>(Val&: VH)->getArgOperand(i: 0) << "\n";
284
285 return PreservedAnalyses::all();
286}
287
288void AssumptionCacheTracker::FunctionCallbackVH::deleted() {
289 auto I = ACT->AssumptionCaches.find_as(Val: cast<Function>(Val: getValPtr()));
290 if (I != ACT->AssumptionCaches.end())
291 ACT->AssumptionCaches.erase(I);
292 // 'this' now dangles!
293}
294
295AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) {
296 // We probe the function map twice to try and avoid creating a value handle
297 // around the function in common cases. This makes insertion a bit slower,
298 // but if we have to insert we're going to scan the whole function so that
299 // shouldn't matter.
300 auto I = AssumptionCaches.find_as(Val: &F);
301 if (I != AssumptionCaches.end())
302 return *I->second;
303
304 auto *TTIWP = getAnalysisIfAvailable<TargetTransformInfoWrapperPass>();
305 auto *TTI = TTIWP ? &TTIWP->getTTI(F) : nullptr;
306
307 // Ok, build a new cache by scanning the function, insert it and the value
308 // handle into our map, and return the newly populated cache.
309 auto IP = AssumptionCaches.insert(KV: std::make_pair(
310 x: FunctionCallbackVH(&F, this), y: std::make_unique<AssumptionCache>(args&: F, args&: TTI)));
311 assert(IP.second && "Scanning function already in the map?");
312 return *IP.first->second;
313}
314
315AssumptionCache *AssumptionCacheTracker::lookupAssumptionCache(Function &F) {
316 auto I = AssumptionCaches.find_as(Val: &F);
317 if (I != AssumptionCaches.end())
318 return I->second.get();
319 return nullptr;
320}
321
322void AssumptionCacheTracker::verifyAnalysis() const {
323 // FIXME: In the long term the verifier should not be controllable with a
324 // flag. We should either fix all passes to correctly update the assumption
325 // cache and enable the verifier unconditionally or somehow arrange for the
326 // assumption list to be updated automatically by passes.
327 if (!VerifyAssumptionCache)
328 return;
329
330 SmallPtrSet<const CallInst *, 4> AssumptionSet;
331 for (const auto &I : AssumptionCaches) {
332 for (auto &VH : I.second->assumptions())
333 if (VH)
334 AssumptionSet.insert(Ptr: cast<CallInst>(Val&: VH));
335
336 for (const BasicBlock &B : cast<Function>(Val&: *I.first))
337 for (const Instruction &II : B)
338 if (match(&II, m_Intrinsic<Intrinsic::assume>()) &&
339 !AssumptionSet.count(cast<CallInst>(&II)))
340 report_fatal_error(reason: "Assumption in scanned function not in cache");
341 }
342}
343
344AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) {
345 initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry());
346}
347
348AssumptionCacheTracker::~AssumptionCacheTracker() = default;
349
350char AssumptionCacheTracker::ID = 0;
351
352INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker",
353 "Assumption Cache Tracker", false, true)
354

source code of llvm/lib/Analysis/AssumptionCache.cpp