[RISCV] Fix mgather -> riscv.masked.strided.load combine not extending indices (...
[llvm-project.git] / llvm / lib / CodeGen / MachineLoopUtils.cpp
blob0e8335d4974d722419b85a08a005beb78948bff8
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 //===----------------------------------------------------------------------===//
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;
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);
24 } // namespace
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());
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);
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);
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 const TargetRegisterClass *ConstrainRegClass =
67 MRI.constrainRegClass(R, MRI.getRegClass(Use->getReg()));
68 assert(ConstrainRegClass &&
69 "Expected a valid constrained register class!");
70 (void)ConstrainRegClass;
71 Use->setReg(R);
77 for (auto I = NewBB->getFirstNonPHI(); I != NewBB->end(); ++I)
78 for (MachineOperand &MO : I->uses())
79 if (MO.isReg() && Remaps.count(MO.getReg()))
80 MO.setReg(Remaps[MO.getReg()]);
82 for (auto I = NewBB->begin(); I->isPHI(); ++I) {
83 MachineInstr &MI = *I;
84 unsigned LoopRegIdx = 3, InitRegIdx = 1;
85 if (MI.getOperand(2).getMBB() != Preheader)
86 std::swap(LoopRegIdx, InitRegIdx);
87 MachineInstr &OrigPhi = findEquivalentInstruction(MI, Loop);
88 assert(OrigPhi.isPHI());
89 if (Direction == LPD_Front) {
90 // When peeling front, we are only left with the initial value from the
91 // preheader.
92 Register R = MI.getOperand(LoopRegIdx).getReg();
93 if (Remaps.count(R))
94 R = Remaps[R];
95 OrigPhi.getOperand(InitRegIdx).setReg(R);
96 MI.removeOperand(LoopRegIdx + 1);
97 MI.removeOperand(LoopRegIdx + 0);
98 } else {
99 // When peeling back, the initial value is the loop-carried value from
100 // the original loop.
101 Register LoopReg = OrigPhi.getOperand(LoopRegIdx).getReg();
102 MI.getOperand(LoopRegIdx).setReg(LoopReg);
103 MI.removeOperand(InitRegIdx + 1);
104 MI.removeOperand(InitRegIdx + 0);
108 DebugLoc DL;
109 if (Direction == LPD_Front) {
110 Preheader->ReplaceUsesOfBlockWith(Loop, NewBB);
111 NewBB->addSuccessor(Loop);
112 Loop->replacePhiUsesWith(Preheader, NewBB);
113 Preheader->updateTerminator(Loop);
114 TII->removeBranch(*NewBB);
115 TII->insertBranch(*NewBB, Loop, nullptr, {}, DL);
116 } else {
117 Loop->replaceSuccessor(Exit, NewBB);
118 Exit->replacePhiUsesWith(Loop, NewBB);
119 NewBB->addSuccessor(Exit);
121 MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
122 SmallVector<MachineOperand, 4> Cond;
123 bool CanAnalyzeBr = !TII->analyzeBranch(*Loop, TBB, FBB, Cond);
124 (void)CanAnalyzeBr;
125 assert(CanAnalyzeBr && "Must be able to analyze the loop branch!");
126 TII->removeBranch(*Loop);
127 TII->insertBranch(*Loop, TBB == Exit ? NewBB : TBB,
128 FBB == Exit ? NewBB : FBB, Cond, DL);
129 if (TII->removeBranch(*NewBB) > 0)
130 TII->insertBranch(*NewBB, Exit, nullptr, {}, DL);
133 return NewBB;