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 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
99void 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
112void 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
139void AssumptionCache::AffectedValueCallbackVH::deleted() {
140 AC->AffectedValues.erase(Val: getValPtr());
141 // 'this' now dangles!
142}
143
144void 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
156void 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
168void 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
187void 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
221AssumptionCache AssumptionAnalysis::run(Function &F,
222 FunctionAnalysisManager &FAM) {
223 auto &TTI = FAM.getResult<TargetIRAnalysis>(IR&: F);
224 return AssumptionCache(F, &TTI);
225}
226
227AnalysisKey AssumptionAnalysis::Key;
228
229PreservedAnalyses 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
241void 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
248AssumptionCache &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
268AssumptionCache *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
275void 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
297AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) {
298 initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry());
299}
300
301AssumptionCacheTracker::~AssumptionCacheTracker() = default;
302
303char AssumptionCacheTracker::ID = 0;
304
305INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker",
306 "Assumption Cache Tracker", false, true)
307