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 /// Deterministically compare OperandBundleDef.
34 /// The ordering is:
35 /// - by the attribute's name aka operand bundle tag, (doesn't change)
36 /// - then by the numeric Value of the argument, (doesn't change)
37 /// - lastly by the Name of the current Value it WasOn. (may change)
38 /// This order is deterministic and allows looking for the right kind of
39 /// attribute with binary search. However finding the right WasOn needs to be
40 /// done via linear search because values can get replaced.
41 bool isLowerOpBundle(const OperandBundleDef &LHS, const OperandBundleDef &RHS) {
42   auto getTuple = [](const OperandBundleDef &Op) {
43     return std::make_tuple(
44         Op.getTag(),
45         Op.input_size() <= ABA_Argument
46             ? 0
47             : cast<ConstantInt>(*(Op.input_begin() + ABA_Argument))
48                   ->getZExtValue(),
49         Op.input_size() <= ABA_WasOn
50             ? StringRef("")
51             : (*(Op.input_begin() + ABA_WasOn))->getName());
52   };
53   return getTuple(LHS) < getTuple(RHS);
54 }
55 
56 bool isUsefullToPreserve(Attribute::AttrKind Kind) {
57   switch (Kind) {
58     case Attribute::NonNull:
59     case Attribute::Alignment:
60     case Attribute::Dereferenceable:
61     case Attribute::DereferenceableOrNull:
62     case Attribute::Cold:
63       return true;
64     default:
65       return false;
66   }
67 }
68 
69 /// This class contain all knowledge that have been gather while building an
70 /// llvm.assume and the function to manipulate it.
71 struct AssumeBuilderState {
72   Module *M;
73 
74   using MapKey = std::pair<Value *, Attribute::AttrKind>;
75   SmallDenseMap<MapKey, unsigned, 8> AssumedKnowledgeMap;
76   Instruction *InsertBeforeInstruction = nullptr;
77 
78   AssumeBuilderState(Module *M) : M(M) {}
79 
80   void addKnowledge(RetainedKnowledge RK) {
81     MapKey Key{RK.WasOn, RK.AttrKind};
82     auto Lookup = AssumedKnowledgeMap.find(Key);
83     if (Lookup == AssumedKnowledgeMap.end()) {
84       AssumedKnowledgeMap[Key] = RK.ArgValue;
85       return;
86     }
87     assert(((Lookup->second == 0 && RK.ArgValue == 0) ||
88             (Lookup->second != 0 && RK.ArgValue != 0)) &&
89            "inconsistent argument value");
90 
91     /// This is only desirable because for all attributes taking an argument
92     /// higher is better.
93     Lookup->second = std::max(Lookup->second, RK.ArgValue);
94   }
95 
96   void addAttribute(Attribute Attr, Value *WasOn) {
97     if (Attr.isTypeAttribute() || Attr.isStringAttribute() ||
98         (!ShouldPreserveAllAttributes &&
99          !isUsefullToPreserve(Attr.getKindAsEnum())))
100       return;
101     unsigned AttrArg = 0;
102     if (Attr.isIntAttribute())
103       AttrArg = Attr.getValueAsInt();
104     addKnowledge({Attr.getKindAsEnum(), AttrArg, WasOn});
105   }
106 
107   void addCall(const CallBase *Call) {
108     auto addAttrList = [&](AttributeList AttrList) {
109       for (unsigned Idx = AttributeList::FirstArgIndex;
110            Idx < AttrList.getNumAttrSets(); Idx++)
111         for (Attribute Attr : AttrList.getAttributes(Idx))
112           addAttribute(Attr, Call->getArgOperand(Idx - 1));
113       for (Attribute Attr : AttrList.getFnAttributes())
114         addAttribute(Attr, nullptr);
115     };
116     addAttrList(Call->getAttributes());
117     if (Function *Fn = Call->getCalledFunction())
118       addAttrList(Fn->getAttributes());
119   }
120 
121   IntrinsicInst *build() {
122     if (AssumedKnowledgeMap.empty())
123       return nullptr;
124     Function *FnAssume = Intrinsic::getDeclaration(M, Intrinsic::assume);
125     LLVMContext &C = M->getContext();
126     SmallVector<OperandBundleDef, 8> OpBundle;
127     for (auto &MapElem : AssumedKnowledgeMap) {
128       SmallVector<Value *, 2> Args;
129       if (MapElem.first.first)
130         Args.push_back(MapElem.first.first);
131 
132       /// This is only valid because for all attribute that currently exist a
133       /// value of 0 is useless. and should not be preserved.
134       if (MapElem.second)
135         Args.push_back(ConstantInt::get(Type::getInt64Ty(M->getContext()),
136                                         MapElem.second));
137       OpBundle.push_back(OperandBundleDefT<Value *>(
138           std::string(Attribute::getNameFromAttrKind(MapElem.first.second)),
139           Args));
140     }
141     llvm::sort(OpBundle, isLowerOpBundle);
142     return cast<IntrinsicInst>(CallInst::Create(
143         FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}), OpBundle));
144   }
145 
146   void addAccessedPtr(Instruction *MemInst, Value *Pointer, Type *AccType,
147                       MaybeAlign MA) {
148     unsigned DerefSize = MemInst->getModule()
149                              ->getDataLayout()
150                              .getTypeStoreSize(AccType)
151                              .getKnownMinSize();
152     if (DerefSize != 0) {
153       addKnowledge({Attribute::Dereferenceable, DerefSize, Pointer});
154       if (!NullPointerIsDefined(MemInst->getFunction(),
155                                 Pointer->getType()->getPointerAddressSpace()))
156         addKnowledge({Attribute::NonNull, 0u, Pointer});
157     }
158     if (MA.valueOrOne() > 1)
159       addKnowledge(
160           {Attribute::Alignment, unsigned(MA.valueOrOne().value()), Pointer});
161   }
162 
163   void addInstruction(Instruction *I) {
164     if (auto *Call = dyn_cast<CallBase>(I))
165       return addCall(Call);
166     if (auto *Load = dyn_cast<LoadInst>(I))
167       return addAccessedPtr(I, Load->getPointerOperand(), Load->getType(),
168                             Load->getAlign());
169     if (auto *Store = dyn_cast<StoreInst>(I))
170       return addAccessedPtr(I, Store->getPointerOperand(),
171                             Store->getValueOperand()->getType(),
172                             Store->getAlign());
173     // TODO: Add support for the other Instructions.
174     // TODO: Maybe we should look around and merge with other llvm.assume.
175   }
176 };
177 
178 } // namespace
179 
180 IntrinsicInst *llvm::buildAssumeFromInst(Instruction *I) {
181   if (!EnableKnowledgeRetention)
182     return nullptr;
183   AssumeBuilderState Builder(I->getModule());
184   Builder.addInstruction(I);
185   return Builder.build();
186 }
187 
188 void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC) {
189   if (!EnableKnowledgeRetention)
190     return;
191   AssumeBuilderState Builder(I->getModule());
192   Builder.InsertBeforeInstruction = I;
193   Builder.addInstruction(I);
194   if (IntrinsicInst *Intr = Builder.build()) {
195     Intr->insertBefore(I);
196     if (AC)
197       AC->registerAssumption(Intr);
198   }
199 }
200 
201 PreservedAnalyses AssumeBuilderPass::run(Function &F,
202                                          FunctionAnalysisManager &AM) {
203   for (Instruction &I : instructions(F))
204     if (Instruction *Assume = buildAssumeFromInst(&I))
205       Assume->insertBefore(&I);
206   return PreservedAnalyses::all();
207 }
208