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/IR/BasicBlock.h"
19 #include "llvm/IR/Function.h"
20 #include "llvm/IR/InstrTypes.h"
21 #include "llvm/IR/Instruction.h"
22 #include "llvm/IR/Instructions.h"
23 #include "llvm/IR/Intrinsics.h"
24 #include "llvm/IR/PassManager.h"
25 #include "llvm/IR/PatternMatch.h"
26 #include "llvm/Pass.h"
27 #include "llvm/Support/Casting.h"
28 #include "llvm/Support/CommandLine.h"
29 #include "llvm/Support/ErrorHandling.h"
30 #include "llvm/Support/raw_ostream.h"
31 #include <algorithm>
32 #include <cassert>
33 #include <utility>
34 
35 using namespace llvm;
36 using namespace llvm::PatternMatch;
37 
38 static cl::opt<bool>
39     VerifyAssumptionCache("verify-assumption-cache", cl::Hidden,
40                           cl::desc("Enable verification of assumption cache"),
41                           cl::init(false));
42 
43 SmallVector<WeakTrackingVH, 1> &
44 AssumptionCache::getOrInsertAffectedValues(Value *V) {
45   // Try using find_as first to avoid creating extra value handles just for the
46   // purpose of doing the lookup.
47   auto AVI = AffectedValues.find_as(V);
48   if (AVI != AffectedValues.end())
49     return AVI->second;
50 
51   auto AVIP = AffectedValues.insert(
52       {AffectedValueCallbackVH(V, this), SmallVector<WeakTrackingVH, 1>()});
53   return AVIP.first->second;
54 }
55 
56 void AssumptionCache::updateAffectedValues(CallInst *CI) {
57   // Note: This code must be kept in-sync with the code in
58   // computeKnownBitsFromAssume in ValueTracking.
59 
60   SmallVector<Value *, 16> Affected;
61   auto AddAffected = [&Affected](Value *V) {
62     if (isa<Argument>(V)) {
63       Affected.push_back(V);
64     } else if (auto *I = dyn_cast<Instruction>(V)) {
65       Affected.push_back(I);
66 
67       // Peek through unary operators to find the source of the condition.
68       Value *Op;
69       if (match(I, m_BitCast(m_Value(Op))) ||
70           match(I, m_PtrToInt(m_Value(Op))) ||
71           match(I, m_Not(m_Value(Op)))) {
72         if (isa<Instruction>(Op) || isa<Argument>(Op))
73           Affected.push_back(Op);
74       }
75     }
76   };
77 
78   Value *Cond = CI->getArgOperand(0), *A, *B;
79   AddAffected(Cond);
80 
81   CmpInst::Predicate Pred;
82   if (match(Cond, m_ICmp(Pred, m_Value(A), m_Value(B)))) {
83     AddAffected(A);
84     AddAffected(B);
85 
86     if (Pred == ICmpInst::ICMP_EQ) {
87       // For equality comparisons, we handle the case of bit inversion.
88       auto AddAffectedFromEq = [&AddAffected](Value *V) {
89         Value *A;
90         if (match(V, m_Not(m_Value(A)))) {
91           AddAffected(A);
92           V = A;
93         }
94 
95         Value *B;
96         ConstantInt *C;
97         // (A & B) or (A | B) or (A ^ B).
98         if (match(V, m_BitwiseLogic(m_Value(A), m_Value(B)))) {
99           AddAffected(A);
100           AddAffected(B);
101         // (A << C) or (A >>_s C) or (A >>_u C) where C is some constant.
102         } else if (match(V, m_Shift(m_Value(A), m_ConstantInt(C)))) {
103           AddAffected(A);
104         }
105       };
106 
107       AddAffectedFromEq(A);
108       AddAffectedFromEq(B);
109     }
110   }
111 
112   for (auto &AV : Affected) {
113     auto &AVV = getOrInsertAffectedValues(AV);
114     if (std::find(AVV.begin(), AVV.end(), CI) == AVV.end())
115       AVV.push_back(CI);
116   }
117 }
118 
119 void AssumptionCache::AffectedValueCallbackVH::deleted() {
120   auto AVI = AC->AffectedValues.find(getValPtr());
121   if (AVI != AC->AffectedValues.end())
122     AC->AffectedValues.erase(AVI);
123   // 'this' now dangles!
124 }
125 
126 void AssumptionCache::copyAffectedValuesInCache(Value *OV, Value *NV) {
127   auto &NAVV = getOrInsertAffectedValues(NV);
128   auto AVI = AffectedValues.find(OV);
129   if (AVI == AffectedValues.end())
130     return;
131 
132   for (auto &A : AVI->second)
133     if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end())
134       NAVV.push_back(A);
135 }
136 
137 void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) {
138   if (!isa<Instruction>(NV) && !isa<Argument>(NV))
139     return;
140 
141   // Any assumptions that affected this value now affect the new value.
142 
143   AC->copyAffectedValuesInCache(getValPtr(), NV);
144   // 'this' now might dangle! If the AffectedValues map was resized to add an
145   // entry for NV then this object might have been destroyed in favor of some
146   // copy in the grown map.
147 }
148 
149 void AssumptionCache::scanFunction() {
150   assert(!Scanned && "Tried to scan the function twice!");
151   assert(AssumeHandles.empty() && "Already have assumes when scanning!");
152 
153   // Go through all instructions in all blocks, add all calls to @llvm.assume
154   // to this cache.
155   for (BasicBlock &B : F)
156     for (Instruction &II : B)
157       if (match(&II, m_Intrinsic<Intrinsic::assume>()))
158         AssumeHandles.push_back(&II);
159 
160   // Mark the scan as complete.
161   Scanned = true;
162 
163   // Update affected values.
164   for (auto &A : AssumeHandles)
165     updateAffectedValues(cast<CallInst>(A));
166 }
167 
168 void AssumptionCache::registerAssumption(CallInst *CI) {
169   assert(match(CI, m_Intrinsic<Intrinsic::assume>()) &&
170          "Registered call does not call @llvm.assume");
171 
172   // If we haven't scanned the function yet, just drop this assumption. It will
173   // be found when we scan later.
174   if (!Scanned)
175     return;
176 
177   AssumeHandles.push_back(CI);
178 
179 #ifndef NDEBUG
180   assert(CI->getParent() &&
181          "Cannot register @llvm.assume call not in a basic block");
182   assert(&F == CI->getParent()->getParent() &&
183          "Cannot register @llvm.assume call not in this function");
184 
185   // We expect the number of assumptions to be small, so in an asserts build
186   // check that we don't accumulate duplicates and that all assumptions point
187   // to the same function.
188   SmallPtrSet<Value *, 16> AssumptionSet;
189   for (auto &VH : AssumeHandles) {
190     if (!VH)
191       continue;
192 
193     assert(&F == cast<Instruction>(VH)->getParent()->getParent() &&
194            "Cached assumption not inside this function!");
195     assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) &&
196            "Cached something other than a call to @llvm.assume!");
197     assert(AssumptionSet.insert(VH).second &&
198            "Cache contains multiple copies of a call!");
199   }
200 #endif
201 
202   updateAffectedValues(CI);
203 }
204 
205 AnalysisKey AssumptionAnalysis::Key;
206 
207 PreservedAnalyses AssumptionPrinterPass::run(Function &F,
208                                              FunctionAnalysisManager &AM) {
209   AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
210 
211   OS << "Cached assumptions for function: " << F.getName() << "\n";
212   for (auto &VH : AC.assumptions())
213     if (VH)
214       OS << "  " << *cast<CallInst>(VH)->getArgOperand(0) << "\n";
215 
216   return PreservedAnalyses::all();
217 }
218 
219 void AssumptionCacheTracker::FunctionCallbackVH::deleted() {
220   auto I = ACT->AssumptionCaches.find_as(cast<Function>(getValPtr()));
221   if (I != ACT->AssumptionCaches.end())
222     ACT->AssumptionCaches.erase(I);
223   // 'this' now dangles!
224 }
225 
226 AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) {
227   // We probe the function map twice to try and avoid creating a value handle
228   // around the function in common cases. This makes insertion a bit slower,
229   // but if we have to insert we're going to scan the whole function so that
230   // shouldn't matter.
231   auto I = AssumptionCaches.find_as(&F);
232   if (I != AssumptionCaches.end())
233     return *I->second;
234 
235   // Ok, build a new cache by scanning the function, insert it and the value
236   // handle into our map, and return the newly populated cache.
237   auto IP = AssumptionCaches.insert(std::make_pair(
238       FunctionCallbackVH(&F, this), llvm::make_unique<AssumptionCache>(F)));
239   assert(IP.second && "Scanning function already in the map?");
240   return *IP.first->second;
241 }
242 
243 void AssumptionCacheTracker::verifyAnalysis() const {
244   // FIXME: In the long term the verifier should not be controllable with a
245   // flag. We should either fix all passes to correctly update the assumption
246   // cache and enable the verifier unconditionally or somehow arrange for the
247   // assumption list to be updated automatically by passes.
248   if (!VerifyAssumptionCache)
249     return;
250 
251   SmallPtrSet<const CallInst *, 4> AssumptionSet;
252   for (const auto &I : AssumptionCaches) {
253     for (auto &VH : I.second->assumptions())
254       if (VH)
255         AssumptionSet.insert(cast<CallInst>(VH));
256 
257     for (const BasicBlock &B : cast<Function>(*I.first))
258       for (const Instruction &II : B)
259         if (match(&II, m_Intrinsic<Intrinsic::assume>()) &&
260             !AssumptionSet.count(cast<CallInst>(&II)))
261           report_fatal_error("Assumption in scanned function not in cache");
262   }
263 }
264 
265 AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) {
266   initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry());
267 }
268 
269 AssumptionCacheTracker::~AssumptionCacheTracker() = default;
270 
271 char AssumptionCacheTracker::ID = 0;
272 
273 INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker",
274                 "Assumption Cache Tracker", false, true)
275