xref: /llvm-project/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp (revision 813f438baaa9638529023b2218875e01ea037735)
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/Analysis/AssumeBundleQueries.h"
11 #include "llvm/Analysis/AssumptionCache.h"
12 #include "llvm/ADT/DenseSet.h"
13 #include "llvm/IR/Function.h"
14 #include "llvm/IR/InstIterator.h"
15 #include "llvm/IR/IntrinsicInst.h"
16 #include "llvm/IR/Module.h"
17 #include "llvm/Support/CommandLine.h"
18 
19 using namespace llvm;
20 
21 cl::opt<bool> ShouldPreserveAllAttributes(
22     "assume-preserve-all", cl::init(false), cl::Hidden,
23     cl::desc("enable preservation of all attrbitues. even those that are "
24              "unlikely to be usefull"));
25 
26 cl::opt<bool> EnableKnowledgeRetention(
27     "enable-knowledge-retention", cl::init(false), cl::Hidden,
28     cl::desc(
29         "enable preservation of attributes throughout code transformation"));
30 
31 namespace {
32 
33 struct AssumedKnowledge {
34   const char *Name;
35   Value *Argument;
36   enum {
37     None,
38     Empty,
39     Tombstone,
40   };
41   /// Contain the argument and a flag if needed.
42   llvm::PointerIntPair<Value *, 2> WasOn;
43 };
44 
45 } // namespace
46 
47 namespace llvm {
48 
49 template <> struct DenseMapInfo<AssumedKnowledge> {
50   static AssumedKnowledge getEmptyKey() {
51     return {nullptr, nullptr, {nullptr, AssumedKnowledge::Empty}};
52   }
53   static AssumedKnowledge getTombstoneKey() {
54     return {nullptr, nullptr, {nullptr, AssumedKnowledge::Tombstone}};
55   }
56   static unsigned getHashValue(const AssumedKnowledge &AK) {
57     return hash_combine(AK.Name, AK.Argument, AK.WasOn.getPointer());
58   }
59   static bool isEqual(const AssumedKnowledge &LHS,
60                       const AssumedKnowledge &RHS) {
61     return LHS.WasOn == RHS.WasOn && LHS.Name == RHS.Name &&
62            LHS.Argument == RHS.Argument;
63   }
64 };
65 
66 } // namespace llvm
67 
68 namespace {
69 
70 /// Deterministically compare OperandBundleDef.
71 /// The ordering is:
72 /// - by the attribute's name aka operand bundle tag, (doesn't change)
73 /// - then by the numeric Value of the argument, (doesn't change)
74 /// - lastly by the Name of the current Value it WasOn. (may change)
75 /// This order is deterministic and allows looking for the right kind of
76 /// attribute with binary search. However finding the right WasOn needs to be
77 /// done via linear search because values can get replaced.
78 bool isLowerOpBundle(const OperandBundleDef &LHS, const OperandBundleDef &RHS) {
79   auto getTuple = [](const OperandBundleDef &Op) {
80     return std::make_tuple(
81         Op.getTag(),
82         Op.input_size() <= ABA_Argument
83             ? 0
84             : cast<ConstantInt>(*(Op.input_begin() + ABA_Argument))
85                   ->getZExtValue(),
86         Op.input_size() <= ABA_WasOn
87             ? StringRef("")
88             : (*(Op.input_begin() + ABA_WasOn))->getName());
89   };
90   return getTuple(LHS) < getTuple(RHS);
91 }
92 
93 bool isUsefullToPreserve(Attribute::AttrKind Kind) {
94   switch (Kind) {
95     case Attribute::NonNull:
96     case Attribute::Alignment:
97     case Attribute::Dereferenceable:
98     case Attribute::DereferenceableOrNull:
99     case Attribute::Cold:
100       return true;
101     default:
102       return false;
103   }
104 }
105 
106 /// This class contain all knowledge that have been gather while building an
107 /// llvm.assume and the function to manipulate it.
108 struct AssumeBuilderState {
109   Module *M;
110 
111   SmallDenseSet<AssumedKnowledge, 8> AssumedKnowledgeSet;
112 
113   AssumeBuilderState(Module *M) : M(M) {}
114 
115   void addAttribute(Attribute Attr, Value *WasOn) {
116     if (!ShouldPreserveAllAttributes &&
117         (Attr.isTypeAttribute() || Attr.isStringAttribute() ||
118          !isUsefullToPreserve(Attr.getKindAsEnum())))
119       return;
120     StringRef Name;
121     Value *AttrArg = nullptr;
122     if (Attr.isStringAttribute())
123       Name = Attr.getKindAsString();
124     else
125       Name = Attribute::getNameFromAttrKind(Attr.getKindAsEnum());
126     if (Attr.isIntAttribute())
127       AttrArg = ConstantInt::get(Type::getInt64Ty(M->getContext()),
128                                  Attr.getValueAsInt());
129     AssumedKnowledgeSet.insert(
130         {Name.data(), AttrArg, {WasOn, AssumedKnowledge::None}});
131   }
132 
133   void addCall(const CallBase *Call) {
134     auto addAttrList = [&](AttributeList AttrList) {
135       for (unsigned Idx = AttributeList::FirstArgIndex;
136            Idx < AttrList.getNumAttrSets(); Idx++)
137         for (Attribute Attr : AttrList.getAttributes(Idx))
138           addAttribute(Attr, Call->getArgOperand(Idx - 1));
139       for (Attribute Attr : AttrList.getFnAttributes())
140         addAttribute(Attr, nullptr);
141     };
142     addAttrList(Call->getAttributes());
143     if (Function *Fn = Call->getCalledFunction())
144       addAttrList(Fn->getAttributes());
145   }
146 
147   IntrinsicInst *build() {
148     if (AssumedKnowledgeSet.empty())
149       return nullptr;
150     Function *FnAssume = Intrinsic::getDeclaration(M, Intrinsic::assume);
151     LLVMContext &C = M->getContext();
152     SmallVector<OperandBundleDef, 8> OpBundle;
153     for (const AssumedKnowledge &Elem : AssumedKnowledgeSet) {
154       SmallVector<Value *, 2> Args;
155       assert(Attribute::getAttrKindFromName(Elem.Name) ==
156                  Attribute::AttrKind::None ||
157              static_cast<bool>(Elem.Argument) ==
158                  Attribute::doesAttrKindHaveArgument(
159                      Attribute::getAttrKindFromName(Elem.Name)));
160       if (Elem.WasOn.getPointer())
161         Args.push_back(Elem.WasOn.getPointer());
162       if (Elem.Argument)
163         Args.push_back(Elem.Argument);
164       OpBundle.push_back(OperandBundleDefT<Value *>(Elem.Name, Args));
165     }
166     llvm::sort(OpBundle, isLowerOpBundle);
167     return cast<IntrinsicInst>(CallInst::Create(
168         FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}), OpBundle));
169   }
170 
171   void addAttr(Attribute::AttrKind Kind, Value *Val, unsigned Argument = 0) {
172     AssumedKnowledge AK;
173     AK.Name = Attribute::getNameFromAttrKind(Kind).data();
174     AK.WasOn.setPointer(Val);
175     if (Attribute::doesAttrKindHaveArgument(Kind)) {
176       AK.Argument =
177           ConstantInt::get(Type::getInt64Ty(M->getContext()), Argument);
178     } else {
179       AK.Argument = nullptr;
180       assert(Argument == 0 && "there should be no argument");
181     }
182     AssumedKnowledgeSet.insert(AK);
183   };
184 
185   void addAccessedPtr(Instruction *MemInst, Value *Pointer, Type *AccType,
186                       MaybeAlign MA) {
187     uint64_t DerefSize = MemInst->getModule()
188                              ->getDataLayout()
189                              .getTypeStoreSize(AccType)
190                              .getKnownMinSize();
191     if (DerefSize != 0) {
192       addAttr(Attribute::Dereferenceable, Pointer, DerefSize);
193       if (!NullPointerIsDefined(MemInst->getFunction(),
194                                 Pointer->getType()->getPointerAddressSpace()))
195         addAttr(Attribute::NonNull, Pointer);
196     }
197     if (MA.valueOrOne() > 1)
198       addAttr(Attribute::Alignment, Pointer, MA.valueOrOne().value());
199   }
200 
201   void addInstruction(Instruction *I) {
202     if (auto *Call = dyn_cast<CallBase>(I))
203       return addCall(Call);
204     if (auto *Load = dyn_cast<LoadInst>(I))
205       return addAccessedPtr(I, Load->getPointerOperand(), Load->getType(),
206                             Load->getAlign());
207     if (auto *Store = dyn_cast<StoreInst>(I))
208       return addAccessedPtr(I, Store->getPointerOperand(),
209                             Store->getValueOperand()->getType(),
210                             Store->getAlign());
211     // TODO: Add support for the other Instructions.
212     // TODO: Maybe we should look around and merge with other llvm.assume.
213   }
214 };
215 
216 } // namespace
217 
218 IntrinsicInst *llvm::buildAssumeFromInst(Instruction *I) {
219   if (!EnableKnowledgeRetention)
220     return nullptr;
221   AssumeBuilderState Builder(I->getModule());
222   Builder.addInstruction(I);
223   return Builder.build();
224 }
225 
226 void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC) {
227   if (IntrinsicInst *Intr = buildAssumeFromInst(I)) {
228     Intr->insertBefore(I);
229     if (AC)
230       AC->registerAssumption(Intr);
231   }
232 }
233 
234 PreservedAnalyses AssumeBuilderPass::run(Function &F,
235                                          FunctionAnalysisManager &AM) {
236   for (Instruction &I : instructions(F))
237     if (Instruction *Assume = buildAssumeFromInst(&I))
238       Assume->insertBefore(&I);
239   return PreservedAnalyses::all();
240 }
241