xref: /llvm-project/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp (revision 5957e058e41c90503d0fcaf595707f627042c218)
1 //===- AssumeBundleBuilder.cpp - tools to preserve informations -*- 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 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
10 #include "llvm/ADT/MapVector.h"
11 #include "llvm/Analysis/AssumeBundleQueries.h"
12 #include "llvm/Analysis/AssumptionCache.h"
13 #include "llvm/Analysis/ValueTracking.h"
14 #include "llvm/IR/Dominators.h"
15 #include "llvm/IR/Function.h"
16 #include "llvm/IR/InstIterator.h"
17 #include "llvm/IR/IntrinsicInst.h"
18 #include "llvm/IR/Module.h"
19 #include "llvm/Support/CommandLine.h"
20 
21 using namespace llvm;
22 
23 cl::opt<bool> ShouldPreserveAllAttributes(
24     "assume-preserve-all", cl::init(false), cl::Hidden,
25     cl::desc("enable preservation of all attrbitues. even those that are "
26              "unlikely to be usefull"));
27 
28 cl::opt<bool> EnableKnowledgeRetention(
29     "enable-knowledge-retention", cl::init(false), cl::Hidden,
30     cl::desc(
31         "enable preservation of attributes throughout code transformation"));
32 
33 namespace {
34 
35 bool isUsefullToPreserve(Attribute::AttrKind Kind) {
36   switch (Kind) {
37     case Attribute::NonNull:
38     case Attribute::Alignment:
39     case Attribute::Dereferenceable:
40     case Attribute::DereferenceableOrNull:
41     case Attribute::Cold:
42       return true;
43     default:
44       return false;
45   }
46 }
47 
48 /// This class contain all knowledge that have been gather while building an
49 /// llvm.assume and the function to manipulate it.
50 struct AssumeBuilderState {
51   Module *M;
52 
53   using MapKey = std::pair<Value *, Attribute::AttrKind>;
54   SmallMapVector<MapKey, unsigned, 8> AssumedKnowledgeMap;
55   Instruction *InsertBeforeInstruction = nullptr;
56   AssumptionCache* AC = nullptr;
57   DominatorTree* DT = nullptr;
58 
59   AssumeBuilderState(Module *M, Instruction *I = nullptr,
60                      AssumptionCache *AC = nullptr, DominatorTree *DT = nullptr)
61       : M(M), InsertBeforeInstruction(I), AC(AC), DT(DT) {}
62 
63   bool tryToPreserveWithoutAddingAssume(RetainedKnowledge RK) {
64     if (!InsertBeforeInstruction || !AC || !RK.WasOn)
65       return false;
66     bool HasBeenPreserved = false;
67     Use* ToUpdate = nullptr;
68     getKnowledgeForValue(
69         RK.WasOn, {RK.AttrKind}, AC,
70         [&](RetainedKnowledge RKOther, Instruction *Assume,
71             const CallInst::BundleOpInfo *Bundle) {
72           if (!isValidAssumeForContext(Assume, InsertBeforeInstruction, DT))
73             return false;
74           if (RKOther.ArgValue >= RK.ArgValue) {
75             HasBeenPreserved = true;
76             return true;
77           } else if (isValidAssumeForContext(InsertBeforeInstruction, Assume,
78                                              DT)) {
79             HasBeenPreserved = true;
80             IntrinsicInst *Intr = cast<IntrinsicInst>(Assume);
81             ToUpdate = &Intr->op_begin()[Bundle->Begin + ABA_Argument];
82             return true;
83           }
84           return false;
85         });
86     if (ToUpdate)
87       ToUpdate->set(
88           ConstantInt::get(Type::getInt64Ty(M->getContext()), RK.ArgValue));
89     return HasBeenPreserved;
90   }
91 
92   void addKnowledge(RetainedKnowledge RK) {
93     if (tryToPreserveWithoutAddingAssume(RK))
94       return;
95     MapKey Key{RK.WasOn, RK.AttrKind};
96     auto Lookup = AssumedKnowledgeMap.find(Key);
97     if (Lookup == AssumedKnowledgeMap.end()) {
98       AssumedKnowledgeMap[Key] = RK.ArgValue;
99       return;
100     }
101     assert(((Lookup->second == 0 && RK.ArgValue == 0) ||
102             (Lookup->second != 0 && RK.ArgValue != 0)) &&
103            "inconsistent argument value");
104 
105     /// This is only desirable because for all attributes taking an argument
106     /// higher is better.
107     Lookup->second = std::max(Lookup->second, RK.ArgValue);
108   }
109 
110   void addAttribute(Attribute Attr, Value *WasOn) {
111     if (Attr.isTypeAttribute() || Attr.isStringAttribute() ||
112         (!ShouldPreserveAllAttributes &&
113          !isUsefullToPreserve(Attr.getKindAsEnum())))
114       return;
115     unsigned AttrArg = 0;
116     if (Attr.isIntAttribute())
117       AttrArg = Attr.getValueAsInt();
118     addKnowledge({Attr.getKindAsEnum(), AttrArg, WasOn});
119   }
120 
121   void addCall(const CallBase *Call) {
122     auto addAttrList = [&](AttributeList AttrList) {
123       for (unsigned Idx = AttributeList::FirstArgIndex;
124            Idx < AttrList.getNumAttrSets(); Idx++)
125         for (Attribute Attr : AttrList.getAttributes(Idx))
126           addAttribute(Attr, Call->getArgOperand(Idx - 1));
127       for (Attribute Attr : AttrList.getFnAttributes())
128         addAttribute(Attr, nullptr);
129     };
130     addAttrList(Call->getAttributes());
131     if (Function *Fn = Call->getCalledFunction())
132       addAttrList(Fn->getAttributes());
133   }
134 
135   IntrinsicInst *build() {
136     if (AssumedKnowledgeMap.empty())
137       return nullptr;
138     Function *FnAssume = Intrinsic::getDeclaration(M, Intrinsic::assume);
139     LLVMContext &C = M->getContext();
140     SmallVector<OperandBundleDef, 8> OpBundle;
141     for (auto &MapElem : AssumedKnowledgeMap) {
142       SmallVector<Value *, 2> Args;
143       if (MapElem.first.first)
144         Args.push_back(MapElem.first.first);
145 
146       /// This is only valid because for all attribute that currently exist a
147       /// value of 0 is useless. and should not be preserved.
148       if (MapElem.second)
149         Args.push_back(ConstantInt::get(Type::getInt64Ty(M->getContext()),
150                                         MapElem.second));
151       OpBundle.push_back(OperandBundleDefT<Value *>(
152           std::string(Attribute::getNameFromAttrKind(MapElem.first.second)),
153           Args));
154     }
155     return cast<IntrinsicInst>(CallInst::Create(
156         FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}), OpBundle));
157   }
158 
159   void addAccessedPtr(Instruction *MemInst, Value *Pointer, Type *AccType,
160                       MaybeAlign MA) {
161     unsigned DerefSize = MemInst->getModule()
162                              ->getDataLayout()
163                              .getTypeStoreSize(AccType)
164                              .getKnownMinSize();
165     if (DerefSize != 0) {
166       addKnowledge({Attribute::Dereferenceable, DerefSize, Pointer});
167       if (!NullPointerIsDefined(MemInst->getFunction(),
168                                 Pointer->getType()->getPointerAddressSpace()))
169         addKnowledge({Attribute::NonNull, 0u, Pointer});
170     }
171     if (MA.valueOrOne() > 1)
172       addKnowledge(
173           {Attribute::Alignment, unsigned(MA.valueOrOne().value()), Pointer});
174   }
175 
176   void addInstruction(Instruction *I) {
177     if (auto *Call = dyn_cast<CallBase>(I))
178       return addCall(Call);
179     if (auto *Load = dyn_cast<LoadInst>(I))
180       return addAccessedPtr(I, Load->getPointerOperand(), Load->getType(),
181                             Load->getAlign());
182     if (auto *Store = dyn_cast<StoreInst>(I))
183       return addAccessedPtr(I, Store->getPointerOperand(),
184                             Store->getValueOperand()->getType(),
185                             Store->getAlign());
186     // TODO: Add support for the other Instructions.
187     // TODO: Maybe we should look around and merge with other llvm.assume.
188   }
189 };
190 
191 } // namespace
192 
193 IntrinsicInst *llvm::buildAssumeFromInst(Instruction *I) {
194   if (!EnableKnowledgeRetention)
195     return nullptr;
196   AssumeBuilderState Builder(I->getModule());
197   Builder.addInstruction(I);
198   return Builder.build();
199 }
200 
201 void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC, DominatorTree* DT) {
202   if (!EnableKnowledgeRetention)
203     return;
204   AssumeBuilderState Builder(I->getModule(), I, AC, DT);
205   Builder.addInstruction(I);
206   if (IntrinsicInst *Intr = Builder.build()) {
207     Intr->insertBefore(I);
208     if (AC)
209       AC->registerAssumption(Intr);
210   }
211 }
212 
213 PreservedAnalyses AssumeBuilderPass::run(Function &F,
214                                          FunctionAnalysisManager &AM) {
215   AssumptionCache* AC = AM.getCachedResult<AssumptionAnalysis>(F);
216   DominatorTree* DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
217   for (Instruction &I : instructions(F))
218     salvageKnowledge(&I, AC, DT);
219   return PreservedAnalyses::all();
220 }
221