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