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 and 10 // @llvm.experimental.guard intrinsics in 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_ICmp(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 } 132 } 133 134 if (TTI) { 135 const Value *Ptr; 136 unsigned AS; 137 std::tie(Ptr, AS) = TTI->getPredicatedAddrSpace(Cond); 138 if (Ptr) 139 AddAffected(const_cast<Value *>(Ptr->stripInBoundsOffsets())); 140 } 141 } 142 143 void AssumptionCache::updateAffectedValues(CondGuardInst *CI) { 144 SmallVector<AssumptionCache::ResultElem, 16> Affected; 145 findAffectedValues(CI, TTI, Affected); 146 147 for (auto &AV : Affected) { 148 auto &AVV = getOrInsertAffectedValues(AV.Assume); 149 if (llvm::none_of(AVV, [&](ResultElem &Elem) { 150 return Elem.Assume == CI && Elem.Index == AV.Index; 151 })) 152 AVV.push_back({CI, AV.Index}); 153 } 154 } 155 156 void AssumptionCache::unregisterAssumption(CondGuardInst *CI) { 157 SmallVector<AssumptionCache::ResultElem, 16> Affected; 158 findAffectedValues(CI, TTI, Affected); 159 160 for (auto &AV : Affected) { 161 auto AVI = AffectedValues.find_as(AV.Assume); 162 if (AVI == AffectedValues.end()) 163 continue; 164 bool Found = false; 165 bool HasNonnull = false; 166 for (ResultElem &Elem : AVI->second) { 167 if (Elem.Assume == CI) { 168 Found = true; 169 Elem.Assume = nullptr; 170 } 171 HasNonnull |= !!Elem.Assume; 172 if (HasNonnull && Found) 173 break; 174 } 175 assert(Found && "already unregistered or incorrect cache state"); 176 if (!HasNonnull) 177 AffectedValues.erase(AVI); 178 } 179 180 erase_value(AssumeHandles, CI); 181 } 182 183 void AssumptionCache::AffectedValueCallbackVH::deleted() { 184 AC->AffectedValues.erase(getValPtr()); 185 // 'this' now dangles! 186 } 187 188 void AssumptionCache::transferAffectedValuesInCache(Value *OV, Value *NV) { 189 auto &NAVV = getOrInsertAffectedValues(NV); 190 auto AVI = AffectedValues.find(OV); 191 if (AVI == AffectedValues.end()) 192 return; 193 194 for (auto &A : AVI->second) 195 if (!llvm::is_contained(NAVV, A)) 196 NAVV.push_back(A); 197 AffectedValues.erase(OV); 198 } 199 200 void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) { 201 if (!isa<Instruction>(NV) && !isa<Argument>(NV)) 202 return; 203 204 // Any assumptions that affected this value now affect the new value. 205 206 AC->transferAffectedValuesInCache(getValPtr(), NV); 207 // 'this' now might dangle! If the AffectedValues map was resized to add an 208 // entry for NV then this object might have been destroyed in favor of some 209 // copy in the grown map. 210 } 211 212 void AssumptionCache::scanFunction() { 213 assert(!Scanned && "Tried to scan the function twice!"); 214 assert(AssumeHandles.empty() && "Already have assumes when scanning!"); 215 216 // Go through all instructions in all blocks, add all calls to @llvm.assume 217 // to this cache. 218 for (BasicBlock &B : F) 219 for (Instruction &I : B) 220 if (isa<CondGuardInst>(&I)) 221 AssumeHandles.push_back({&I, ExprResultIdx}); 222 223 // Mark the scan as complete. 224 Scanned = true; 225 226 // Update affected values. 227 for (auto &A : AssumeHandles) 228 updateAffectedValues(cast<CondGuardInst>(A)); 229 } 230 231 void AssumptionCache::registerAssumption(CondGuardInst *CI) { 232 // If we haven't scanned the function yet, just drop this assumption. It will 233 // be found when we scan later. 234 if (!Scanned) 235 return; 236 237 AssumeHandles.push_back({CI, ExprResultIdx}); 238 239 #ifndef NDEBUG 240 assert(CI->getParent() && 241 "Cannot a register CondGuardInst not in a basic block"); 242 assert(&F == CI->getParent()->getParent() && 243 "Cannot a register CondGuardInst not in this function"); 244 245 // We expect the number of assumptions to be small, so in an asserts build 246 // check that we don't accumulate duplicates and that all assumptions point 247 // to the same function. 248 SmallPtrSet<Value *, 16> AssumptionSet; 249 for (auto &VH : AssumeHandles) { 250 if (!VH) 251 continue; 252 253 assert(&F == cast<Instruction>(VH)->getParent()->getParent() && 254 "Cached assumption not inside this function!"); 255 assert(isa<CondGuardInst>(VH) && 256 "Cached something other than CondGuardInst!"); 257 assert(AssumptionSet.insert(VH).second && 258 "Cache contains multiple copies of a call!"); 259 } 260 #endif 261 262 updateAffectedValues(CI); 263 } 264 265 AssumptionCache AssumptionAnalysis::run(Function &F, 266 FunctionAnalysisManager &FAM) { 267 auto &TTI = FAM.getResult<TargetIRAnalysis>(F); 268 return AssumptionCache(F, &TTI); 269 } 270 271 AnalysisKey AssumptionAnalysis::Key; 272 273 PreservedAnalyses AssumptionPrinterPass::run(Function &F, 274 FunctionAnalysisManager &AM) { 275 AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F); 276 277 OS << "Cached assumptions for function: " << F.getName() << "\n"; 278 for (auto &VH : AC.assumptions()) 279 if (VH) 280 OS << " " << *cast<CallInst>(VH)->getArgOperand(0) << "\n"; 281 282 return PreservedAnalyses::all(); 283 } 284 285 void AssumptionCacheTracker::FunctionCallbackVH::deleted() { 286 auto I = ACT->AssumptionCaches.find_as(cast<Function>(getValPtr())); 287 if (I != ACT->AssumptionCaches.end()) 288 ACT->AssumptionCaches.erase(I); 289 // 'this' now dangles! 290 } 291 292 AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) { 293 // We probe the function map twice to try and avoid creating a value handle 294 // around the function in common cases. This makes insertion a bit slower, 295 // but if we have to insert we're going to scan the whole function so that 296 // shouldn't matter. 297 auto I = AssumptionCaches.find_as(&F); 298 if (I != AssumptionCaches.end()) 299 return *I->second; 300 301 auto *TTIWP = getAnalysisIfAvailable<TargetTransformInfoWrapperPass>(); 302 auto *TTI = TTIWP ? &TTIWP->getTTI(F) : nullptr; 303 304 // Ok, build a new cache by scanning the function, insert it and the value 305 // handle into our map, and return the newly populated cache. 306 auto IP = AssumptionCaches.insert(std::make_pair( 307 FunctionCallbackVH(&F, this), std::make_unique<AssumptionCache>(F, TTI))); 308 assert(IP.second && "Scanning function already in the map?"); 309 return *IP.first->second; 310 } 311 312 AssumptionCache *AssumptionCacheTracker::lookupAssumptionCache(Function &F) { 313 auto I = AssumptionCaches.find_as(&F); 314 if (I != AssumptionCaches.end()) 315 return I->second.get(); 316 return nullptr; 317 } 318 319 void AssumptionCacheTracker::verifyAnalysis() const { 320 // FIXME: In the long term the verifier should not be controllable with a 321 // flag. We should either fix all passes to correctly update the assumption 322 // cache and enable the verifier unconditionally or somehow arrange for the 323 // assumption list to be updated automatically by passes. 324 if (!VerifyAssumptionCache) 325 return; 326 327 SmallPtrSet<const CallInst *, 4> AssumptionSet; 328 for (const auto &I : AssumptionCaches) { 329 for (auto &VH : I.second->assumptions()) 330 if (VH) 331 AssumptionSet.insert(cast<CallInst>(VH)); 332 333 for (const BasicBlock &B : cast<Function>(*I.first)) 334 for (const Instruction &II : B) 335 if (match(&II, m_Intrinsic<Intrinsic::assume>()) && 336 !AssumptionSet.count(cast<CallInst>(&II))) 337 report_fatal_error("Assumption in scanned function not in cache"); 338 } 339 } 340 341 AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) { 342 initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry()); 343 } 344 345 AssumptionCacheTracker::~AssumptionCacheTracker() = default; 346 347 char AssumptionCacheTracker::ID = 0; 348 349 INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker", 350 "Assumption Cache Tracker", false, true) 351