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