1//===- ControlFlowUtils.cpp - Control Flow Utilities -----------------------==//
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// Utilities to manipulate the CFG and restore SSA for the new control flow.
10//
11//===----------------------------------------------------------------------===//
12
13#include "llvm/Transforms/Utils/ControlFlowUtils.h"
14#include "llvm/ADT/SetVector.h"
15#include "llvm/Analysis/DomTreeUpdater.h"
16#include "llvm/Analysis/LoopInfo.h"
17#include "llvm/IR/Constants.h"
18#include "llvm/IR/Instructions.h"
19#include "llvm/IR/ValueHandle.h"
20#include "llvm/Transforms/Utils/Local.h"
21
22#define DEBUG_TYPE "control-flow-hub"
23
24using namespace llvm;
25
26using BBPredicates = DenseMap<BasicBlock *, Instruction *>;
27using EdgeDescriptor = ControlFlowHub::BranchDescriptor;
28
29// Redirects the terminator of the incoming block to the first guard block in
30// the hub. Returns the branch condition from `BB` if it exits.
31// - If only one of Succ0 or Succ1 is not null, the corresponding branch
32// successor is redirected to the FirstGuardBlock.
33// - Else both are not null, and branch is replaced with an unconditional
34// branch to the FirstGuardBlock.
35static Value *redirectToHub(BasicBlock *BB, BasicBlock *Succ0,
36 BasicBlock *Succ1, BasicBlock *FirstGuardBlock) {
37 assert(isa<BranchInst>(BB->getTerminator()) &&
38 "Only support branch terminator.");
39 auto *Branch = cast<BranchInst>(Val: BB->getTerminator());
40 auto *Condition = Branch->isConditional() ? Branch->getCondition() : nullptr;
41
42 assert(Succ0 || Succ1);
43
44 if (Branch->isUnconditional()) {
45 assert(Succ0 == Branch->getSuccessor(0));
46 assert(!Succ1);
47 Branch->setSuccessor(idx: 0, NewSucc: FirstGuardBlock);
48 } else {
49 assert(!Succ1 || Succ1 == Branch->getSuccessor(1));
50 if (Succ0 && !Succ1) {
51 Branch->setSuccessor(idx: 0, NewSucc: FirstGuardBlock);
52 } else if (Succ1 && !Succ0) {
53 Branch->setSuccessor(idx: 1, NewSucc: FirstGuardBlock);
54 } else {
55 Branch->eraseFromParent();
56 BranchInst::Create(IfTrue: FirstGuardBlock, InsertBefore: BB);
57 }
58 }
59
60 return Condition;
61}
62
63// Setup the branch instructions for guard blocks.
64//
65// Each guard block terminates in a conditional branch that transfers
66// control to the corresponding outgoing block or the next guard
67// block. The last guard block has two outgoing blocks as successors.
68static void setupBranchForGuard(ArrayRef<BasicBlock *> GuardBlocks,
69 ArrayRef<BasicBlock *> Outgoing,
70 BBPredicates &GuardPredicates) {
71 assert(Outgoing.size() > 1);
72 assert(GuardBlocks.size() == Outgoing.size() - 1);
73 int I = 0;
74 for (int E = GuardBlocks.size() - 1; I != E; ++I) {
75 BasicBlock *Out = Outgoing[I];
76 BranchInst::Create(IfTrue: Out, IfFalse: GuardBlocks[I + 1], Cond: GuardPredicates[Out],
77 InsertBefore: GuardBlocks[I]);
78 }
79 BasicBlock *Out = Outgoing[I];
80 BranchInst::Create(IfTrue: Out, IfFalse: Outgoing[I + 1], Cond: GuardPredicates[Out],
81 InsertBefore: GuardBlocks[I]);
82}
83
84// Assign an index to each outgoing block. At the corresponding guard
85// block, compute the branch condition by comparing this index.
86static void calcPredicateUsingInteger(ArrayRef<EdgeDescriptor> Branches,
87 ArrayRef<BasicBlock *> Outgoing,
88 ArrayRef<BasicBlock *> GuardBlocks,
89 BBPredicates &GuardPredicates) {
90 LLVMContext &Context = GuardBlocks.front()->getContext();
91 BasicBlock *FirstGuardBlock = GuardBlocks.front();
92 Type *Int32Ty = Type::getInt32Ty(C&: Context);
93
94 auto *Phi = PHINode::Create(Ty: Int32Ty, NumReservedValues: Branches.size(), NameStr: "merged.bb.idx",
95 InsertBefore: FirstGuardBlock);
96
97 for (auto [BB, Succ0, Succ1] : Branches) {
98 Value *Condition = redirectToHub(BB, Succ0, Succ1, FirstGuardBlock);
99 Value *IncomingId = nullptr;
100 if (Succ0 && Succ1 && Succ0 != Succ1) {
101 auto Succ0Iter = find(Range&: Outgoing, Val: Succ0);
102 auto Succ1Iter = find(Range&: Outgoing, Val: Succ1);
103 Value *Id0 =
104 ConstantInt::get(Ty: Int32Ty, V: std::distance(first: Outgoing.begin(), last: Succ0Iter));
105 Value *Id1 =
106 ConstantInt::get(Ty: Int32Ty, V: std::distance(first: Outgoing.begin(), last: Succ1Iter));
107 IncomingId = SelectInst::Create(C: Condition, S1: Id0, S2: Id1, NameStr: "target.bb.idx",
108 InsertBefore: BB->getTerminator()->getIterator());
109 } else {
110 // Get the index of the non-null successor, or when both successors
111 // are the same block, use that block's index directly.
112 auto SuccIter = Succ0 ? find(Range&: Outgoing, Val: Succ0) : find(Range&: Outgoing, Val: Succ1);
113 IncomingId =
114 ConstantInt::get(Ty: Int32Ty, V: std::distance(first: Outgoing.begin(), last: SuccIter));
115 }
116 Phi->addIncoming(V: IncomingId, BB);
117 }
118
119 for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
120 BasicBlock *Out = Outgoing[I];
121 LLVM_DEBUG(dbgs() << "Creating integer guard for " << Out->getName()
122 << "\n");
123 auto *Cmp = ICmpInst::Create(Op: Instruction::ICmp, Pred: ICmpInst::ICMP_EQ, S1: Phi,
124 S2: ConstantInt::get(Ty: Int32Ty, V: I),
125 Name: Out->getName() + ".predicate", InsertBefore: GuardBlocks[I]);
126 GuardPredicates[Out] = Cmp;
127 }
128}
129
130// Determine the branch condition to be used at each guard block from the
131// original boolean values.
132static void calcPredicateUsingBooleans(
133 ArrayRef<EdgeDescriptor> Branches, ArrayRef<BasicBlock *> Outgoing,
134 SmallVectorImpl<BasicBlock *> &GuardBlocks, BBPredicates &GuardPredicates,
135 SmallVectorImpl<WeakVH> &DeletionCandidates) {
136 LLVMContext &Context = GuardBlocks.front()->getContext();
137 auto *BoolTrue = ConstantInt::getTrue(Context);
138 auto *BoolFalse = ConstantInt::getFalse(Context);
139 BasicBlock *FirstGuardBlock = GuardBlocks.front();
140
141 // The predicate for the last outgoing is trivially true, and so we
142 // process only the first N-1 successors.
143 for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
144 BasicBlock *Out = Outgoing[I];
145 LLVM_DEBUG(dbgs() << "Creating boolean guard for " << Out->getName()
146 << "\n");
147
148 auto *Phi =
149 PHINode::Create(Ty: Type::getInt1Ty(C&: Context), NumReservedValues: Branches.size(),
150 NameStr: StringRef("Guard.") + Out->getName(), InsertBefore: FirstGuardBlock);
151 GuardPredicates[Out] = Phi;
152 }
153
154 for (auto [BB, Succ0, Succ1] : Branches) {
155 Value *Condition = redirectToHub(BB, Succ0, Succ1, FirstGuardBlock);
156
157 // Optimization: Consider an incoming block A with both successors
158 // Succ0 and Succ1 in the set of outgoing blocks. The predicates
159 // for Succ0 and Succ1 complement each other. If Succ0 is visited
160 // first in the loop below, control will branch to Succ0 using the
161 // corresponding predicate. But if that branch is not taken, then
162 // control must reach Succ1, which means that the incoming value of
163 // the predicate from `BB` is true for Succ1.
164 bool OneSuccessorDone = false;
165 for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
166 BasicBlock *Out = Outgoing[I];
167 PHINode *Phi = cast<PHINode>(Val: GuardPredicates[Out]);
168 if (Out != Succ0 && Out != Succ1) {
169 Phi->addIncoming(V: BoolFalse, BB);
170 } else if (!Succ0 || !Succ1 || Succ0 == Succ1 || OneSuccessorDone) {
171 // Optimization: When only one successor is an outgoing block,
172 // or both successors are the same block, the incoming predicate
173 // from `BB` is always true.
174 Phi->addIncoming(V: BoolTrue, BB);
175 } else {
176 assert(Succ0 && Succ1);
177 if (Out == Succ0) {
178 Phi->addIncoming(V: Condition, BB);
179 } else {
180 Value *Inverted = invertCondition(Condition);
181 DeletionCandidates.push_back(Elt: Condition);
182 Phi->addIncoming(V: Inverted, BB);
183 }
184 OneSuccessorDone = true;
185 }
186 }
187 }
188}
189
190// Capture the existing control flow as guard predicates, and redirect
191// control flow from \p Incoming block through the \p GuardBlocks to the
192// \p Outgoing blocks.
193//
194// There is one guard predicate for each outgoing block OutBB. The
195// predicate represents whether the hub should transfer control flow
196// to OutBB. These predicates are NOT ORTHOGONAL. The Hub evaluates
197// them in the same order as the Outgoing set-vector, and control
198// branches to the first outgoing block whose predicate evaluates to true.
199//
200// The last guard block has two outgoing blocks as successors since the
201// condition for the final outgoing block is trivially true. So we create one
202// less block (including the first guard block) than the number of outgoing
203// blocks.
204static void convertToGuardPredicates(
205 ArrayRef<EdgeDescriptor> Branches, ArrayRef<BasicBlock *> Outgoing,
206 SmallVectorImpl<BasicBlock *> &GuardBlocks,
207 SmallVectorImpl<WeakVH> &DeletionCandidates, const StringRef Prefix,
208 std::optional<unsigned> MaxControlFlowBooleans) {
209 BBPredicates GuardPredicates;
210 Function *F = Outgoing.front()->getParent();
211
212 for (int I = 0, E = Outgoing.size() - 1; I != E; ++I)
213 GuardBlocks.push_back(
214 Elt: BasicBlock::Create(Context&: F->getContext(), Name: Prefix + ".guard", Parent: F));
215
216 // When we are using an integer to record which target block to jump to, we
217 // are creating less live values, actually we are using one single integer to
218 // store the index of the target block. When we are using booleans to store
219 // the branching information, we need (N-1) boolean values, where N is the
220 // number of outgoing block.
221 if (!MaxControlFlowBooleans || Outgoing.size() <= *MaxControlFlowBooleans)
222 calcPredicateUsingBooleans(Branches, Outgoing, GuardBlocks, GuardPredicates,
223 DeletionCandidates);
224 else
225 calcPredicateUsingInteger(Branches, Outgoing, GuardBlocks, GuardPredicates);
226
227 setupBranchForGuard(GuardBlocks, Outgoing, GuardPredicates);
228}
229
230// After creating a control flow hub, the operands of PHINodes in an outgoing
231// block Out no longer match the predecessors of that block. Predecessors of Out
232// that are incoming blocks to the hub are now replaced by just one edge from
233// the hub. To match this new control flow, the corresponding values from each
234// PHINode must now be moved a new PHINode in the first guard block of the hub.
235//
236// This operation cannot be performed with SSAUpdater, because it involves one
237// new use: If the block Out is in the list of Incoming blocks, then the newly
238// created PHI in the Hub will use itself along that edge from Out to Hub.
239static void reconnectPhis(BasicBlock *Out, BasicBlock *GuardBlock,
240 ArrayRef<EdgeDescriptor> Incoming,
241 BasicBlock *FirstGuardBlock) {
242 auto I = Out->begin();
243 while (I != Out->end() && isa<PHINode>(Val: I)) {
244 auto *Phi = cast<PHINode>(Val&: I);
245 auto *NewPhi =
246 PHINode::Create(Ty: Phi->getType(), NumReservedValues: Incoming.size(),
247 NameStr: Phi->getName() + ".moved", InsertBefore: FirstGuardBlock->begin());
248 bool AllUndef = true;
249 for (auto [BB, Succ0, Succ1] : Incoming) {
250 Value *V = PoisonValue::get(T: Phi->getType());
251 if (Phi->getBasicBlockIndex(BB) != -1) {
252 V = Phi->removeIncomingValue(BB, DeletePHIIfEmpty: false);
253 // When both successors are the same (Succ0 == Succ1), there are two
254 // edges from BB to Out, so we need to remove the second PHI entry too.
255 if (Succ0 && Succ1 && Succ0 == Succ1 &&
256 Phi->getBasicBlockIndex(BB) != -1)
257 Phi->removeIncomingValue(BB, DeletePHIIfEmpty: false);
258 if (BB == Out) {
259 V = NewPhi;
260 }
261 AllUndef &= isa<UndefValue>(Val: V);
262 }
263
264 NewPhi->addIncoming(V, BB);
265 }
266 assert(NewPhi->getNumIncomingValues() == Incoming.size());
267 Value *NewV = NewPhi;
268 if (AllUndef) {
269 NewPhi->eraseFromParent();
270 NewV = PoisonValue::get(T: Phi->getType());
271 }
272 if (Phi->getNumOperands() == 0) {
273 Phi->replaceAllUsesWith(V: NewV);
274 I = Phi->eraseFromParent();
275 continue;
276 }
277 Phi->addIncoming(V: NewV, BB: GuardBlock);
278 ++I;
279 }
280}
281
282std::pair<BasicBlock *, bool> ControlFlowHub::finalize(
283 DomTreeUpdater *DTU, SmallVectorImpl<BasicBlock *> &GuardBlocks,
284 const StringRef Prefix, std::optional<unsigned> MaxControlFlowBooleans) {
285#ifndef NDEBUG
286 SmallPtrSet<BasicBlock *, 8> Incoming;
287#endif
288 SetVector<BasicBlock *> Outgoing;
289
290 for (auto [BB, Succ0, Succ1] : Branches) {
291#ifndef NDEBUG
292 assert(
293 (Incoming.insert(BB).second || isa<CallBrInst>(BB->getTerminator())) &&
294 "Duplicate entry for incoming block.");
295#endif
296 if (Succ0)
297 Outgoing.insert(X: Succ0);
298 if (Succ1)
299 Outgoing.insert(X: Succ1);
300 }
301
302 assert(Outgoing.size() && "No outgoing edges");
303
304 if (Outgoing.size() < 2)
305 return {Outgoing.front(), false};
306
307 SmallVector<DominatorTree::UpdateType, 16> Updates;
308 if (DTU) {
309 for (auto [BB, Succ0, Succ1] : Branches) {
310 if (Succ0)
311 Updates.push_back(Elt: {DominatorTree::Delete, BB, Succ0});
312 // Only add Succ1 if it's different from Succ0 to avoid duplicate updates
313 if (Succ1 && Succ1 != Succ0)
314 Updates.push_back(Elt: {DominatorTree::Delete, BB, Succ1});
315 }
316 }
317
318 SmallVector<WeakVH, 8> DeletionCandidates;
319 convertToGuardPredicates(Branches, Outgoing: Outgoing.getArrayRef(), GuardBlocks,
320 DeletionCandidates, Prefix, MaxControlFlowBooleans);
321 BasicBlock *FirstGuardBlock = GuardBlocks.front();
322
323 // Update the PHINodes in each outgoing block to match the new control flow.
324 for (int I = 0, E = GuardBlocks.size(); I != E; ++I)
325 reconnectPhis(Out: Outgoing[I], GuardBlock: GuardBlocks[I], Incoming: Branches, FirstGuardBlock);
326 // Process the Nth (last) outgoing block with the (N-1)th (last) guard block.
327 reconnectPhis(Out: Outgoing.back(), GuardBlock: GuardBlocks.back(), Incoming: Branches, FirstGuardBlock);
328
329 if (DTU) {
330 int NumGuards = GuardBlocks.size();
331
332 for (auto [BB, Succ0, Succ1] : Branches)
333 Updates.push_back(Elt: {DominatorTree::Insert, BB, FirstGuardBlock});
334
335 for (int I = 0; I != NumGuards - 1; ++I) {
336 Updates.push_back(Elt: {DominatorTree::Insert, GuardBlocks[I], Outgoing[I]});
337 Updates.push_back(
338 Elt: {DominatorTree::Insert, GuardBlocks[I], GuardBlocks[I + 1]});
339 }
340 // The second successor of the last guard block is an outgoing block instead
341 // of having a "next" guard block.
342 Updates.push_back(Elt: {DominatorTree::Insert, GuardBlocks[NumGuards - 1],
343 Outgoing[NumGuards - 1]});
344 Updates.push_back(Elt: {DominatorTree::Insert, GuardBlocks[NumGuards - 1],
345 Outgoing[NumGuards]});
346 DTU->applyUpdates(Updates);
347 }
348
349 for (auto I : DeletionCandidates) {
350 if (I->use_empty())
351 if (auto *Inst = dyn_cast_or_null<Instruction>(Val&: I))
352 Inst->eraseFromParent();
353 }
354
355 return {FirstGuardBlock, true};
356}
357