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/AssumeBundleQueries.h"
15#include "llvm/Analysis/AssumptionCache.h"
16#include "llvm/ADT/STLExtras.h"
17#include "llvm/ADT/SmallPtrSet.h"
18#include "llvm/ADT/SmallVector.h"
19#include "llvm/IR/BasicBlock.h"
20#include "llvm/IR/Function.h"
21#include "llvm/IR/InstrTypes.h"
22#include "llvm/IR/Instruction.h"
23#include "llvm/IR/Instructions.h"
24#include "llvm/IR/Intrinsics.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 <algorithm>
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(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(V);
50  if (AVI != AffectedValues.end())
51    return AVI->second;
52
53  auto AVIP = AffectedValues.insert(
54      {AffectedValueCallbackVH(V, this), SmallVector<ResultElem, 1>()});
55  return AVIP.first->second;
56}
57
58static void
59findAffectedValues(CallInst *CI,
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 AddAffected = [&Affected](Value *V, unsigned Idx =
65                                               AssumptionCache::ExprResultIdx) {
66    if (isa<Argument>(V)) {
67      Affected.push_back({V, Idx});
68    } else if (auto *I = dyn_cast<Instruction>(V)) {
69      Affected.push_back({I, Idx});
70
71      // Peek through unary operators to find the source of the condition.
72      Value *Op;
73      if (match(I, m_BitCast(m_Value(Op))) ||
74          match(I, m_PtrToInt(m_Value(Op))) || match(I, m_Not(m_Value(Op)))) {
75        if (isa<Instruction>(Op) || isa<Argument>(Op))
76          Affected.push_back({Op, Idx});
77      }
78    }
79  };
80
81  for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++) {
82    if (CI->getOperandBundleAt(Idx).Inputs.size() > ABA_WasOn &&
83        CI->getOperandBundleAt(Idx).getTagName() != IgnoreBundleTag)
84      AddAffected(CI->getOperandBundleAt(Idx).Inputs[ABA_WasOn], Idx);
85  }
86
87  Value *Cond = CI->getArgOperand(0), *A, *B;
88  AddAffected(Cond);
89
90  CmpInst::Predicate Pred;
91  if (match(Cond, m_ICmp(Pred, m_Value(A), m_Value(B)))) {
92    AddAffected(A);
93    AddAffected(B);
94
95    if (Pred == ICmpInst::ICMP_EQ) {
96      // For equality comparisons, we handle the case of bit inversion.
97      auto AddAffectedFromEq = [&AddAffected](Value *V) {
98        Value *A;
99        if (match(V, m_Not(m_Value(A)))) {
100          AddAffected(A);
101          V = A;
102        }
103
104        Value *B;
105        ConstantInt *C;
106        // (A & B) or (A | B) or (A ^ B).
107        if (match(V, m_BitwiseLogic(m_Value(A), m_Value(B)))) {
108          AddAffected(A);
109          AddAffected(B);
110        // (A << C) or (A >>_s C) or (A >>_u C) where C is some constant.
111        } else if (match(V, m_Shift(m_Value(A), m_ConstantInt(C)))) {
112          AddAffected(A);
113        }
114      };
115
116      AddAffectedFromEq(A);
117      AddAffectedFromEq(B);
118    }
119  }
120}
121
122void AssumptionCache::updateAffectedValues(CallInst *CI) {
123  SmallVector<AssumptionCache::ResultElem, 16> Affected;
124  findAffectedValues(CI, Affected);
125
126  for (auto &AV : Affected) {
127    auto &AVV = getOrInsertAffectedValues(AV.Assume);
128    if (std::find_if(AVV.begin(), AVV.end(), [&](ResultElem &Elem) {
129          return Elem.Assume == CI && Elem.Index == AV.Index;
130        }) == AVV.end())
131      AVV.push_back({CI, AV.Index});
132  }
133}
134
135void AssumptionCache::unregisterAssumption(CallInst *CI) {
136  SmallVector<AssumptionCache::ResultElem, 16> Affected;
137  findAffectedValues(CI, Affected);
138
139  for (auto &AV : Affected) {
140    auto AVI = AffectedValues.find_as(AV.Assume);
141    if (AVI == AffectedValues.end())
142      continue;
143    bool Found = false;
144    bool HasNonnull = false;
145    for (ResultElem &Elem : AVI->second) {
146      if (Elem.Assume == CI) {
147        Found = true;
148        Elem.Assume = nullptr;
149      }
150      HasNonnull |= !!Elem.Assume;
151      if (HasNonnull && Found)
152        break;
153    }
154    assert(Found && "already unregistered or incorrect cache state");
155    if (!HasNonnull)
156      AffectedValues.erase(AVI);
157  }
158
159  AssumeHandles.erase(
160      remove_if(AssumeHandles, [CI](ResultElem &RE) { return CI == RE; }),
161      AssumeHandles.end());
162}
163
164void AssumptionCache::AffectedValueCallbackVH::deleted() {
165  auto AVI = AC->AffectedValues.find(getValPtr());
166  if (AVI != AC->AffectedValues.end())
167    AC->AffectedValues.erase(AVI);
168  // 'this' now dangles!
169}
170
171void AssumptionCache::transferAffectedValuesInCache(Value *OV, Value *NV) {
172  auto &NAVV = getOrInsertAffectedValues(NV);
173  auto AVI = AffectedValues.find(OV);
174  if (AVI == AffectedValues.end())
175    return;
176
177  for (auto &A : AVI->second)
178    if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end())
179      NAVV.push_back(A);
180  AffectedValues.erase(OV);
181}
182
183void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) {
184  if (!isa<Instruction>(NV) && !isa<Argument>(NV))
185    return;
186
187  // Any assumptions that affected this value now affect the new value.
188
189  AC->transferAffectedValuesInCache(getValPtr(), NV);
190  // 'this' now might dangle! If the AffectedValues map was resized to add an
191  // entry for NV then this object might have been destroyed in favor of some
192  // copy in the grown map.
193}
194
195void AssumptionCache::scanFunction() {
196  assert(!Scanned && "Tried to scan the function twice!");
197  assert(AssumeHandles.empty() && "Already have assumes when scanning!");
198
199  // Go through all instructions in all blocks, add all calls to @llvm.assume
200  // to this cache.
201  for (BasicBlock &B : F)
202    for (Instruction &II : B)
203      if (match(&II, m_Intrinsic<Intrinsic::assume>()))
204        AssumeHandles.push_back({&II, ExprResultIdx});
205
206  // Mark the scan as complete.
207  Scanned = true;
208
209  // Update affected values.
210  for (auto &A : AssumeHandles)
211    updateAffectedValues(cast<CallInst>(A));
212}
213
214void AssumptionCache::registerAssumption(CallInst *CI) {
215  assert(match(CI, m_Intrinsic<Intrinsic::assume>()) &&
216         "Registered call does not call @llvm.assume");
217
218  // If we haven't scanned the function yet, just drop this assumption. It will
219  // be found when we scan later.
220  if (!Scanned)
221    return;
222
223  AssumeHandles.push_back({CI, ExprResultIdx});
224
225#ifndef NDEBUG
226  assert(CI->getParent() &&
227         "Cannot register @llvm.assume call not in a basic block");
228  assert(&F == CI->getParent()->getParent() &&
229         "Cannot register @llvm.assume call not in this function");
230
231  // We expect the number of assumptions to be small, so in an asserts build
232  // check that we don't accumulate duplicates and that all assumptions point
233  // to the same function.
234  SmallPtrSet<Value *, 16> AssumptionSet;
235  for (auto &VH : AssumeHandles) {
236    if (!VH)
237      continue;
238
239    assert(&F == cast<Instruction>(VH)->getParent()->getParent() &&
240           "Cached assumption not inside this function!");
241    assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) &&
242           "Cached something other than a call to @llvm.assume!");
243    assert(AssumptionSet.insert(VH).second &&
244           "Cache contains multiple copies of a call!");
245  }
246#endif
247
248  updateAffectedValues(CI);
249}
250
251AnalysisKey AssumptionAnalysis::Key;
252
253PreservedAnalyses AssumptionPrinterPass::run(Function &F,
254                                             FunctionAnalysisManager &AM) {
255  AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
256
257  OS << "Cached assumptions for function: " << F.getName() << "\n";
258  for (auto &VH : AC.assumptions())
259    if (VH)
260      OS << "  " << *cast<CallInst>(VH)->getArgOperand(0) << "\n";
261
262  return PreservedAnalyses::all();
263}
264
265void AssumptionCacheTracker::FunctionCallbackVH::deleted() {
266  auto I = ACT->AssumptionCaches.find_as(cast<Function>(getValPtr()));
267  if (I != ACT->AssumptionCaches.end())
268    ACT->AssumptionCaches.erase(I);
269  // 'this' now dangles!
270}
271
272AssumptionCache &AssumptionCacheTracker::getAssumptionCache(Function &F) {
273  // We probe the function map twice to try and avoid creating a value handle
274  // around the function in common cases. This makes insertion a bit slower,
275  // but if we have to insert we're going to scan the whole function so that
276  // shouldn't matter.
277  auto I = AssumptionCaches.find_as(&F);
278  if (I != AssumptionCaches.end())
279    return *I->second;
280
281  // Ok, build a new cache by scanning the function, insert it and the value
282  // handle into our map, and return the newly populated cache.
283  auto IP = AssumptionCaches.insert(std::make_pair(
284      FunctionCallbackVH(&F, this), std::make_unique<AssumptionCache>(F)));
285  assert(IP.second && "Scanning function already in the map?");
286  return *IP.first->second;
287}
288
289AssumptionCache *AssumptionCacheTracker::lookupAssumptionCache(Function &F) {
290  auto I = AssumptionCaches.find_as(&F);
291  if (I != AssumptionCaches.end())
292    return I->second.get();
293  return nullptr;
294}
295
296void AssumptionCacheTracker::verifyAnalysis() const {
297  // FIXME: In the long term the verifier should not be controllable with a
298  // flag. We should either fix all passes to correctly update the assumption
299  // cache and enable the verifier unconditionally or somehow arrange for the
300  // assumption list to be updated automatically by passes.
301  if (!VerifyAssumptionCache)
302    return;
303
304  SmallPtrSet<const CallInst *, 4> AssumptionSet;
305  for (const auto &I : AssumptionCaches) {
306    for (auto &VH : I.second->assumptions())
307      if (VH)
308        AssumptionSet.insert(cast<CallInst>(VH));
309
310    for (const BasicBlock &B : cast<Function>(*I.first))
311      for (const Instruction &II : B)
312        if (match(&II, m_Intrinsic<Intrinsic::assume>()) &&
313            !AssumptionSet.count(cast<CallInst>(&II)))
314          report_fatal_error("Assumption in scanned function not in cache");
315  }
316}
317
318AssumptionCacheTracker::AssumptionCacheTracker() : ImmutablePass(ID) {
319  initializeAssumptionCacheTrackerPass(*PassRegistry::getPassRegistry());
320}
321
322AssumptionCacheTracker::~AssumptionCacheTracker() = default;
323
324char AssumptionCacheTracker::ID = 0;
325
326INITIALIZE_PASS(AssumptionCacheTracker, "assumption-cache-tracker",
327                "Assumption Cache Tracker", false, true)
328