1//===------- LoopBoundSplit.cpp - Split Loop Bound --------------*- 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/Scalar/LoopBoundSplit.h"
10#include "llvm/ADT/Sequence.h"
11#include "llvm/Analysis/LoopAnalysisManager.h"
12#include "llvm/Analysis/LoopInfo.h"
13#include "llvm/Analysis/ScalarEvolution.h"
14#include "llvm/Analysis/ScalarEvolutionExpressions.h"
15#include "llvm/IR/PatternMatch.h"
16#include "llvm/Transforms/Scalar/LoopPassManager.h"
17#include "llvm/Transforms/Utils/BasicBlockUtils.h"
18#include "llvm/Transforms/Utils/Cloning.h"
19#include "llvm/Transforms/Utils/LoopSimplify.h"
20#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
21
22#define DEBUG_TYPE "loop-bound-split"
23
24using namespace llvm;
25using namespace PatternMatch;
26
27namespace {
28struct ConditionInfo {
29 /// Branch instruction with this condition
30 CondBrInst *BI = nullptr;
31 /// ICmp instruction with this condition
32 ICmpInst *ICmp = nullptr;
33 /// Preciate info
34 CmpPredicate Pred = ICmpInst::BAD_ICMP_PREDICATE;
35 /// AddRec llvm value
36 Value *AddRecValue = nullptr;
37 /// Non PHI AddRec llvm value
38 Value *NonPHIAddRecValue = nullptr;
39 /// Bound llvm value
40 Value *BoundValue = nullptr;
41 /// AddRec SCEV
42 const SCEVAddRecExpr *AddRecSCEV = nullptr;
43 /// Bound SCEV
44 const SCEV *BoundSCEV = nullptr;
45
46 ConditionInfo() = default;
47};
48} // namespace
49
50static bool calculateUpperBound(const Loop &L, ScalarEvolution &SE,
51 ConditionInfo &Cond, bool IsExitCond) {
52 if (IsExitCond) {
53 const SCEV *ExitCount = SE.getExitCount(L: &L, ExitingBlock: Cond.ICmp->getParent());
54 if (isa<SCEVCouldNotCompute>(Val: ExitCount))
55 return false;
56
57 Cond.BoundSCEV = ExitCount;
58 return true;
59 }
60
61 // For non-exit condtion, if pred is LT, keep existing bound.
62 if (Cond.Pred == ICmpInst::ICMP_SLT || Cond.Pred == ICmpInst::ICMP_ULT)
63 return true;
64
65 // For non-exit condition, if pre is LE, try to convert it to LT.
66 // Range Range
67 // AddRec <= Bound --> AddRec < Bound + 1
68 if (Cond.Pred != ICmpInst::ICMP_ULE && Cond.Pred != ICmpInst::ICMP_SLE)
69 return false;
70
71 if (IntegerType *BoundSCEVIntType =
72 dyn_cast<IntegerType>(Val: Cond.BoundSCEV->getType())) {
73 unsigned BitWidth = BoundSCEVIntType->getBitWidth();
74 APInt Max = ICmpInst::isSigned(Pred: Cond.Pred)
75 ? APInt::getSignedMaxValue(numBits: BitWidth)
76 : APInt::getMaxValue(numBits: BitWidth);
77 const SCEV *MaxSCEV = SE.getConstant(Val: Max);
78 // Check Bound < INT_MAX
79 ICmpInst::Predicate Pred =
80 ICmpInst::isSigned(Pred: Cond.Pred) ? ICmpInst::ICMP_SLT : ICmpInst::ICMP_ULT;
81 if (SE.isKnownPredicate(Pred, LHS: Cond.BoundSCEV, RHS: MaxSCEV)) {
82 const SCEV *BoundPlusOneSCEV =
83 SE.getAddExpr(LHS: Cond.BoundSCEV, RHS: SE.getOne(Ty: BoundSCEVIntType));
84 Cond.BoundSCEV = BoundPlusOneSCEV;
85 Cond.Pred = Pred;
86 return true;
87 }
88 }
89
90 // ToDo: Support ICMP_NE/EQ.
91
92 return false;
93}
94
95/// Check whether \p ICmp compares an induction variable of \p L against a
96/// bound this pass can split on, and describe it in \p Cond if so.
97static bool hasProcessableCondition(const Loop &L, ScalarEvolution &SE,
98 ICmpInst *ICmp, ConditionInfo &Cond,
99 bool IsExitCond) {
100 Cond.ICmp = ICmp;
101 if (!match(V: ICmp, P: m_ICmp(Pred&: Cond.Pred, L: m_Value(V&: Cond.AddRecValue),
102 R: m_Value(V&: Cond.BoundValue))))
103 return false;
104
105 const SCEV *AddRecSCEV = SE.getSCEV(V: Cond.AddRecValue);
106 const SCEV *BoundSCEV = SE.getSCEV(V: Cond.BoundValue);
107 // Locate the recurrence in AddRecSCEV and the bound in BoundSCEV.
108 if (!isa<SCEVAddRecExpr>(Val: AddRecSCEV) && isa<SCEVAddRecExpr>(Val: BoundSCEV)) {
109 std::swap(a&: Cond.AddRecValue, b&: Cond.BoundValue);
110 std::swap(a&: AddRecSCEV, b&: BoundSCEV);
111 Cond.Pred = ICmpInst::getSwappedPredicate(pred: Cond.Pred);
112 }
113
114 // Allowed AddRec as induction variable.
115 Cond.AddRecSCEV = dyn_cast<SCEVAddRecExpr>(Val: AddRecSCEV);
116 if (!Cond.AddRecSCEV)
117 return false;
118
119 // If the induction variable is a PHI node, the value from the backedge is
120 // used instead.
121 Cond.NonPHIAddRecValue = Cond.AddRecValue;
122 if (auto *PN = dyn_cast<PHINode>(Val: Cond.AddRecValue))
123 Cond.NonPHIAddRecValue = PN->getIncomingValueForBlock(BB: L.getLoopLatch());
124
125 // The BoundSCEV should be evaluated at loop entry.
126 Cond.BoundSCEV = BoundSCEV;
127 if (!SE.isAvailableAtLoopEntry(S: Cond.BoundSCEV, L: &L))
128 return false;
129
130 if (!Cond.AddRecSCEV->isAffine())
131 return false;
132
133 // Allowed constant step.
134 const auto *StepRecSCEV =
135 dyn_cast<SCEVConstant>(Val: Cond.AddRecSCEV->getStepRecurrence(SE));
136 if (!StepRecSCEV)
137 return false;
138
139 // Allowed positive step for now.
140 // TODO: Support negative step.
141 ConstantInt *StepCI = StepRecSCEV->getValue();
142 if (StepCI->isNegative() || StepCI->isZero())
143 return false;
144
145 // Calculate upper bound.
146 if (!calculateUpperBound(L, SE, Cond, IsExitCond))
147 return false;
148
149 return true;
150}
151
152static bool isProcessableCondBI(const ScalarEvolution &SE,
153 const CondBrInst *BI) {
154 BasicBlock *TrueSucc = nullptr;
155 BasicBlock *FalseSucc = nullptr;
156 Value *LHS, *RHS;
157 if (!match(V: BI, P: m_Br(C: m_ICmp(L: m_Value(V&: LHS), R: m_Value(V&: RHS)),
158 T: m_BasicBlock(V&: TrueSucc), F: m_BasicBlock(V&: FalseSucc))))
159 return false;
160
161 if (!SE.isSCEVable(Ty: LHS->getType()))
162 return false;
163 assert(SE.isSCEVable(RHS->getType()) && "Expected RHS's type is SCEVable");
164
165 if (TrueSucc == FalseSucc)
166 return false;
167
168 return true;
169}
170
171static bool canSplitLoopBound(const Loop &L, const DominatorTree &DT,
172 ScalarEvolution &SE, ConditionInfo &Cond) {
173 // Skip function with optsize.
174 if (L.getHeader()->getParent()->hasOptSize())
175 return false;
176
177 // Split only innermost loop.
178 if (!L.isInnermost())
179 return false;
180
181 // Check loop is in simplified form.
182 if (!L.isLoopSimplifyForm())
183 return false;
184
185 // Check loop is in LCSSA form.
186 if (!L.isLCSSAForm(DT))
187 return false;
188
189 // Skip loop that cannot be cloned.
190 if (!L.isSafeToClone())
191 return false;
192
193 BasicBlock *ExitingBB = L.getExitingBlock();
194 // Assumed only one exiting block.
195 if (!ExitingBB)
196 return false;
197
198 CondBrInst *ExitingBI = dyn_cast<CondBrInst>(Val: ExitingBB->getTerminator());
199 if (!ExitingBI)
200 return false;
201
202 // Allowed only conditional branch with ICmp.
203 if (!isProcessableCondBI(SE, BI: ExitingBI))
204 return false;
205
206 // Check the condition is processable.
207 ICmpInst *ICmp = cast<ICmpInst>(Val: ExitingBI->getCondition());
208 if (!hasProcessableCondition(L, SE, ICmp, Cond, /*IsExitCond*/ true))
209 return false;
210
211 Cond.BI = ExitingBI;
212 return true;
213}
214
215static bool isProfitableToTransform(const Loop &L, const CondBrInst *BI) {
216 // If the conditional branch splits a loop into two halves, we could
217 // generally say it is profitable.
218 //
219 // ToDo: Add more profitable cases here.
220
221 // Check this branch causes diamond CFG.
222 BasicBlock *Succ0 = BI->getSuccessor(i: 0);
223 BasicBlock *Succ1 = BI->getSuccessor(i: 1);
224
225 BasicBlock *Succ0Succ = Succ0->getSingleSuccessor();
226 BasicBlock *Succ1Succ = Succ1->getSingleSuccessor();
227 if (!Succ0Succ || !Succ1Succ || Succ0Succ != Succ1Succ)
228 return false;
229
230 // ToDo: Calculate each successor's instruction cost.
231
232 return true;
233}
234
235static CondBrInst *findSplitCandidate(const Loop &L, ScalarEvolution &SE,
236 ConditionInfo &ExitingCond,
237 ConditionInfo &SplitCandidateCond) {
238 for (auto *BB : L.blocks()) {
239 // Skip condition of backedge.
240 if (L.getLoopLatch() == BB)
241 continue;
242
243 auto *BI = dyn_cast<CondBrInst>(Val: BB->getTerminator());
244 if (!BI)
245 continue;
246
247 // Check conditional branch with ICmp.
248 if (!isProcessableCondBI(SE, BI))
249 continue;
250
251 // Skip loop invariant condition.
252 if (L.isLoopInvariant(V: BI->getCondition()))
253 continue;
254
255 // Check the condition is processable.
256 ICmpInst *ICmp = cast<ICmpInst>(Val: BI->getCondition());
257 if (!hasProcessableCondition(L, SE, ICmp, Cond&: SplitCandidateCond,
258 /*IsExitCond*/ false))
259 continue;
260
261 if (ExitingCond.BoundSCEV->getType() !=
262 SplitCandidateCond.BoundSCEV->getType())
263 continue;
264
265 // After transformation, we assume the split condition of the pre-loop is
266 // always true. In order to guarantee it, we need to check the start value
267 // of the split cond AddRec satisfies the split condition.
268 if (!SE.isLoopEntryGuardedByCond(L: &L, Pred: SplitCandidateCond.Pred,
269 LHS: SplitCandidateCond.AddRecSCEV->getStart(),
270 RHS: SplitCandidateCond.BoundSCEV))
271 continue;
272
273 SplitCandidateCond.BI = BI;
274 return BI;
275 }
276
277 return nullptr;
278}
279
280static bool splitLoopBound(Loop &L, DominatorTree &DT, LoopInfo &LI,
281 ScalarEvolution &SE, LPMUpdater &U) {
282 ConditionInfo SplitCandidateCond;
283 ConditionInfo ExitingCond;
284
285 // Check we can split this loop's bound.
286 if (!canSplitLoopBound(L, DT, SE, Cond&: ExitingCond))
287 return false;
288
289 if (!findSplitCandidate(L, SE, ExitingCond, SplitCandidateCond))
290 return false;
291
292 if (!isProfitableToTransform(L, BI: SplitCandidateCond.BI))
293 return false;
294
295 // Now, we have a split candidate. Let's build a form as below.
296 // +--------------------+
297 // | preheader |
298 // | set up newbound |
299 // +--------------------+
300 // | /----------------\
301 // +--------v----v------+ |
302 // | header |---\ |
303 // | with true condition| | |
304 // +--------------------+ | |
305 // | | |
306 // +--------v-----------+ | |
307 // | if.then.BB | | |
308 // +--------------------+ | |
309 // | | |
310 // +--------v-----------<---/ |
311 // | latch >----------/
312 // | with newbound |
313 // +--------------------+
314 // |
315 // +--------v-----------+
316 // | preheader2 |--------------\
317 // | if (AddRec i != | |
318 // | org bound) | |
319 // +--------------------+ |
320 // | /----------------\ |
321 // +--------v----v------+ | |
322 // | header2 |---\ | |
323 // | conditional branch | | | |
324 // |with false condition| | | |
325 // +--------------------+ | | |
326 // | | | |
327 // +--------v-----------+ | | |
328 // | if.then.BB2 | | | |
329 // +--------------------+ | | |
330 // | | | |
331 // +--------v-----------<---/ | |
332 // | latch2 >----------/ |
333 // | with org bound | |
334 // +--------v-----------+ |
335 // | |
336 // | +---------------+ |
337 // +--> exit <-------/
338 // +---------------+
339
340 // Let's create post loop.
341 SmallVector<BasicBlock *, 8> PostLoopBlocks;
342 Loop *PostLoop;
343 ValueToValueMapTy VMap;
344 BasicBlock *PreHeader = L.getLoopPreheader();
345 BasicBlock *SplitLoopPH = SplitEdge(From: PreHeader, To: L.getHeader(), DT: &DT, LI: &LI);
346 PostLoop = cloneLoopWithPreheader(Before: L.getExitBlock(), LoopDomBB: SplitLoopPH, OrigLoop: &L, VMap,
347 NameSuffix: ".split", LI: &LI, DT: &DT, Blocks&: PostLoopBlocks);
348 remapInstructionsInBlocks(Blocks: PostLoopBlocks, VMap);
349
350 BasicBlock *PostLoopPreHeader = PostLoop->getLoopPreheader();
351 IRBuilder<> Builder(&PostLoopPreHeader->front());
352
353 // Replace exit branch target of pre-loop by post-loop's preheader.
354 // Note: update the branch here after calling cloneLoopWithPreheader()
355 // to keep the IR valid.
356 if (L.getExitBlock() == ExitingCond.BI->getSuccessor(i: 0))
357 ExitingCond.BI->setSuccessor(idx: 0, NewSucc: PostLoopPreHeader);
358 else
359 ExitingCond.BI->setSuccessor(idx: 1, NewSucc: PostLoopPreHeader);
360
361 // Update dominator tree.
362 DT.changeImmediateDominator(BB: PostLoopPreHeader, NewBB: L.getExitingBlock());
363#ifndef NDEBUG
364 LI.verify();
365#endif
366 // Update phi nodes in header of post-loop.
367 bool isExitingLatch = L.getExitingBlock() == L.getLoopLatch();
368 Value *ExitingCondLCSSAPhi = nullptr;
369 for (PHINode &PN : L.getHeader()->phis()) {
370 // Create LCSSA phi node in preheader of post-loop.
371 PHINode *LCSSAPhi =
372 Builder.CreatePHI(Ty: PN.getType(), NumReservedValues: 1, Name: PN.getName() + ".lcssa");
373 LCSSAPhi->setDebugLoc(PN.getDebugLoc());
374 // If the exiting block is loop latch, the phi does not have the update at
375 // last iteration. In this case, update lcssa phi with value from backedge.
376 LCSSAPhi->addIncoming(
377 V: isExitingLatch ? PN.getIncomingValueForBlock(BB: L.getLoopLatch()) : &PN,
378 BB: L.getExitingBlock());
379
380 // Update the start value of phi node in post-loop with the LCSSA phi node.
381 PHINode *PostLoopPN = cast<PHINode>(Val&: VMap[&PN]);
382 PostLoopPN->setIncomingValueForBlock(BB: PostLoopPreHeader, V: LCSSAPhi);
383
384 // Find PHI with exiting condition from pre-loop. The PHI should be
385 // SCEVAddRecExpr and have same incoming value from backedge with
386 // ExitingCond.
387 //
388 // TODO: Separate SCEV queries from PHI node updates.
389 if (!SE.isSCEVable(Ty: PN.getType()))
390 continue;
391
392 const SCEVAddRecExpr *PhiSCEV = dyn_cast<SCEVAddRecExpr>(Val: SE.getSCEV(V: &PN));
393 if (PhiSCEV && ExitingCond.NonPHIAddRecValue ==
394 PN.getIncomingValueForBlock(BB: L.getLoopLatch()))
395 ExitingCondLCSSAPhi = LCSSAPhi;
396 }
397
398 // Add conditional branch to check we can skip post-loop in its preheader,
399 // and update DT.
400 Instruction *OrigBI = PostLoopPreHeader->getTerminator();
401 ICmpInst::Predicate Pred = ICmpInst::ICMP_NE;
402 Value *Cond =
403 Builder.CreateICmp(P: Pred, LHS: ExitingCondLCSSAPhi, RHS: ExitingCond.BoundValue);
404 Builder.CreateCondBr(Cond, True: PostLoop->getHeader(), False: PostLoop->getExitBlock());
405 OrigBI->eraseFromParent();
406 DT.changeImmediateDominator(BB: PostLoop->getExitBlock(), NewBB: PostLoopPreHeader);
407#ifdef EXPENSIVE_CHECKS
408 assert(DT.verify(DominatorTree::VerificationLevel::Full) &&
409 "DT broken during transformation!");
410#else
411 assert(DT.verify(DominatorTree::VerificationLevel::Fast) &&
412 "DT broken during transformation!");
413#endif
414
415 // Create new loop bound and add it into preheader of pre-loop.
416 const SCEV *NewBoundSCEV = ExitingCond.BoundSCEV;
417 const SCEV *SplitBoundSCEV = SplitCandidateCond.BoundSCEV;
418 NewBoundSCEV = ICmpInst::isSigned(Pred: ExitingCond.Pred)
419 ? SE.getSMinExpr(LHS: NewBoundSCEV, RHS: SplitBoundSCEV)
420 : SE.getUMinExpr(LHS: NewBoundSCEV, RHS: SplitBoundSCEV);
421
422 SCEVExpander Expander(SE, "split");
423 Instruction *InsertPt = SplitLoopPH->getTerminator();
424 Value *NewBoundValue =
425 Expander.expandCodeFor(SH: NewBoundSCEV, Ty: NewBoundSCEV->getType(), I: InsertPt);
426 NewBoundValue->setName("new.bound");
427
428 // Replace exiting bound value of pre-loop NewBound.
429 ExitingCond.ICmp->setOperand(i_nocapture: 1, Val_nocapture: NewBoundValue);
430
431 // Replace SplitCandidateCond.BI's condition of pre-loop by True.
432 LLVMContext &Context = PreHeader->getContext();
433 SplitCandidateCond.BI->setCondition(ConstantInt::getTrue(Context));
434
435 // Replace cloned SplitCandidateCond.BI's condition in post-loop by False.
436 CondBrInst *ClonedSplitCandidateBI =
437 cast<CondBrInst>(Val&: VMap[SplitCandidateCond.BI]);
438 ClonedSplitCandidateBI->setCondition(ConstantInt::getFalse(Context));
439
440 // Update phi node in exit block of post-loop.
441 Builder.SetInsertPoint(TheBB: PostLoopPreHeader, IP: PostLoopPreHeader->begin());
442 for (PHINode &PN : PostLoop->getExitBlock()->phis()) {
443 for (auto i : seq<int>(Begin: 0, End: PN.getNumOperands())) {
444 // Check incoming block is pre-loop's exiting block.
445 if (PN.getIncomingBlock(i) == L.getExitingBlock()) {
446 Value *IncomingValue = PN.getIncomingValue(i);
447
448 // Create LCSSA phi node for incoming value.
449 PHINode *LCSSAPhi =
450 Builder.CreatePHI(Ty: PN.getType(), NumReservedValues: 1, Name: PN.getName() + ".lcssa");
451 LCSSAPhi->setDebugLoc(PN.getDebugLoc());
452 LCSSAPhi->addIncoming(V: IncomingValue, BB: PN.getIncomingBlock(i));
453
454 // Replace pre-loop's exiting block by post-loop's preheader.
455 PN.setIncomingBlock(i, BB: PostLoopPreHeader);
456 // Replace incoming value by LCSSAPhi.
457 PN.setIncomingValue(i, V: LCSSAPhi);
458 // Add a new incoming value with post-loop's exiting block.
459 PN.addIncoming(V: VMap[IncomingValue], BB: PostLoop->getExitingBlock());
460 }
461 }
462 }
463
464 // Invalidate cached SE information.
465 SE.forgetLoop(L: &L);
466
467 // Canonicalize loops.
468 simplifyLoop(L: &L, DT: &DT, LI: &LI, SE: &SE, AC: nullptr, MSSAU: nullptr, PreserveLCSSA: true);
469 simplifyLoop(L: PostLoop, DT: &DT, LI: &LI, SE: &SE, AC: nullptr, MSSAU: nullptr, PreserveLCSSA: true);
470
471 // Add new post-loop to loop pass manager.
472 U.addSiblingLoops(NewSibLoops: PostLoop);
473
474 return true;
475}
476
477PreservedAnalyses LoopBoundSplitPass::run(Loop &L, LoopAnalysisManager &AM,
478 LoopStandardAnalysisResults &AR,
479 LPMUpdater &U) {
480 [[maybe_unused]] Function &F = *L.getHeader()->getParent();
481
482 LLVM_DEBUG(dbgs() << "Spliting bound of loop in " << F.getName() << ": " << L
483 << "\n");
484
485 if (!splitLoopBound(L, DT&: AR.DT, LI&: AR.LI, SE&: AR.SE, U))
486 return PreservedAnalyses::all();
487
488 assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
489 AR.LI.verify();
490
491 return getLoopPassPreservedAnalyses();
492}
493