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