lib/CodeGen/MachineLoopUtils.cpp

   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   // FIXME: Add DenseMapInfo trait for Register so we can use it as a key.
  45   DenseMap<unsigned, Register> Remaps;
  46   auto InsertPt = NewBB->end();
  47   for (MachineInstr &MI : *Loop) {
  48     MachineInstr *NewMI = MF.CloneMachineInstr(&MI);
  49     NewBB->insert(InsertPt, NewMI);
  50     for (MachineOperand &MO : NewMI->defs()) {
  51       Register OrigR = MO.getReg();
  52       if (OrigR.isPhysical())
  53         continue;
  54       Register &R = Remaps[OrigR];
  55       R = MRI.createVirtualRegister(MRI.getRegClass(OrigR));
  56       MO.setReg(R);
  57
  58       if (Direction == LPD_Back) {
  59         // Replace all uses outside the original loop with the new register.
  60         // FIXME: is the use_iterator stable enough to mutate register uses
  61         // while iterating?
  62         SmallVector<MachineOperand *, 4> Uses;
  63         for (auto &Use : MRI.use_operands(OrigR))
  64           if (Use.getParent()->getParent() != Loop)
  65             Uses.push_back(&Use);
  66         for (auto *Use : Uses) {
  67           MRI.constrainRegClass(R, MRI.getRegClass(Use->getReg()));
  68           Use->setReg(R);
  69         }
  70       }
  71     }
  72   }
  73
  74   for (auto I = NewBB->getFirstNonPHI(); I != NewBB->end(); ++I)
  75     for (MachineOperand &MO : I->uses())
  76       if (MO.isReg() && Remaps.count(MO.getReg()))
  77         MO.setReg(Remaps[MO.getReg()]);
  78
  79   for (auto I = NewBB->begin(); I->isPHI(); ++I) {
  80     MachineInstr &MI = *I;
  81     unsigned LoopRegIdx = 3, InitRegIdx = 1;
  82     if (MI.getOperand(2).getMBB() != Preheader)
  83       std::swap(LoopRegIdx, InitRegIdx);
  84     MachineInstr &OrigPhi = findEquivalentInstruction(MI, Loop);
  85     assert(OrigPhi.isPHI());
  86     if (Direction == LPD_Front) {
  87       // When peeling front, we are only left with the initial value from the
  88       // preheader.
  89       Register R = MI.getOperand(LoopRegIdx).getReg();
  90       if (Remaps.count(R))
  91         R = Remaps[R];
  92       OrigPhi.getOperand(InitRegIdx).setReg(R);
  93       MI.RemoveOperand(LoopRegIdx + 1);
  94       MI.RemoveOperand(LoopRegIdx + 0);
  95     } else {
  96       // When peeling back, the initial value is the loop-carried value from
  97       // the original loop.
  98       Register LoopReg = OrigPhi.getOperand(LoopRegIdx).getReg();
  99       MI.getOperand(LoopRegIdx).setReg(LoopReg);
 100       MI.RemoveOperand(InitRegIdx + 1);
 101       MI.RemoveOperand(InitRegIdx + 0);
 102     }
 103   }
 104
 105   DebugLoc DL;
 106   if (Direction == LPD_Front) {
 107     Preheader->replaceSuccessor(Loop, NewBB);
 108     NewBB->addSuccessor(Loop);
 109     Loop->replacePhiUsesWith(Preheader, NewBB);
 110     if (TII->removeBranch(*Preheader) > 0)
 111       TII->insertBranch(*Preheader, NewBB, nullptr, {}, DL);
 112     TII->removeBranch(*NewBB);
 113     TII->insertBranch(*NewBB, Loop, nullptr, {}, DL);
 114   } else {
 115     Loop->replaceSuccessor(Exit, NewBB);
 116     Exit->replacePhiUsesWith(Loop, NewBB);
 117     NewBB->addSuccessor(Exit);
 118
 119     MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
 120     SmallVector<MachineOperand, 4> Cond;
 121     bool CanAnalyzeBr = !TII->analyzeBranch(*Loop, TBB, FBB, Cond);
 122     (void)CanAnalyzeBr;
 123     assert(CanAnalyzeBr && "Must be able to analyze the loop branch!");
 124     TII->removeBranch(*Loop);
 125     TII->insertBranch(*Loop, TBB == Exit ? NewBB : TBB,
 126                       FBB == Exit ? NewBB : FBB, Cond, DL);
 127     if (TII->removeBranch(*NewBB) > 0)
 128       TII->insertBranch(*NewBB, Exit, nullptr, {}, DL);
 129   }
 130
 131   return NewBB;
 132 }