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 | |
37 | using namespace llvm; |
38 | using namespace llvm::PatternMatch; |
39 | |
40 | static cl::opt<bool> |
41 | VerifyAssumptionCache("verify-assumption-cache" , cl::Hidden, |
42 | cl::desc("Enable verification of assumption cache" ), |
43 | cl::init(Val: false)); |
44 | |
45 | SmallVector<AssumptionCache::ResultElem, 1> & |
46 | AssumptionCache::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 | |
58 | static void |
59 | findAffectedValues(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 InsertAffected = [&Affected](Value *V) { |
65 | Affected.push_back(Elt: {.Assume: V, .Index: AssumptionCache::ExprResultIdx}); |
66 | }; |
67 | |
68 | auto AddAffectedVal = [&Affected](Value *V, unsigned Idx) { |
69 | if (isa<Argument>(Val: V) || isa<GlobalValue>(Val: V) || isa<Instruction>(Val: V)) { |
70 | Affected.push_back(Elt: {.Assume: V, .Index: Idx}); |
71 | } |
72 | }; |
73 | |
74 | for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++) { |
75 | OperandBundleUse Bundle = CI->getOperandBundleAt(Index: Idx); |
76 | if (Bundle.getTagName() == "separate_storage" ) { |
77 | assert(Bundle.Inputs.size() == 2 && |
78 | "separate_storage must have two args" ); |
79 | AddAffectedVal(getUnderlyingObject(V: Bundle.Inputs[0]), Idx); |
80 | AddAffectedVal(getUnderlyingObject(V: Bundle.Inputs[1]), Idx); |
81 | } else if (Bundle.Inputs.size() > ABA_WasOn && |
82 | Bundle.getTagName() != IgnoreBundleTag) |
83 | AddAffectedVal(Bundle.Inputs[ABA_WasOn], Idx); |
84 | } |
85 | |
86 | Value *Cond = CI->getArgOperand(i: 0); |
87 | findValuesAffectedByCondition(Cond, /*IsAssume=*/true, InsertAffected); |
88 | |
89 | if (TTI) { |
90 | const Value *Ptr; |
91 | unsigned AS; |
92 | std::tie(args&: Ptr, args&: AS) = TTI->getPredicatedAddrSpace(V: Cond); |
93 | if (Ptr) |
94 | AddAffectedVal(const_cast<Value *>(Ptr->stripInBoundsOffsets()), |
95 | AssumptionCache::ExprResultIdx); |
96 | } |
97 | } |
98 | |
99 | void AssumptionCache::updateAffectedValues(AssumeInst *CI) { |
100 | SmallVector<AssumptionCache::ResultElem, 16> Affected; |
101 | findAffectedValues(CI, TTI, Affected); |
102 | |
103 | for (auto &AV : Affected) { |
104 | auto &AVV = getOrInsertAffectedValues(V: AV.Assume); |
105 | if (llvm::none_of(Range&: AVV, P: [&](ResultElem &Elem) { |
106 | return Elem.Assume == CI && Elem.Index == AV.Index; |
107 | })) |
108 | AVV.push_back(Elt: {.Assume: CI, .Index: AV.Index}); |
109 | } |
110 | } |
111 | |
112 | void AssumptionCache::unregisterAssumption(AssumeInst *CI) { |
113 | SmallVector<AssumptionCache::ResultElem, 16> Affected; |
114 | findAffectedValues(CI, TTI, Affected); |
115 | |
116 | for (auto &AV : Affected) { |
117 | auto AVI = AffectedValues.find_as(Val: AV.Assume); |
118 | if (AVI == AffectedValues.end()) |
119 | continue; |
120 | bool Found = false; |
121 | bool HasNonnull = false; |
122 | for (ResultElem &Elem : AVI->second) { |
123 | if (Elem.Assume == CI) { |
124 | Found = true; |
125 | Elem.Assume = nullptr; |
126 | } |
127 | HasNonnull |= !!Elem.Assume; |
128 | if (HasNonnull && Found) |
129 | break; |
130 | } |
131 | assert(Found && "already unregistered or incorrect cache state" ); |
132 | if (!HasNonnull) |
133 | AffectedValues.erase(I: AVI); |
134 | } |
135 | |
136 | llvm::erase(C&: AssumeHandles, V: CI); |
137 | } |
138 | |
139 | void AssumptionCache::AffectedValueCallbackVH::deleted() { |
140 | AC->AffectedValues.erase(Val: getValPtr()); |
141 | // 'this' now dangles! |
142 | } |
143 | |
144 | void AssumptionCache::transferAffectedValuesInCache(Value *OV, Value *NV) { |
145 | auto &NAVV = getOrInsertAffectedValues(V: NV); |
146 | auto AVI = AffectedValues.find(Val: OV); |
147 | if (AVI == AffectedValues.end()) |
148 | return; |
149 | |
150 | for (auto &A : AVI->second) |
151 | if (!llvm::is_contained(Range&: NAVV, Element: A)) |
152 | NAVV.push_back(Elt: A); |
153 | AffectedValues.erase(Val: OV); |
154 | } |
155 | |
156 | void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) { |
157 | if (!isa<Instruction>(Val: NV) && !isa<Argument>(Val: NV)) |
158 | return; |
159 | |
160 | // Any assumptions that affected this value now affect the new value. |
161 | |
162 | AC->transferAffectedValuesInCache(OV: getValPtr(), NV); |
163 | // 'this' now might dangle! If the AffectedValues map was resized to add an |
164 | // entry for NV then this object might have been destroyed in favor of some |
165 | // copy in the grown map. |
166 | } |
167 | |
168 | void AssumptionCache::scanFunction() { |
169 | assert(!Scanned && "Tried to scan the function twice!" ); |
170 | assert(AssumeHandles.empty() && "Already have assumes when scanning!" ); |
171 | |
172 | // Go through all instructions in all blocks, add all calls to @llvm.assume |
173 | // to this cache. |
174 | for (BasicBlock &B : F) |
175 | for (Instruction &I : B) |
176 | if (isa<AssumeInst>(Val: &I)) |
177 | AssumeHandles.push_back(Elt: {.Assume: &I, .Index: ExprResultIdx}); |
178 | |
179 | // Mark the scan as complete. |
180 | Scanned = true; |
181 | |
182 | // Update affected values. |
183 | for (auto &A : AssumeHandles) |
184 | updateAffectedValues(CI: cast<AssumeInst>(Val&: A)); |
185 | } |
186 | |
187 | void AssumptionCache::registerAssumption(AssumeInst *CI) { |
188 | // If we haven't scanned the function yet, just drop this assumption. It will |
189 | // be found when we scan later. |
190 | if (!Scanned) |
191 | return; |
192 | |
193 | AssumeHandles.push_back(Elt: {.Assume: CI, .Index: ExprResultIdx}); |
194 | |
195 | #ifndef NDEBUG |
196 | assert(CI->getParent() && |
197 | "Cannot register @llvm.assume call not in a basic block" ); |
198 | assert(&F == CI->getParent()->getParent() && |
199 | "Cannot register @llvm.assume call not in this function" ); |
200 | |
201 | // We expect the number of assumptions to be small, so in an asserts build |
202 | // check that we don't accumulate duplicates and that all assumptions point |
203 | // to the same function. |
204 | SmallPtrSet<Value *, 16> AssumptionSet; |
205 | for (auto &VH : AssumeHandles) { |
206 | if (!VH) |
207 | continue; |
208 | |
209 | assert(&F == cast<Instruction>(VH)->getParent()->getParent() && |
210 | "Cached assumption not inside this function!" ); |
211 | assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) && |
212 | "Cached something other than a call to @llvm.assume!" ); |
213 | assert(AssumptionSet.insert(VH).second && |
214 | "Cache contains multiple copies of a call!" ); |
215 | } |
216 | #endif |
217 | |
218 | updateAffectedValues(CI); |
219 | } |
220 | |
221 | AssumptionCache AssumptionAnalysis::run(Function &F, |
222 | FunctionAnalysisManager &FAM) { |
223 | auto &TTI = FAM.getResult<TargetIRAnalysis>(IR&: F); |
224 | return AssumptionCache(F, &TTI); |
225 | } |
226 | |
227 | AnalysisKey AssumptionAnalysis::Key; |
228 | |
229 | PreservedAnalyses AssumptionPrinterPass::run(Function &F, |
230 | FunctionAnalysisManager &AM) { |
231 | AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(IR&: F); |
232 | |
233 | OS << "Cached assumptions for function: " << F.getName() << "\n" ; |
234 | for (auto &VH : AC.assumptions()) |
235 | if (VH) |
236 | OS << " " << *cast<CallInst>(Val&: VH)->getArgOperand(i: 0) << "\n" ; |
237 | |
238 | return PreservedAnalyses::all(); |
239 | } |
240 | |
241 | void AssumptionCacheTracker::FunctionCallbackVH::deleted() { |
242 | auto I = ACT->AssumptionCaches.find_as(Val: cast<Function>(Val: getValPtr())); |
243 | if (I != ACT->AssumptionCaches.end()) |
244 | ACT->AssumptionCaches.erase(I); |
245 | // 'this' now dangles! |
246 | } |
247 | |
248 | AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) { |
249 | // We probe the function map twice to try and avoid creating a value handle |
250 | // around the function in common cases. This makes insertion a bit slower, |
251 | // but if we have to insert we're going to scan the whole function so that |
252 | // shouldn't matter. |
253 | auto I = AssumptionCaches.find_as(Val: &F); |
254 | if (I != AssumptionCaches.end()) |
255 | return *I->second; |
256 | |
257 | auto *TTIWP = getAnalysisIfAvailable<TargetTransformInfoWrapperPass>(); |
258 | auto *TTI = TTIWP ? &TTIWP->getTTI(F) : nullptr; |
259 | |
260 | // Ok, build a new cache by scanning the function, insert it and the value |
261 | // handle into our map, and return the newly populated cache. |
262 | auto IP = AssumptionCaches.insert(KV: std::make_pair( |
263 | x: FunctionCallbackVH(&F, this), y: std::make_unique<AssumptionCache>(args&: F, args&: TTI))); |
264 | assert(IP.second && "Scanning function already in the map?" ); |
265 | return *IP.first->second; |
266 | } |
267 | |
268 | AssumptionCache *AssumptionCacheTracker::lookupAssumptionCache(Function &F) { |
269 | auto I = AssumptionCaches.find_as(Val: &F); |
270 | if (I != AssumptionCaches.end()) |
271 | return I->second.get(); |
272 | return nullptr; |
273 | } |
274 | |
275 | void AssumptionCacheTracker::verifyAnalysis() const { |
276 | // FIXME: In the long term the verifier should not be controllable with a |
277 | // flag. We should either fix all passes to correctly update the assumption |
278 | // cache and enable the verifier unconditionally or somehow arrange for the |
279 | // assumption list to be updated automatically by passes. |
280 | if (!VerifyAssumptionCache) |
281 | return; |
282 | |
283 | SmallPtrSet<const CallInst *, 4> AssumptionSet; |
284 | for (const auto &I : AssumptionCaches) { |
285 | for (auto &VH : I.second->assumptions()) |
286 | if (VH) |
287 | AssumptionSet.insert(Ptr: cast<CallInst>(Val&: VH)); |
288 | |
289 | for (const BasicBlock &B : cast<Function>(Val&: *I.first)) |
290 | for (const Instruction &II : B) |
291 | if (match(V: &II, P: m_Intrinsic<Intrinsic::assume>()) && |
292 | !AssumptionSet.count(Ptr: cast<CallInst>(Val: &II))) |
293 | report_fatal_error(reason: "Assumption in scanned function not in cache" ); |
294 | } |
295 | } |
296 | |
297 | AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) { |
298 | initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry()); |
299 | } |
300 | |
301 | AssumptionCacheTracker::~AssumptionCacheTracker() = default; |
302 | |
303 | char AssumptionCacheTracker::ID = 0; |
304 | |
305 | INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker" , |
306 | "Assumption Cache Tracker" , false, true) |
307 | |