1//===- llvm/Analysis/AssumptionCache.h - Track @llvm.assume -----*- C++ -*-===//
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 (allowing assumptions within any function to be
11// found cheaply by other parts of the optimizer).
12//
13//===----------------------------------------------------------------------===//
14
15#ifndef LLVM_ANALYSIS_ASSUMPTIONCACHE_H
16#define LLVM_ANALYSIS_ASSUMPTIONCACHE_H
17
18#include "llvm/ADT/ArrayRef.h"
19#include "llvm/ADT/DenseMap.h"
20#include "llvm/ADT/DenseMapInfo.h"
21#include "llvm/ADT/SmallVector.h"
22#include "llvm/IR/PassManager.h"
23#include "llvm/IR/ValueHandle.h"
24#include "llvm/Pass.h"
25#include <memory>
26
27namespace llvm {
28
29class AssumeInst;
30class Function;
31class raw_ostream;
32class TargetTransformInfo;
33class Value;
34
35/// A cache of \@llvm.assume calls within a function.
36///
37/// This cache provides fast lookup of assumptions within a function by caching
38/// them and amortizing the cost of scanning for them across all queries. Passes
39/// that create new assumptions are required to call registerAssumption() to
40/// register any new \@llvm.assume calls that they create. Deletions of
41/// \@llvm.assume calls do not require special handling.
42class AssumptionCache {
43public:
44  /// Value of ResultElem::Index indicating that the argument to the call of the
45  /// llvm.assume.
46  enum : unsigned { ExprResultIdx = std::numeric_limits<unsigned>::max() };
47
48  struct ResultElem {
49    WeakVH Assume;
50
51    /// contains either ExprResultIdx or the index of the operand bundle
52    /// containing the knowledge.
53    unsigned Index;
54    operator Value *() const { return Assume; }
55  };
56
57private:
58  /// The function for which this cache is handling assumptions.
59  ///
60  /// We track this to lazily populate our assumptions.
61  Function &F;
62
63  TargetTransformInfo *TTI;
64
65  /// Vector of weak value handles to calls of the \@llvm.assume
66  /// intrinsic.
67  SmallVector<ResultElem, 4> AssumeHandles;
68
69  class AffectedValueCallbackVH final : public CallbackVH {
70    AssumptionCache *AC;
71
72    void deleted() override;
73    void allUsesReplacedWith(Value *) override;
74
75  public:
76    using DMI = DenseMapInfo<Value *>;
77
78    AffectedValueCallbackVH(Value *V, AssumptionCache *AC = nullptr)
79        : CallbackVH(V), AC(AC) {}
80  };
81
82  friend AffectedValueCallbackVH;
83
84  /// A map of values about which an assumption might be providing
85  /// information to the relevant set of assumptions.
86  using AffectedValuesMap =
87      DenseMap<AffectedValueCallbackVH, SmallVector<ResultElem, 1>,
88               AffectedValueCallbackVH::DMI>;
89  AffectedValuesMap AffectedValues;
90
91  /// Get the vector of assumptions which affect a value from the cache.
92  SmallVector<ResultElem, 1> &getOrInsertAffectedValues(Value *V);
93
94  /// Move affected values in the cache for OV to be affected values for NV.
95  void transferAffectedValuesInCache(Value *OV, Value *NV);
96
97  /// Flag tracking whether we have scanned the function yet.
98  ///
99  /// We want to be as lazy about this as possible, and so we scan the function
100  /// at the last moment.
101  bool Scanned = false;
102
103  /// Scan the function for assumptions and add them to the cache.
104  void scanFunction();
105
106public:
107  /// Construct an AssumptionCache from a function by scanning all of
108  /// its instructions.
109  AssumptionCache(Function &F, TargetTransformInfo *TTI = nullptr)
110      : F(F), TTI(TTI) {}
111
112  /// This cache is designed to be self-updating and so it should never be
113  /// invalidated.
114  bool invalidate(Function &, const PreservedAnalyses &,
115                  FunctionAnalysisManager::Invalidator &) {
116    return false;
117  }
118
119  /// Add an \@llvm.assume intrinsic to this function's cache.
120  ///
121  /// The call passed in must be an instruction within this function and must
122  /// not already be in the cache.
123  void registerAssumption(AssumeInst *CI);
124
125  /// Remove an \@llvm.assume intrinsic from this function's cache if it has
126  /// been added to the cache earlier.
127  void unregisterAssumption(AssumeInst *CI);
128
129  /// Update the cache of values being affected by this assumption (i.e.
130  /// the values about which this assumption provides information).
131  void updateAffectedValues(AssumeInst *CI);
132
133  /// Clear the cache of \@llvm.assume intrinsics for a function.
134  ///
135  /// It will be re-scanned the next time it is requested.
136  void clear() {
137    AssumeHandles.clear();
138    AffectedValues.clear();
139    Scanned = false;
140  }
141
142  /// Access the list of assumption handles currently tracked for this
143  /// function.
144  ///
145  /// Note that these produce weak handles that may be null. The caller must
146  /// handle that case.
147  /// FIXME: We should replace this with pointee_iterator<filter_iterator<...>>
148  /// when we can write that to filter out the null values. Then caller code
149  /// will become simpler.
150  MutableArrayRef<ResultElem> assumptions() {
151    if (!Scanned)
152      scanFunction();
153    return AssumeHandles;
154  }
155
156  /// Access the list of assumptions which affect this value.
157  MutableArrayRef<ResultElem> assumptionsFor(const Value *V) {
158    if (!Scanned)
159      scanFunction();
160
161    auto AVI = AffectedValues.find_as(const_cast<Value *>(V));
162    if (AVI == AffectedValues.end())
163      return MutableArrayRef<ResultElem>();
164
165    return AVI->second;
166  }
167};
168
169/// A function analysis which provides an \c AssumptionCache.
170///
171/// This analysis is intended for use with the new pass manager and will vend
172/// assumption caches for a given function.
173class AssumptionAnalysis : public AnalysisInfoMixin<AssumptionAnalysis> {
174  friend AnalysisInfoMixin<AssumptionAnalysis>;
175
176  static AnalysisKey Key;
177
178public:
179  using Result = AssumptionCache;
180
181  AssumptionCache run(Function &F, FunctionAnalysisManager &);
182};
183
184/// Printer pass for the \c AssumptionAnalysis results.
185class AssumptionPrinterPass : public PassInfoMixin<AssumptionPrinterPass> {
186  raw_ostream &OS;
187
188public:
189  explicit AssumptionPrinterPass(raw_ostream &OS) : OS(OS) {}
190
191  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
192
193  static bool isRequired() { return true; }
194};
195
196/// An immutable pass that tracks lazily created \c AssumptionCache
197/// objects.
198///
199/// This is essentially a workaround for the legacy pass manager's weaknesses
200/// which associates each assumption cache with Function and clears it if the
201/// function is deleted. The nature of the AssumptionCache is that it is not
202/// invalidated by any changes to the function body and so this is sufficient
203/// to be conservatively correct.
204class AssumptionCacheTracker : public ImmutablePass {
205  /// A callback value handle applied to function objects, which we use to
206  /// delete our cache of intrinsics for a function when it is deleted.
207  class FunctionCallbackVH final : public CallbackVH {
208    AssumptionCacheTracker *ACT;
209
210    void deleted() override;
211
212  public:
213    using DMI = DenseMapInfo<Value *>;
214
215    FunctionCallbackVH(Value *V, AssumptionCacheTracker *ACT = nullptr)
216        : CallbackVH(V), ACT(ACT) {}
217  };
218
219  friend FunctionCallbackVH;
220
221  using FunctionCallsMap =
222      DenseMap<FunctionCallbackVH, std::unique_ptr<AssumptionCache>,
223               FunctionCallbackVH::DMI>;
224
225  FunctionCallsMap AssumptionCaches;
226
227public:
228  /// Get the cached assumptions for a function.
229  ///
230  /// If no assumptions are cached, this will scan the function. Otherwise, the
231  /// existing cache will be returned.
232  AssumptionCache &getAssumptionCache(Function &F);
233
234  /// Return the cached assumptions for a function if it has already been
235  /// scanned. Otherwise return nullptr.
236  AssumptionCache *lookupAssumptionCache(Function &F);
237
238  AssumptionCacheTracker();
239  ~AssumptionCacheTracker() override;
240
241  void releaseMemory() override {
242    verifyAnalysis();
243    AssumptionCaches.shrink_and_clear();
244  }
245
246  void verifyAnalysis() const override;
247
248  bool doFinalization(Module &) override {
249    verifyAnalysis();
250    return false;
251  }
252
253  static char ID; // Pass identification, replacement for typeid
254};
255
256template<> struct simplify_type<AssumptionCache::ResultElem> {
257  using SimpleType = Value *;
258
259  static SimpleType getSimplifiedValue(AssumptionCache::ResultElem &Val) {
260    return Val;
261  }
262};
263template<> struct simplify_type<const AssumptionCache::ResultElem> {
264  using SimpleType = /*const*/ Value *;
265
266  static SimpleType getSimplifiedValue(const AssumptionCache::ResultElem &Val) {
267    return Val;
268  }
269};
270
271} // end namespace llvm
272
273#endif // LLVM_ANALYSIS_ASSUMPTIONCACHE_H
274