[AMDGPU] Split waterfall loop exec manipulation

Split waterfall loops into multiple blocks so that exec mask
manipulation (s_and_saveexec) does not occur in the middle of
a block.

VGPR live range optimizer is updated to handle waterfall loops
spanning multiple blocks.

Reviewed By: ruiling

Differential Revision: https://reviews.llvm.org/D122200

Author: perlfu

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

@@ -741,16 +741,19 @@ bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
   // To insert the loop we need to split the block. Move everything before this
   // point to a new block, and insert a new empty block before this instruction.
   MachineBasicBlock *LoopBB = MF->CreateMachineBasicBlock();
+  MachineBasicBlock *BodyBB = MF->CreateMachineBasicBlock();
   MachineBasicBlock *RemainderBB = MF->CreateMachineBasicBlock();
   MachineBasicBlock *RestoreExecBB = MF->CreateMachineBasicBlock();
   MachineFunction::iterator MBBI(MBB);
   ++MBBI;
   MF->insert(MBBI, LoopBB);
+  MF->insert(MBBI, BodyBB);
   MF->insert(MBBI, RestoreExecBB);
   MF->insert(MBBI, RemainderBB);
 
-  LoopBB->addSuccessor(RestoreExecBB);
-  LoopBB->addSuccessor(LoopBB);
+  LoopBB->addSuccessor(BodyBB);
+  BodyBB->addSuccessor(RestoreExecBB);
+  BodyBB->addSuccessor(LoopBB);
 
   // Move the rest of the block into a new block.
   RemainderBB->transferSuccessorsAndUpdatePHIs(&MBB);
@@ -762,26 +765,26 @@ bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
   B.setInsertPt(*LoopBB, LoopBB->end());
 
   B.buildInstr(TargetOpcode::PHI)
-    .addDef(PhiExec)
-    .addReg(InitSaveExecReg)
-    .addMBB(&MBB)
-    .addReg(NewExec)
-    .addMBB(LoopBB);
+      .addDef(PhiExec)
+      .addReg(InitSaveExecReg)
+      .addMBB(&MBB)
+      .addReg(NewExec)
+      .addMBB(BodyBB);
 
   const DebugLoc &DL = B.getDL();
 
   MachineInstr &FirstInst = *Range.begin();
 
-  // Move the instruction into the loop. Note we moved everything after
+  // Move the instruction into the loop body. Note we moved everything after
   // Range.end() already into a new block, so Range.end() is no longer valid.
-  LoopBB->splice(LoopBB->end(), &MBB, Range.begin(), MBB.end());
+  BodyBB->splice(BodyBB->end(), &MBB, Range.begin(), MBB.end());
 
   // Figure out the iterator range after splicing the instructions.
   MachineBasicBlock::iterator NewBegin = FirstInst.getIterator();
-  auto NewEnd = LoopBB->end();
+  auto NewEnd = BodyBB->end();
 
-  MachineBasicBlock::iterator I = Range.begin();
-  B.setInsertPt(*LoopBB, I);
+  MachineBasicBlock::iterator I = LoopBB->end();
+  B.setMBB(*LoopBB);
 
   Register CondReg;
 
@@ -813,7 +816,7 @@ bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
         B.setMBB(MBB);
         OpReg = B.buildCopy(OpTy, OpReg).getReg(0);
         MRI.setRegBank(OpReg, AMDGPU::VGPRRegBank);
-        B.setInstr(*I);
+        B.setMBB(*LoopBB);
       }
 
       unsigned OpSize = OpTy.getSizeInBits();
@@ -879,7 +882,7 @@ bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
           for (unsigned PieceIdx = 0; PieceIdx != NumPieces; ++PieceIdx)
             UnmergePieces.push_back(Unmerge.getReg(PieceIdx));
         }
-        B.setInstr(*I);
+        B.setMBB(*LoopBB);
 
         for (Register UnmergePiece : UnmergePieces) {
           Register CurrentLaneOpReg;
@@ -978,7 +981,7 @@ bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
 
   MRI.setSimpleHint(NewExec, CondReg);
 
-  B.setInsertPt(*LoopBB, LoopBB->end());
+  B.setInsertPt(*BodyBB, BodyBB->end());
 
   // Update EXEC, switch all done bits to 0 and all todo bits to 1.
   B.buildInstr(XorTermOpc)

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -5464,7 +5464,8 @@ void SIInstrInfo::legalizeGenericOperand(MachineBasicBlock &InsertMBB,
 static void
 emitLoadSRsrcFromVGPRLoop(const SIInstrInfo &TII, MachineRegisterInfo &MRI,
                           MachineBasicBlock &OrigBB, MachineBasicBlock &LoopBB,
-                          const DebugLoc &DL, MachineOperand &Rsrc) {
+                          MachineBasicBlock &BodyBB, const DebugLoc &DL,
+                          MachineOperand &Rsrc) {
   MachineFunction &MF = *OrigBB.getParent();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
@@ -5557,14 +5558,14 @@ emitLoadSRsrcFromVGPRLoop(const SIInstrInfo &TII, MachineRegisterInfo &MRI,
       .addReg(CondReg, RegState::Kill);
 
   // The original instruction is here; we insert the terminators after it.
-  I = LoopBB.end();
+  I = BodyBB.end();
 
   // Update EXEC, switch all done bits to 0 and all todo bits to 1.
-  BuildMI(LoopBB, I, DL, TII.get(XorTermOpc), Exec)
+  BuildMI(BodyBB, I, DL, TII.get(XorTermOpc), Exec)
       .addReg(Exec)
       .addReg(SaveExec);
 
-  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::SI_WATERFALL_LOOP)).addMBB(&LoopBB);
+  BuildMI(BodyBB, I, DL, TII.get(AMDGPU::SI_WATERFALL_LOOP)).addMBB(&LoopBB);
 }
 
 // Build a waterfall loop around \p MI, replacing the VGPR \p Rsrc register
@@ -5611,44 +5612,48 @@ loadSRsrcFromVGPR(const SIInstrInfo &TII, MachineInstr &MI,
   // To insert the loop we need to split the block. Move everything after this
   // point to a new block, and insert a new empty block between the two.
   MachineBasicBlock *LoopBB = MF.CreateMachineBasicBlock();
+  MachineBasicBlock *BodyBB = MF.CreateMachineBasicBlock();
   MachineBasicBlock *RemainderBB = MF.CreateMachineBasicBlock();
   MachineFunction::iterator MBBI(MBB);
   ++MBBI;
 
   MF.insert(MBBI, LoopBB);
+  MF.insert(MBBI, BodyBB);
   MF.insert(MBBI, RemainderBB);
 
-  LoopBB->addSuccessor(LoopBB);
-  LoopBB->addSuccessor(RemainderBB);
+  LoopBB->addSuccessor(BodyBB);
+  BodyBB->addSuccessor(LoopBB);
+  BodyBB->addSuccessor(RemainderBB);
 
-  // Move Begin to MI to the LoopBB, and the remainder of the block to
+  // Move Begin to MI to the BodyBB, and the remainder of the block to
   // RemainderBB.
   RemainderBB->transferSuccessorsAndUpdatePHIs(&MBB);
   RemainderBB->splice(RemainderBB->begin(), &MBB, End, MBB.end());
-  LoopBB->splice(LoopBB->begin(), &MBB, Begin, MBB.end());
+  BodyBB->splice(BodyBB->begin(), &MBB, Begin, MBB.end());
 
   MBB.addSuccessor(LoopBB);
 
   // Update dominators. We know that MBB immediately dominates LoopBB, that
-  // LoopBB immediately dominates RemainderBB, and that RemainderBB immediately
-  // dominates all of the successors transferred to it from MBB that MBB used
-  // to properly dominate.
+  // LoopBB immediately dominates BodyBB, and BodyBB immediately dominates
+  // RemainderBB. RemainderBB immediately dominates all of the successors
+  // transferred to it from MBB that MBB used to properly dominate.
   if (MDT) {
     MDT->addNewBlock(LoopBB, &MBB);
-    MDT->addNewBlock(RemainderBB, LoopBB);
+    MDT->addNewBlock(BodyBB, LoopBB);
+    MDT->addNewBlock(RemainderBB, BodyBB);
     for (auto &Succ : RemainderBB->successors()) {
       if (MDT->properlyDominates(&MBB, Succ)) {
         MDT->changeImmediateDominator(Succ, RemainderBB);
       }
     }
   }
 
-  emitLoadSRsrcFromVGPRLoop(TII, MRI, MBB, *LoopBB, DL, Rsrc);
+  emitLoadSRsrcFromVGPRLoop(TII, MRI, MBB, *LoopBB, *BodyBB, DL, Rsrc);
 
   // Restore the EXEC mask
   MachineBasicBlock::iterator First = RemainderBB->begin();
   BuildMI(*RemainderBB, First, DL, TII.get(MovExecOpc), Exec).addReg(SaveExec);
-  return LoopBB;
+  return BodyBB;
 }
 
 // Extract pointer from Rsrc and return a zero-value Rsrc replacement.

llvm/lib/Target/AMDGPU/SIOptimizeVGPRLiveRange.cpp

@@ -112,8 +112,10 @@ class SIOptimizeVGPRLiveRange : public MachineFunctionPass {
                             SmallVectorImpl<Register> &CandidateRegs) const;
 
   void collectWaterfallCandidateRegisters(
-      MachineBasicBlock *Loop,
-      SmallSetVector<Register, 16> &CandidateRegs) const;
+      MachineBasicBlock *LoopHeader, MachineBasicBlock *LoopEnd,
+      SmallSetVector<Register, 16> &CandidateRegs,
+      SmallSetVector<MachineBasicBlock *, 2> &Blocks,
+      SmallVectorImpl<MachineInstr *> &Instructions) const;
 
   void findNonPHIUsesInBlock(Register Reg, MachineBasicBlock *MBB,
                              SmallVectorImpl<MachineInstr *> &Uses) const;
@@ -131,7 +133,10 @@ class SIOptimizeVGPRLiveRange : public MachineFunctionPass {
                     MachineBasicBlock *Flow, MachineBasicBlock *Endif,
                     SmallSetVector<MachineBasicBlock *, 16> &ElseBlocks) const;
 
-  void optimizeWaterfallLiveRange(Register Reg, MachineBasicBlock *If) const;
+  void optimizeWaterfallLiveRange(
+      Register Reg, MachineBasicBlock *LoopHeader,
+      SmallSetVector<MachineBasicBlock *, 2> &LoopBlocks,
+      SmallVectorImpl<MachineInstr *> &Instructions) const;
 
   SIOptimizeVGPRLiveRange() : MachineFunctionPass(ID) {}
 
@@ -323,12 +328,30 @@ void SIOptimizeVGPRLiveRange::collectCandidateRegisters(
 /// Collect the registers used in the waterfall loop block that are defined
 /// before.
 void SIOptimizeVGPRLiveRange::collectWaterfallCandidateRegisters(
-    MachineBasicBlock *Loop,
-    SmallSetVector<Register, 16> &CandidateRegs) const {
+    MachineBasicBlock *LoopHeader, MachineBasicBlock *LoopEnd,
+    SmallSetVector<Register, 16> &CandidateRegs,
+    SmallSetVector<MachineBasicBlock *, 2> &Blocks,
+    SmallVectorImpl<MachineInstr *> &Instructions) const {
+
+  // Collect loop instructions, potentially spanning multiple blocks
+  auto *MBB = LoopHeader;
+  for (;;) {
+    Blocks.insert(MBB);
+    for (auto &MI : *MBB) {
+      if (MI.isDebugInstr())
+        continue;
+      Instructions.push_back(&MI);
+    }
+    if (MBB == LoopEnd)
+      break;
+    assert(MBB->pred_size() == 1 ||
+           (MBB == LoopHeader && MBB->pred_size() == 2));
+    assert(MBB->succ_size() == 1);
+    MBB = *MBB->succ_begin();
+  }
 
-  for (auto &MI : Loop->instrs()) {
-    if (MI.isDebugInstr())
-      continue;
+  for (auto *I : Instructions) {
+    auto &MI = *I;
 
     for (auto &MO : MI.operands()) {
       if (!MO.isReg() || !MO.getReg() || MO.isDef())
@@ -340,16 +363,17 @@ void SIOptimizeVGPRLiveRange::collectWaterfallCandidateRegisters(
         continue;
 
       if (MO.readsReg()) {
-        const MachineBasicBlock *DefMBB = MRI->getVRegDef(MOReg)->getParent();
+        MachineBasicBlock *DefMBB = MRI->getVRegDef(MOReg)->getParent();
         // Make sure the value is defined before the LOOP block
-        if (DefMBB != Loop && !CandidateRegs.contains(MOReg)) {
+        if (!Blocks.contains(DefMBB) && !CandidateRegs.contains(MOReg)) {
           // If the variable is used after the loop, the register coalescer will
           // merge the newly created register and remove the phi node again.
           // Just do nothing in that case.
           LiveVariables::VarInfo &OldVarInfo = LV->getVarInfo(MOReg);
           bool IsUsed = false;
-          for (auto *Succ : Loop->successors()) {
-            if (Succ != Loop && OldVarInfo.isLiveIn(*Succ, MOReg, *MRI)) {
+          for (auto *Succ : LoopEnd->successors()) {
+            if (!Blocks.contains(Succ) &&
+                OldVarInfo.isLiveIn(*Succ, MOReg, *MRI)) {
               IsUsed = true;
               break;
             }
@@ -513,7 +537,9 @@ void SIOptimizeVGPRLiveRange::optimizeLiveRange(
 }
 
 void SIOptimizeVGPRLiveRange::optimizeWaterfallLiveRange(
-    Register Reg, MachineBasicBlock *Loop) const {
+    Register Reg, MachineBasicBlock *LoopHeader,
+    SmallSetVector<MachineBasicBlock *, 2> &Blocks,
+    SmallVectorImpl<MachineInstr *> &Instructions) const {
   // Insert a new PHI, marking the value from the last loop iteration undef.
   LLVM_DEBUG(dbgs() << "Optimizing " << printReg(Reg, TRI) << '\n');
   const auto *RC = MRI->getRegClass(Reg);
@@ -525,15 +551,16 @@ void SIOptimizeVGPRLiveRange::optimizeWaterfallLiveRange(
   for (auto &O : make_early_inc_range(MRI->use_operands(Reg))) {
     auto *UseMI = O.getParent();
     auto *UseBlock = UseMI->getParent();
-    // Replace uses in Loop block
-    if (UseBlock == Loop)
+    // Replace uses in Loop blocks
+    if (Blocks.contains(UseBlock))
       O.setReg(NewReg);
   }
 
-  MachineInstrBuilder PHI = BuildMI(*Loop, Loop->getFirstNonPHI(), DebugLoc(),
-                                    TII->get(TargetOpcode::PHI), NewReg);
-  for (auto *Pred : Loop->predecessors()) {
-    if (Pred == Loop)
+  MachineInstrBuilder PHI =
+      BuildMI(*LoopHeader, LoopHeader->getFirstNonPHI(), DebugLoc(),
+              TII->get(TargetOpcode::PHI), NewReg);
+  for (auto *Pred : LoopHeader->predecessors()) {
+    if (Blocks.contains(Pred))
       PHI.addReg(UndefReg, RegState::Undef).addMBB(Pred);
     else
       PHI.addReg(Reg).addMBB(Pred);
@@ -542,21 +569,36 @@ void SIOptimizeVGPRLiveRange::optimizeWaterfallLiveRange(
   LiveVariables::VarInfo &NewVarInfo = LV->getVarInfo(NewReg);
   LiveVariables::VarInfo &OldVarInfo = LV->getVarInfo(Reg);
 
-  // collectWaterfallCandidateRegisters only collects registers that are dead
-  // after the loop. So we know that the old reg is not live throughout the
-  // whole block anymore.
-  OldVarInfo.AliveBlocks.reset(Loop->getNumber());
-
-  // Mark the last use as kill
-  for (auto &MI : reverse(Loop->instrs())) {
-    if (MI.readsRegister(NewReg, TRI)) {
-      MI.addRegisterKilled(NewReg, TRI);
-      NewVarInfo.Kills.push_back(&MI);
+  // Find last use and mark as kill
+  MachineInstr *Kill = nullptr;
+  for (auto *MI : reverse(Instructions)) {
+    if (MI->readsRegister(NewReg, TRI)) {
+      MI->addRegisterKilled(NewReg, TRI);
+      NewVarInfo.Kills.push_back(MI);
+      Kill = MI;
       break;
     }
   }
-  assert(!NewVarInfo.Kills.empty() &&
-         "Failed to find last usage of register in loop");
+  assert(Kill && "Failed to find last usage of register in loop");
+
+  MachineBasicBlock *KillBlock = Kill->getParent();
+  bool PostKillBlock = false;
+  for (auto *Block : Blocks) {
+    auto BBNum = Block->getNumber();
+
+    // collectWaterfallCandidateRegisters only collects registers that are dead
+    // after the loop. So we know that the old reg is no longer live throughout
+    // the waterfall loop.
+    OldVarInfo.AliveBlocks.reset(BBNum);
+
+    // The new register is live up to (and including) the block that kills it.
+    PostKillBlock |= (Block == KillBlock);
+    if (PostKillBlock) {
+      NewVarInfo.AliveBlocks.reset(BBNum);
+    } else if (Block != LoopHeader) {
+      NewVarInfo.AliveBlocks.set(BBNum);
+    }
+  }
 }
 
 char SIOptimizeVGPRLiveRange::ID = 0;
@@ -620,15 +662,22 @@ bool SIOptimizeVGPRLiveRange::runOnMachineFunction(MachineFunction &MF) {
         for (auto Reg : CandidateRegs)
           optimizeLiveRange(Reg, &MBB, IfTarget, Endif, ElseBlocks);
       } else if (MI.getOpcode() == AMDGPU::SI_WATERFALL_LOOP) {
+        auto *LoopHeader = MI.getOperand(0).getMBB();
+        auto *LoopEnd = &MBB;
+
         LLVM_DEBUG(dbgs() << "Checking Waterfall loop: "
-                          << printMBBReference(MBB) << '\n');
+                          << printMBBReference(*LoopHeader) << '\n');
 
         SmallSetVector<Register, 16> CandidateRegs;
-        collectWaterfallCandidateRegisters(&MBB, CandidateRegs);
+        SmallVector<MachineInstr *, 16> Instructions;
+        SmallSetVector<MachineBasicBlock *, 2> Blocks;
+
+        collectWaterfallCandidateRegisters(LoopHeader, LoopEnd, CandidateRegs,
+                                           Blocks, Instructions);
         MadeChange |= !CandidateRegs.empty();
         // Now we are safe to optimize.
         for (auto Reg : CandidateRegs)
-          optimizeWaterfallLiveRange(Reg, &MBB);
+          optimizeWaterfallLiveRange(Reg, LoopHeader, Blocks, Instructions);
       }
     }
   }