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 
27 namespace llvm {
28 
29 class AssumeInst;
30 class Function;
31 class raw_ostream;
32 class TargetTransformInfo;
33 class 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.
42 class AssumptionCache {
43 public:
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 
57 private:
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 
106 public:
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.
173 class AssumptionAnalysis : public AnalysisInfoMixin<AssumptionAnalysis> {
174   friend AnalysisInfoMixin<AssumptionAnalysis>;
175 
176   static AnalysisKey Key;
177 
178 public:
179   using Result = AssumptionCache;
180 
181   AssumptionCache run(Function &F, FunctionAnalysisManager &);
182 };
183 
184 /// Printer pass for the \c AssumptionAnalysis results.
185 class AssumptionPrinterPass : public PassInfoMixin<AssumptionPrinterPass> {
186   raw_ostream &OS;
187 
188 public:
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.
204 class 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 
227 public:
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 
256 template<> struct simplify_type<AssumptionCache::ResultElem> {
257   using SimpleType = Value *;
258 
259   static SimpleType getSimplifiedValue(AssumptionCache::ResultElem &Val) {
260     return Val;
261   }
262 };
263 template<> 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