1 //=- MachineLoopUtils.cpp - Functions for manipulating loops ----------------=//
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/CodeGen/MachineLoopUtils.h"
10 #include "llvm/CodeGen/MachineBasicBlock.h"
11 #include "llvm/CodeGen/MachineRegisterInfo.h"
12 #include "llvm/CodeGen/TargetInstrInfo.h"
13 using namespace llvm;
14 
15 namespace {
16 // MI's parent and BB are clones of each other. Find the equivalent copy of MI
17 // in BB.
18 MachineInstr &findEquivalentInstruction(MachineInstr &MI,
19                                         MachineBasicBlock *BB) {
20   MachineBasicBlock *PB = MI.getParent();
21   unsigned Offset = std::distance(PB->instr_begin(), MachineBasicBlock::instr_iterator(MI));
22   return *std::next(BB->instr_begin(), Offset);
23 }
24 } // namespace
25 
26 MachineBasicBlock *llvm::PeelSingleBlockLoop(LoopPeelDirection Direction,
27                                              MachineBasicBlock *Loop,
28                                              MachineRegisterInfo &MRI,
29                                              const TargetInstrInfo *TII) {
30   MachineFunction &MF = *Loop->getParent();
31   MachineBasicBlock *Preheader = *Loop->pred_begin();
32   if (Preheader == Loop)
33     Preheader = *std::next(Loop->pred_begin());
34   MachineBasicBlock *Exit = *Loop->succ_begin();
35   if (Exit == Loop)
36     Exit = *std::next(Loop->succ_begin());
37 
38   MachineBasicBlock *NewBB = MF.CreateMachineBasicBlock(Loop->getBasicBlock());
39   if (Direction == LPD_Front)
40     MF.insert(Loop->getIterator(), NewBB);
41   else
42     MF.insert(std::next(Loop->getIterator()), NewBB);
43 
44   DenseMap<Register, Register> Remaps;
45   auto InsertPt = NewBB->end();
46   for (MachineInstr &MI : *Loop) {
47     MachineInstr *NewMI = MF.CloneMachineInstr(&MI);
48     NewBB->insert(InsertPt, NewMI);
49     for (MachineOperand &MO : NewMI->defs()) {
50       Register OrigR = MO.getReg();
51       if (OrigR.isPhysical())
52         continue;
53       Register &R = Remaps[OrigR];
54       R = MRI.createVirtualRegister(MRI.getRegClass(OrigR));
55       MO.setReg(R);
56 
57       if (Direction == LPD_Back) {
58         // Replace all uses outside the original loop with the new register.
59         // FIXME: is the use_iterator stable enough to mutate register uses
60         // while iterating?
61         SmallVector<MachineOperand *, 4> Uses;
62         for (auto &Use : MRI.use_operands(OrigR))
63           if (Use.getParent()->getParent() != Loop)
64             Uses.push_back(&Use);
65         for (auto *Use : Uses) {
66           MRI.constrainRegClass(R, MRI.getRegClass(Use->getReg()));
67           Use->setReg(R);
68         }
69       }
70     }
71   }
72 
73   for (auto I = NewBB->getFirstNonPHI(); I != NewBB->end(); ++I)
74     for (MachineOperand &MO : I->uses())
75       if (MO.isReg() && Remaps.count(MO.getReg()))
76         MO.setReg(Remaps[MO.getReg()]);
77 
78   for (auto I = NewBB->begin(); I->isPHI(); ++I) {
79     MachineInstr &MI = *I;
80     unsigned LoopRegIdx = 3, InitRegIdx = 1;
81     if (MI.getOperand(2).getMBB() != Preheader)
82       std::swap(LoopRegIdx, InitRegIdx);
83     MachineInstr &OrigPhi = findEquivalentInstruction(MI, Loop);
84     assert(OrigPhi.isPHI());
85     if (Direction == LPD_Front) {
86       // When peeling front, we are only left with the initial value from the
87       // preheader.
88       Register R = MI.getOperand(LoopRegIdx).getReg();
89       if (Remaps.count(R))
90         R = Remaps[R];
91       OrigPhi.getOperand(InitRegIdx).setReg(R);
92       MI.removeOperand(LoopRegIdx + 1);
93       MI.removeOperand(LoopRegIdx + 0);
94     } else {
95       // When peeling back, the initial value is the loop-carried value from
96       // the original loop.
97       Register LoopReg = OrigPhi.getOperand(LoopRegIdx).getReg();
98       MI.getOperand(LoopRegIdx).setReg(LoopReg);
99       MI.removeOperand(InitRegIdx + 1);
100       MI.removeOperand(InitRegIdx + 0);
101     }
102   }
103 
104   DebugLoc DL;
105   if (Direction == LPD_Front) {
106     Preheader->ReplaceUsesOfBlockWith(Loop, NewBB);
107     NewBB->addSuccessor(Loop);
108     Loop->replacePhiUsesWith(Preheader, NewBB);
109     Preheader->updateTerminator(Loop);
110     TII->removeBranch(*NewBB);
111     TII->insertBranch(*NewBB, Loop, nullptr, {}, DL);
112   } else {
113     Loop->replaceSuccessor(Exit, NewBB);
114     Exit->replacePhiUsesWith(Loop, NewBB);
115     NewBB->addSuccessor(Exit);
116 
117     MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
118     SmallVector<MachineOperand, 4> Cond;
119     bool CanAnalyzeBr = !TII->analyzeBranch(*Loop, TBB, FBB, Cond);
120     (void)CanAnalyzeBr;
121     assert(CanAnalyzeBr && "Must be able to analyze the loop branch!");
122     TII->removeBranch(*Loop);
123     TII->insertBranch(*Loop, TBB == Exit ? NewBB : TBB,
124                       FBB == Exit ? NewBB : FBB, Cond, DL);
125     if (TII->removeBranch(*NewBB) > 0)
126       TII->insertBranch(*NewBB, Exit, nullptr, {}, DL);
127   }
128 
129   return NewBB;
130 }
131