[llvm-branch-commits] [llvm] AMDGPU/GlobalISel: Add waterfall lowering in regbanklegalize (PR #142790)

[Pierre van Houtryve via llvm-branch-commits]

================
@@ -57,6 +57,226 @@ void RegBankLegalizeHelper::findRuleAndApplyMapping(MachineInstr &MI) {
   lower(MI, Mapping, WaterfallSgprs);
 }
 
+bool RegBankLegalizeHelper::executeInWaterfallLoop(
+    MachineIRBuilder &B, iterator_range<MachineBasicBlock::iterator> Range,
+    SmallSet<Register, 4> &SGPROperandRegs) {
+  // Track use registers which have already been expanded with a readfirstlane
+  // sequence. This may have multiple uses if moving a sequence.
+  DenseMap<Register, Register> WaterfalledRegMap;
+
+  MachineBasicBlock &MBB = B.getMBB();
+  MachineFunction &MF = B.getMF();
+
+  const SIRegisterInfo *TRI = ST.getRegisterInfo();
+  const TargetRegisterClass *WaveRC = TRI->getWaveMaskRegClass();
+  unsigned MovExecOpc, MovExecTermOpc, XorTermOpc, AndSaveExecOpc, ExecReg;
+  if (ST.isWave32()) {
+    MovExecOpc = AMDGPU::S_MOV_B32;
+    MovExecTermOpc = AMDGPU::S_MOV_B32_term;
+    XorTermOpc = AMDGPU::S_XOR_B32_term;
+    AndSaveExecOpc = AMDGPU::S_AND_SAVEEXEC_B32;
+    ExecReg = AMDGPU::EXEC_LO;
+  } else {
+    MovExecOpc = AMDGPU::S_MOV_B64;
+    MovExecTermOpc = AMDGPU::S_MOV_B64_term;
+    XorTermOpc = AMDGPU::S_XOR_B64_term;
+    AndSaveExecOpc = AMDGPU::S_AND_SAVEEXEC_B64;
+    ExecReg = AMDGPU::EXEC;
+  }
+
+#ifndef NDEBUG
+  const int OrigRangeSize = std::distance(Range.begin(), Range.end());
+#endif
+
+  MachineRegisterInfo &MRI = *B.getMRI();
+  Register SaveExecReg = MRI.createVirtualRegister(WaveRC);
+  Register InitSaveExecReg = MRI.createVirtualRegister(WaveRC);
+
+  // Don't bother using generic instructions/registers for the exec mask.
+  B.buildInstr(TargetOpcode::IMPLICIT_DEF).addDef(InitSaveExecReg);
+
+  Register SavedExec = MRI.createVirtualRegister(WaveRC);
+
+  // 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 *RestoreExecBB = MF.CreateMachineBasicBlock();
+  MachineBasicBlock *RemainderBB = 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(BodyBB);
+  BodyBB->addSuccessor(RestoreExecBB);
+  BodyBB->addSuccessor(LoopBB);
+
+  // Move the rest of the block into a new block.
+  RemainderBB->transferSuccessorsAndUpdatePHIs(&MBB);
+  RemainderBB->splice(RemainderBB->begin(), &MBB, Range.end(), MBB.end());
+
+  MBB.addSuccessor(LoopBB);
+  RestoreExecBB->addSuccessor(RemainderBB);
+
+  B.setInsertPt(*LoopBB, LoopBB->end());
+
+  // +-MBB:------------+
+  // | ...             |
+  // | %0 = G_INST_1   |
+  // | %Dst = MI %Vgpr |
+  // | %1 = G_INST_2   |
+  // | ...             |
+  // +-----------------+
+  // ->
+  // +-MBB-------------------------------+
+  // | ...                               |
+  // | %0 = G_INST_1                     |
+  // | %SaveExecReg = S_MOV_B32 $exec_lo |
+  // +----------------|------------------+
+  //                  |                         /------------------------------|
+  //                  V                        V                               |
+  // +-LoopBB---------------------------------------------------------------+  |
+  // | %CurrentLaneReg:sgpr(s32) = READFIRSTLANE %Vgpr                      |  |
+  // |   instead of executing for each lane, see if other lanes had         |  |
+  // |   same value for %Vgpr and execute for them also.                    |  |
+  // | %CondReg:vcc(s1) = G_ICMP eq %CurrentLaneReg, %Vgpr                  |  |
+  // | %CondRegLM:sreg_32 = ballot %CondReg // copy vcc to sreg32 lane mask |  |
+  // | %SavedExec = S_AND_SAVEEXEC_B32 %CondRegLM                           |  |
+  // |   exec is active for lanes with the same "CurrentLane value" in Vgpr |  |
+  // +----------------|-----------------------------------------------------+  |
+  //                  V                                                        |
+  // +-BodyBB------------------------------------------------------------+     |
+  // | %Dst = MI %CurrentLaneReg:sgpr(s32)                               |     |
+  // |   executed only for active lanes and written to Dst               |     |
+  // | $exec = S_XOR_B32 $exec, %SavedExec                               |     |
+  // |   set active lanes to 0 in SavedExec, lanes that did not write to |     |
+  // |   Dst yet, and set this as new exec (for READFIRSTLANE and ICMP)  |     |
+  // | SI_WATERFALL_LOOP LoopBB                                          |-----|
+  // +----------------|--------------------------------------------------+
+  //                  V
+  // +-RestoreExecBB--------------------------+
+  // | $exec_lo = S_MOV_B32_term %SaveExecReg |
+  // +----------------|-----------------------+
+  //                  V
+  // +-RemainderBB:----------------------+
+  // | %1 = G_INST_2                     |
+  // | ...                               |
+  // +---------------------------------- +
+
+  // 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.
+  BodyBB->splice(BodyBB->end(), &MBB, Range.begin(), MBB.end());
+
+  // Figure out the iterator range after splicing the instructions.
+  MachineBasicBlock::iterator NewBegin = Range.begin()->getIterator();
+  auto NewEnd = BodyBB->end();
+  assert(std::distance(NewBegin, NewEnd) == OrigRangeSize);
+
+  B.setMBB(*LoopBB);
+  Register CondReg;
+
+  for (MachineInstr &MI : make_range(NewBegin, NewEnd)) {
+    for (MachineOperand &Op : MI.all_uses()) {
+      Register OldReg = Op.getReg();
+      if (!SGPROperandRegs.count(OldReg))
+        continue;
+
+      // See if we already processed this register in another instruction in
+      // the sequence.
+      auto OldVal = WaterfalledRegMap.find(OldReg);
+      if (OldVal != WaterfalledRegMap.end()) {
+        Op.setReg(OldVal->second);
+        continue;
+      }
+
+      Register OpReg = Op.getReg();
+      LLT OpTy = MRI.getType(OpReg);
+
+      // TODO: support for agpr
+      assert(MRI.getRegBank(OpReg) == VgprRB);
+      Register CurrentLaneReg = MRI.createVirtualRegister({SgprRB, OpTy});
+      buildReadFirstLane(B, CurrentLaneReg, OpReg, RBI);
+
+      // Build the comparison(s), CurrentLaneReg == OpReg.
+      unsigned OpSize = OpTy.getSizeInBits();
+      bool Is64 = OpSize % 64 == 0;
+      unsigned PartSize = Is64 ? 64 : 32;
+      LLT PartTy = LLT::scalar(PartSize);
+      unsigned NumParts = OpSize / PartSize;
+      SmallVector<Register, 8> OpParts;
+      SmallVector<Register, 8> CurrentLaneParts;
+
+      if (NumParts == 1) {
+        OpParts.push_back(OpReg);
+        CurrentLaneParts.push_back(CurrentLaneReg);
+      } else {
+        auto UnmergeOp = B.buildUnmerge({VgprRB, PartTy}, OpReg);
+        auto UnmergeCurrLane = B.buildUnmerge({SgprRB, PartTy}, CurrentLaneReg);
+        for (unsigned i = 0; i < NumParts; ++i) {
+          OpParts.push_back(UnmergeOp.getReg(i));
+          CurrentLaneParts.push_back(UnmergeCurrLane.getReg(i));
+        }
+      }
+
+      for (unsigned i = 0; i < NumParts; ++i) {
+        Register CmpReg = MRI.createVirtualRegister(VccRB_S1);
+        B.buildICmp(CmpInst::ICMP_EQ, CmpReg, CurrentLaneParts[i], OpParts[i]);
+
+        if (!CondReg) {
+          CondReg = CmpReg;
+        } else {
+          CondReg = B.buildAnd(VccRB_S1, CondReg, CmpReg).getReg(0);
+        }
----------------
Pierre-vh wrote:

```suggestion
        if (!CondReg) 
          CondReg = CmpReg;
        else 
          CondReg = B.buildAnd(VccRB_S1, CondReg, CmpReg).getReg(0);
    
```

https://github.com/llvm/llvm-project/pull/142790