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