1 //===-- MoveAutoInit.cpp - move auto-init inst closer to their use site----===//
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 pass moves instruction maked as auto-init closer to the basic block that
10 // use it, eventually removing it from some control path of the function.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "llvm/Transforms/Utils/MoveAutoInit.h"
15 #include "llvm/ADT/STLExtras.h"
16 #include "llvm/ADT/Statistic.h"
17 #include "llvm/ADT/StringSet.h"
18 #include "llvm/Analysis/MemorySSA.h"
19 #include "llvm/Analysis/MemorySSAUpdater.h"
20 #include "llvm/Analysis/ValueTracking.h"
21 #include "llvm/IR/DebugInfo.h"
22 #include "llvm/IR/Dominators.h"
23 #include "llvm/IR/IRBuilder.h"
24 #include "llvm/IR/Instructions.h"
25 #include "llvm/IR/IntrinsicInst.h"
26 #include "llvm/Support/CommandLine.h"
27 #include "llvm/Transforms/Utils.h"
28 #include "llvm/Transforms/Utils/LoopUtils.h"
29 
30 using namespace llvm;
31 
32 #define DEBUG_TYPE "move-auto-init"
33 
34 STATISTIC(NumMoved, "Number of instructions moved");
35 
36 static cl::opt<unsigned> MoveAutoInitThreshold(
37     "move-auto-init-threshold", cl::Hidden, cl::init(128),
38     cl::desc("Maximum instructions to analyze per moved initialization"));
39 
40 static bool hasAutoInitMetadata(const Instruction &I) {
41   return I.hasMetadata(LLVMContext::MD_annotation) &&
42          any_of(I.getMetadata(LLVMContext::MD_annotation)->operands(),
43                 [](const MDOperand &Op) { return Op.equalsStr("auto-init"); });
44 }
45 
46 static std::optional<MemoryLocation> writeToAlloca(const Instruction &I) {
47   MemoryLocation ML;
48   if (auto *MI = dyn_cast<MemIntrinsic>(&I))
49     ML = MemoryLocation::getForDest(MI);
50   else if (auto *SI = dyn_cast<StoreInst>(&I))
51     ML = MemoryLocation::get(SI);
52   else
53     assert(false && "memory location set");
54 
55   if (isa<AllocaInst>(getUnderlyingObject(ML.Ptr)))
56     return ML;
57   else
58     return {};
59 }
60 
61 /// Finds a BasicBlock in the CFG where instruction `I` can be moved to while
62 /// not changing the Memory SSA ordering and being guarded by at least one
63 /// condition.
64 static BasicBlock *usersDominator(const MemoryLocation &ML, Instruction *I,
65                                   DominatorTree &DT, MemorySSA &MSSA) {
66   BasicBlock *CurrentDominator = nullptr;
67   MemoryUseOrDef &IMA = *MSSA.getMemoryAccess(I);
68   BatchAAResults AA(MSSA.getAA());
69 
70   SmallPtrSet<MemoryAccess *, 8> Visited;
71 
72   auto AsMemoryAccess = [](User *U) { return cast<MemoryAccess>(U); };
73   SmallVector<MemoryAccess *> WorkList(map_range(IMA.users(), AsMemoryAccess));
74 
75   while (!WorkList.empty()) {
76     MemoryAccess *MA = WorkList.pop_back_val();
77     if (!Visited.insert(MA).second)
78       continue;
79 
80     if (Visited.size() > MoveAutoInitThreshold)
81       return nullptr;
82 
83     bool FoundClobberingUser = false;
84     if (auto *M = dyn_cast<MemoryUseOrDef>(MA)) {
85       Instruction *MI = M->getMemoryInst();
86 
87       // If this memory instruction may not clobber `I`, we can skip it.
88       // LifetimeEnd is a valid user, but we do not want it in the user
89       // dominator.
90       if (AA.getModRefInfo(MI, ML) != ModRefInfo::NoModRef &&
91           !MI->isLifetimeStartOrEnd() && MI != I) {
92         FoundClobberingUser = true;
93         CurrentDominator = CurrentDominator
94                                ? DT.findNearestCommonDominator(CurrentDominator,
95                                                                MI->getParent())
96                                : MI->getParent();
97       }
98     }
99     if (!FoundClobberingUser) {
100       auto UsersAsMemoryAccesses = map_range(MA->users(), AsMemoryAccess);
101       append_range(WorkList, UsersAsMemoryAccesses);
102     }
103   }
104   return CurrentDominator;
105 }
106 
107 static bool runMoveAutoInit(Function &F, DominatorTree &DT, MemorySSA &MSSA) {
108   BasicBlock &EntryBB = F.getEntryBlock();
109   SmallVector<std::pair<Instruction *, BasicBlock *>> JobList;
110 
111   //
112   // Compute movable instructions.
113   //
114   for (Instruction &I : EntryBB) {
115     if (!hasAutoInitMetadata(I))
116       continue;
117 
118     std::optional<MemoryLocation> ML = writeToAlloca(I);
119     if (!ML)
120       continue;
121 
122     if (I.isVolatile())
123       continue;
124 
125     BasicBlock *UsersDominator = usersDominator(ML.value(), &I, DT, MSSA);
126     if (!UsersDominator)
127       continue;
128 
129     if (UsersDominator == &EntryBB)
130       continue;
131 
132     // Traverse the CFG to detect cycles `UsersDominator` would be part of.
133     SmallPtrSet<BasicBlock *, 8> TransitiveSuccessors;
134     SmallVector<BasicBlock *> WorkList(successors(UsersDominator));
135     bool HasCycle = false;
136     while (!WorkList.empty()) {
137       BasicBlock *CurrBB = WorkList.pop_back_val();
138       if (CurrBB == UsersDominator)
139         // No early exit because we want to compute the full set of transitive
140         // successors.
141         HasCycle = true;
142       for (BasicBlock *Successor : successors(CurrBB)) {
143         if (!TransitiveSuccessors.insert(Successor).second)
144           continue;
145         WorkList.push_back(Successor);
146       }
147     }
148 
149     // Don't insert if that could create multiple execution of I,
150     // but we can insert it in the non back-edge predecessors, if it exists.
151     if (HasCycle) {
152       BasicBlock *UsersDominatorHead = UsersDominator;
153       while (BasicBlock *UniquePredecessor =
154                  UsersDominatorHead->getUniquePredecessor())
155         UsersDominatorHead = UniquePredecessor;
156 
157       if (UsersDominatorHead == &EntryBB)
158         continue;
159 
160       BasicBlock *DominatingPredecessor = nullptr;
161       for (BasicBlock *Pred : predecessors(UsersDominatorHead)) {
162         // If one of the predecessor of the dominator also transitively is a
163         // successor, moving to the dominator would do the inverse of loop
164         // hoisting, and we don't want that.
165         if (TransitiveSuccessors.count(Pred))
166           continue;
167 
168         DominatingPredecessor =
169             DominatingPredecessor
170                 ? DT.findNearestCommonDominator(DominatingPredecessor, Pred)
171                 : Pred;
172       }
173 
174       if (!DominatingPredecessor || DominatingPredecessor == &EntryBB)
175         continue;
176 
177       UsersDominator = DominatingPredecessor;
178     }
179 
180     // CatchSwitchInst blocks can only have one instruction, so they are not
181     // good candidates for insertion.
182     while (isa<CatchSwitchInst>(UsersDominator->getFirstInsertionPt())) {
183       for (BasicBlock *Pred : predecessors(UsersDominator))
184         UsersDominator = DT.findNearestCommonDominator(UsersDominator, Pred);
185     }
186 
187     // We finally found a place where I can be moved while not introducing extra
188     // execution, and guarded by at least one condition.
189     if (UsersDominator != &EntryBB)
190       JobList.emplace_back(&I, UsersDominator);
191   }
192 
193   //
194   // Perform the actual substitution.
195   //
196   if (JobList.empty())
197     return false;
198 
199   MemorySSAUpdater MSSAU(&MSSA);
200 
201   // Reverse insertion to respect relative order between instructions:
202   // if two instructions are moved from the same BB to the same BB, we insert
203   // the second one in the front, then the first on top of it.
204   for (auto &Job : reverse(JobList)) {
205     Job.first->moveBefore(&*Job.second->getFirstInsertionPt());
206     MSSAU.moveToPlace(MSSA.getMemoryAccess(Job.first), Job.first->getParent(),
207                       MemorySSA::InsertionPlace::Beginning);
208   }
209 
210   if (VerifyMemorySSA)
211     MSSA.verifyMemorySSA();
212 
213   NumMoved += JobList.size();
214 
215   return true;
216 }
217 
218 PreservedAnalyses MoveAutoInitPass::run(Function &F,
219                                         FunctionAnalysisManager &AM) {
220 
221   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
222   auto &MSSA = AM.getResult<MemorySSAAnalysis>(F).getMSSA();
223   if (!runMoveAutoInit(F, DT, MSSA))
224     return PreservedAnalyses::all();
225 
226   PreservedAnalyses PA;
227   PA.preserve<DominatorTreeAnalysis>();
228   PA.preserve<MemorySSAAnalysis>();
229   PA.preserveSet<CFGAnalyses>();
230   return PA;
231 }
232