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 static void findAffectedValues(CallInst *CI, 57 SmallVectorImpl<Value *> &Affected) { 58 // Note: This code must be kept in-sync with the code in 59 // computeKnownBitsFromAssume in ValueTracking. 60 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 113 void AssumptionCache::updateAffectedValues(CallInst *CI) { 114 SmallVector<Value *, 16> Affected; 115 findAffectedValues(CI, Affected); 116 117 for (auto &AV : Affected) { 118 auto &AVV = getOrInsertAffectedValues(AV); 119 if (std::find(AVV.begin(), AVV.end(), CI) == AVV.end()) 120 AVV.push_back(CI); 121 } 122 } 123 124 void AssumptionCache::unregisterAssumption(CallInst *CI) { 125 SmallVector<Value *, 16> Affected; 126 findAffectedValues(CI, Affected); 127 128 for (auto &AV : Affected) { 129 auto AVI = AffectedValues.find_as(AV); 130 if (AVI != AffectedValues.end()) 131 AffectedValues.erase(AVI); 132 } 133 remove_if(AssumeHandles, [CI](WeakTrackingVH &VH) { return CI == VH; }); 134 } 135 136 void AssumptionCache::AffectedValueCallbackVH::deleted() { 137 auto AVI = AC->AffectedValues.find(getValPtr()); 138 if (AVI != AC->AffectedValues.end()) 139 AC->AffectedValues.erase(AVI); 140 // 'this' now dangles! 141 } 142 143 void AssumptionCache::copyAffectedValuesInCache(Value *OV, Value *NV) { 144 auto &NAVV = getOrInsertAffectedValues(NV); 145 auto AVI = AffectedValues.find(OV); 146 if (AVI == AffectedValues.end()) 147 return; 148 149 for (auto &A : AVI->second) 150 if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end()) 151 NAVV.push_back(A); 152 } 153 154 void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) { 155 if (!isa<Instruction>(NV) && !isa<Argument>(NV)) 156 return; 157 158 // Any assumptions that affected this value now affect the new value. 159 160 AC->copyAffectedValuesInCache(getValPtr(), NV); 161 // 'this' now might dangle! If the AffectedValues map was resized to add an 162 // entry for NV then this object might have been destroyed in favor of some 163 // copy in the grown map. 164 } 165 166 void AssumptionCache::scanFunction() { 167 assert(!Scanned && "Tried to scan the function twice!"); 168 assert(AssumeHandles.empty() && "Already have assumes when scanning!"); 169 170 // Go through all instructions in all blocks, add all calls to @llvm.assume 171 // to this cache. 172 for (BasicBlock &B : F) 173 for (Instruction &II : B) 174 if (match(&II, m_Intrinsic<Intrinsic::assume>())) 175 AssumeHandles.push_back(&II); 176 177 // Mark the scan as complete. 178 Scanned = true; 179 180 // Update affected values. 181 for (auto &A : AssumeHandles) 182 updateAffectedValues(cast<CallInst>(A)); 183 } 184 185 void AssumptionCache::registerAssumption(CallInst *CI) { 186 assert(match(CI, m_Intrinsic<Intrinsic::assume>()) && 187 "Registered call does not call @llvm.assume"); 188 189 // If we haven't scanned the function yet, just drop this assumption. It will 190 // be found when we scan later. 191 if (!Scanned) 192 return; 193 194 AssumeHandles.push_back(CI); 195 196 #ifndef NDEBUG 197 assert(CI->getParent() && 198 "Cannot register @llvm.assume call not in a basic block"); 199 assert(&F == CI->getParent()->getParent() && 200 "Cannot register @llvm.assume call not in this function"); 201 202 // We expect the number of assumptions to be small, so in an asserts build 203 // check that we don't accumulate duplicates and that all assumptions point 204 // to the same function. 205 SmallPtrSet<Value *, 16> AssumptionSet; 206 for (auto &VH : AssumeHandles) { 207 if (!VH) 208 continue; 209 210 assert(&F == cast<Instruction>(VH)->getParent()->getParent() && 211 "Cached assumption not inside this function!"); 212 assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) && 213 "Cached something other than a call to @llvm.assume!"); 214 assert(AssumptionSet.insert(VH).second && 215 "Cache contains multiple copies of a call!"); 216 } 217 #endif 218 219 updateAffectedValues(CI); 220 } 221 222 AnalysisKey AssumptionAnalysis::Key; 223 224 PreservedAnalyses AssumptionPrinterPass::run(Function &F, 225 FunctionAnalysisManager &AM) { 226 AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F); 227 228 OS << "Cached assumptions for function: " << F.getName() << "\n"; 229 for (auto &VH : AC.assumptions()) 230 if (VH) 231 OS << " " << *cast<CallInst>(VH)->getArgOperand(0) << "\n"; 232 233 return PreservedAnalyses::all(); 234 } 235 236 void AssumptionCacheTracker::FunctionCallbackVH::deleted() { 237 auto I = ACT->AssumptionCaches.find_as(cast<Function>(getValPtr())); 238 if (I != ACT->AssumptionCaches.end()) 239 ACT->AssumptionCaches.erase(I); 240 // 'this' now dangles! 241 } 242 243 AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) { 244 // We probe the function map twice to try and avoid creating a value handle 245 // around the function in common cases. This makes insertion a bit slower, 246 // but if we have to insert we're going to scan the whole function so that 247 // shouldn't matter. 248 auto I = AssumptionCaches.find_as(&F); 249 if (I != AssumptionCaches.end()) 250 return *I->second; 251 252 // Ok, build a new cache by scanning the function, insert it and the value 253 // handle into our map, and return the newly populated cache. 254 auto IP = AssumptionCaches.insert(std::make_pair( 255 FunctionCallbackVH(&F, this), llvm::make_unique<AssumptionCache>(F))); 256 assert(IP.second && "Scanning function already in the map?"); 257 return *IP.first->second; 258 } 259 260 AssumptionCache *AssumptionCacheTracker::lookupAssumptionCache(Function &F) { 261 auto I = AssumptionCaches.find_as(&F); 262 if (I != AssumptionCaches.end()) 263 return I->second.get(); 264 return nullptr; 265 } 266 267 void AssumptionCacheTracker::verifyAnalysis() const { 268 // FIXME: In the long term the verifier should not be controllable with a 269 // flag. We should either fix all passes to correctly update the assumption 270 // cache and enable the verifier unconditionally or somehow arrange for the 271 // assumption list to be updated automatically by passes. 272 if (!VerifyAssumptionCache) 273 return; 274 275 SmallPtrSet<const CallInst *, 4> AssumptionSet; 276 for (const auto &I : AssumptionCaches) { 277 for (auto &VH : I.second->assumptions()) 278 if (VH) 279 AssumptionSet.insert(cast<CallInst>(VH)); 280 281 for (const BasicBlock &B : cast<Function>(*I.first)) 282 for (const Instruction &II : B) 283 if (match(&II, m_Intrinsic<Intrinsic::assume>()) && 284 !AssumptionSet.count(cast<CallInst>(&II))) 285 report_fatal_error("Assumption in scanned function not in cache"); 286 } 287 } 288 289 AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) { 290 initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry()); 291 } 292 293 AssumptionCacheTracker::~AssumptionCacheTracker() = default; 294 295 char AssumptionCacheTracker::ID = 0; 296 297 INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker", 298 "Assumption Cache Tracker", false, true) 299