AMDGPU: Implement waterfall loop for MIMG instructions with 256-bit SRsrc

[nhaehnle]

Could you make this generic over the VGPR register class instead? That code duplication is rather annoying.

[arsenm]

Desperately needs tests

[cfang]

Add waterloop tests, and update existing tests.

[cfang]

In D82603#2119746, @nhaehnle wrote:

Could you make this generic over the VGPR register class instead? That code duplication is rather annoying.

I agree it does not look nice with the duplication. However, the existing implementation is easy to follow,
and I could not figure out a way that can easily generalize over vgpr class without make the code messy.

We will continue to work on to come out with a generalized solution. But in the mean time, we need
the fix for a correctness issue which has been bothering us for a long time.

[arsenm]

In D82603#2173645, @cfang wrote:

In D82603#2119746, @nhaehnle wrote:

Could you make this generic over the VGPR register class instead? That code duplication is rather annoying.

I agree it does not look nice with the duplication. However, the existing implementation is easy to follow,
and I could not figure out a way that can easily generalize over vgpr class without make the code messy.

We will continue to work on to come out with a generalized solution. But in the mean time, we need
the fix for a correctness issue which has been bothering us for a long time.

GlobalISel already has a generalized implementation, it just relies on generic registers. I'd rather copy something from there and dirty it up with register classes

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

4986

llvm_unreachable

[cfang]

Update based on the reviewers' request to generate the waterfall loop for resource registers of
arbitrary sizes.

Update a few LIT tests, mainly because of the changes in instruction order, or additional move instructions under -O0.

[cfang]

A minor change to use the register directly in stead of getting it from the instruction.

NOTE: this is based on the version we have generalized the resource register size based on the implementation from GISel.

[cfang]

Ping. Is there any further suggestions to move this patch ahead? Thanks.

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

5075

What about the SGPR offset? I guess this was broken before anyway

[cfang]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

5075

Do you mean we will have to legalize the case with SGPR Offset? If you can come up with a test case with SGPR offset, we can fix the broken, in a separate patch of course. Thanks.

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

     ImpDef = Def && Def->isImplicitDef();
   }
   if (!RI.isSGPRClass(DstRC) && !Copy->readsRegister(AMDGPU::EXEC, &RI) &&
       !ImpDef)
     Copy->addOperand(MachineOperand::CreateReg(AMDGPU::EXEC, false, true));
 }
 
 // Emit the actual waterfall loop, executing the wrapped instruction for each
-// unique value of \p Rsrc across all lanes. In the best case we execute 1
+// unique value of \p Rsrc (4 dwords) across all lanes. In the best case we execute 1
 // iteration, in the worst case we execute 64 (once per lane).
 static void
-emitLoadSRsrcFromVGPRLoop(const SIInstrInfo &TII, MachineRegisterInfo &MRI,
+emitLoadSRsrcFromVGPR128Loop(const SIInstrInfo &TII, MachineRegisterInfo &MRI,
                           MachineBasicBlock &OrigBB, MachineBasicBlock &LoopBB,
                           const DebugLoc &DL, MachineOperand &Rsrc) {
   MachineFunction &MF = *OrigBB.getParent();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
   unsigned Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
   unsigned SaveExecOpc =
       ST.isWave32() ? AMDGPU::S_AND_SAVEEXEC_B32 : AMDGPU::S_AND_SAVEEXEC_B64;

 
   // Update EXEC, switch all done bits to 0 and all todo bits to 1.
   BuildMI(LoopBB, I, DL, TII.get(XorTermOpc), Exec)
       .addReg(Exec)
       .addReg(SaveExec);
   BuildMI(LoopBB, I, DL, TII.get(AMDGPU::S_CBRANCH_EXECNZ)).addMBB(&LoopBB);
 }
 
+
+// Emit the actual waterfall loop, executing the wrapped instruction for each
+// unique value of \p Rsrc (8 dwords) across all lanes.
+static void
+emitLoadSRsrcFromVGPR256Loop(const SIInstrInfo &TII, MachineRegisterInfo &MRI,
+                          MachineBasicBlock &OrigBB, MachineBasicBlock &LoopBB,
+                          const DebugLoc &DL, MachineOperand &Rsrc) {
+  MachineFunction &MF = *OrigBB.getParent();
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  const SIRegisterInfo *TRI = ST.getRegisterInfo();
+  unsigned Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
+  unsigned SaveExecOpc =
+      ST.isWave32() ? AMDGPU::S_AND_SAVEEXEC_B32 : AMDGPU::S_AND_SAVEEXEC_B64;
+  unsigned XorTermOpc =
+      ST.isWave32() ? AMDGPU::S_XOR_B32_term : AMDGPU::S_XOR_B64_term;
+  unsigned AndOpc =
+      ST.isWave32() ? AMDGPU::S_AND_B32 : AMDGPU::S_AND_B64;
+  const auto *BoolXExecRC = TRI->getRegClass(AMDGPU::SReg_1_XEXECRegClassID);
+
+  MachineBasicBlock::iterator I = LoopBB.begin();
+
+  Register VRsrc = Rsrc.getReg();
+  unsigned VRsrcUndef = getUndefRegState(Rsrc.isUndef());
+
+  Register SaveExec = MRI.createVirtualRegister(BoolXExecRC);
+  Register CondReg0 = MRI.createVirtualRegister(BoolXExecRC);
+  Register CondReg1 = MRI.createVirtualRegister(BoolXExecRC);
+  Register CondReg2 = MRI.createVirtualRegister(BoolXExecRC);
+  Register CondReg3 = MRI.createVirtualRegister(BoolXExecRC);
+  Register AndCond0 = MRI.createVirtualRegister(BoolXExecRC);
+  Register AndCond1 = MRI.createVirtualRegister(BoolXExecRC);
+  Register AndCond = MRI.createVirtualRegister(BoolXExecRC);
+  Register SRsrcSub0 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub1 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub2 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub3 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub4 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub5 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub6 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrcSub7 = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+  Register SRsrc = MRI.createVirtualRegister(&AMDGPU::SGPR_256RegClass);
+
+  // Beginning of the loop, read the next Rsrc variant.
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub0)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub0);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub1)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub1);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub2)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub2);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub3)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub3);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub4)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub4);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub5)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub5);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub6)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub6);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_READFIRSTLANE_B32), SRsrcSub7)
+      .addReg(VRsrc, VRsrcUndef, AMDGPU::sub7);
+
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::REG_SEQUENCE), SRsrc)
+      .addReg(SRsrcSub0)
+      .addImm(AMDGPU::sub0)
+      .addReg(SRsrcSub1)
+      .addImm(AMDGPU::sub1)
+      .addReg(SRsrcSub2)
+      .addImm(AMDGPU::sub2)
+      .addReg(SRsrcSub3)
+      .addImm(AMDGPU::sub3)
+      .addReg(SRsrcSub4)
+      .addImm(AMDGPU::sub4)
+      .addReg(SRsrcSub5)
+      .addImm(AMDGPU::sub5)
+      .addReg(SRsrcSub6)
+      .addImm(AMDGPU::sub6)
+      .addReg(SRsrcSub7)
+      .addImm(AMDGPU::sub7);
+
+  // Update Rsrc operand to use the SGPR Rsrc.
+  Rsrc.setReg(SRsrc);
+  Rsrc.setIsKill(true);
+
+  // Identify all lanes with identical Rsrc operands in their VGPRs.
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_CMP_EQ_U64_e64), CondReg0)
+      .addReg(SRsrc, 0, AMDGPU::sub0_sub1)
+      .addReg(VRsrc, 0, AMDGPU::sub0_sub1);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_CMP_EQ_U64_e64), CondReg1)
+      .addReg(SRsrc, 0, AMDGPU::sub2_sub3)
+      .addReg(VRsrc, 0, AMDGPU::sub2_sub3);
+  BuildMI(LoopBB, I, DL, TII.get(AndOpc), AndCond0)
+      .addReg(CondReg0)
+      .addReg(CondReg1);
+
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_CMP_EQ_U64_e64), CondReg2)
+      .addReg(SRsrc, 0, AMDGPU::sub4_sub5)
+      .addReg(VRsrc, 0, AMDGPU::sub4_sub5);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_CMP_EQ_U64_e64), CondReg3)
+      .addReg(SRsrc, 0, AMDGPU::sub6_sub7)
+      .addReg(VRsrc, 0, AMDGPU::sub6_sub7);
+  BuildMI(LoopBB, I, DL, TII.get(AndOpc), AndCond1)
+      .addReg(CondReg2)
+      .addReg(CondReg3);
+
+  BuildMI(LoopBB, I, DL, TII.get(AndOpc), AndCond)
+      .addReg(AndCond0)
+      .addReg(AndCond1);
+
+  MRI.setSimpleHint(SaveExec, AndCond);
+
+  // Update EXEC to matching lanes, saving original to SaveExec.
+  BuildMI(LoopBB, I, DL, TII.get(SaveExecOpc), SaveExec)
+      .addReg(AndCond, RegState::Kill);
+
+  // The original instruction is here; we insert the terminators after it.
+  I = LoopBB.end();
+
+  // Update EXEC, switch all done bits to 0 and all todo bits to 1.
+  BuildMI(LoopBB, I, DL, TII.get(XorTermOpc), Exec)
+      .addReg(Exec)
+      .addReg(SaveExec);
+  BuildMI(LoopBB, I, DL, TII.get(AMDGPU::S_CBRANCH_EXECNZ)).addMBB(&LoopBB);
+}
+
 // Build a waterfall loop around \p MI, replacing the VGPR \p Rsrc register
 // with SGPRs by iterating over all unique values across all lanes.
 static void loadSRsrcFromVGPR(const SIInstrInfo &TII, MachineInstr &MI,
                               MachineOperand &Rsrc, MachineDominatorTree *MDT) {
   MachineBasicBlock &MBB = *MI.getParent();
   MachineFunction &MF = *MBB.getParent();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();

     MDT->addNewBlock(RemainderBB, LoopBB);
     for (auto &Succ : RemainderBB->successors()) {
       if (MDT->properlyDominates(&MBB, Succ)) {
         MDT->changeImmediateDominator(Succ, RemainderBB);
       }
     }
   }
 
-  emitLoadSRsrcFromVGPRLoop(TII, MRI, MBB, *LoopBB, DL, Rsrc);
+  // Emit waterfall loop based on Rsrc size.
+  const TargetRegisterClass *OpRC = MRI.getRegClass(Rsrc.getReg());
+  if (OpRC == &AMDGPU::VReg_128RegClass)
+    emitLoadSRsrcFromVGPR128Loop(TII, MRI, MBB, *LoopBB, DL, Rsrc);
+  else if (OpRC == &AMDGPU::VReg_256RegClass)
+    emitLoadSRsrcFromVGPR256Loop(TII, MRI, MBB, *LoopBB, DL, Rsrc);
+  else // Not implemented yet.
+    assert(false && "waterfall loop not imlemented for Rsrc size");
 
   // Restore the EXEC mask
   MachineBasicBlock::iterator First = RemainderBB->begin();
   BuildMI(*RemainderBB, First, DL, TII.get(MovExecOpc), Exec).addReg(SaveExec);
 }
 
 // Extract pointer from Rsrc and return a zero-value Rsrc replacement.
 static std::tuple<unsigned, unsigned>
 extractRsrcPtr(const SIInstrInfo &TII, MachineInstr &MI, MachineOperand &Rsrc) {
   MachineBasicBlock &MBB = *MI.getParent();
   MachineFunction &MF = *MBB.getParent();
   MachineRegisterInfo &MRI = MF.getRegInfo();
 
   // Extract the ptr from the resource descriptor.
   unsigned RsrcPtr =
       TII.buildExtractSubReg(MI, MRI, Rsrc, &AMDGPU::VReg_128RegClass,
                              AMDGPU::sub0_sub1, &AMDGPU::VReg_64RegClass);
 
   // Create an empty resource descriptor
   Register Zero64 = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);

     return;
   }
 
   // Legalize SI_INIT_M0
   if (MI.getOpcode() == AMDGPU::SI_INIT_M0) {
     MachineOperand &Src = MI.getOperand(0);
     if (Src.isReg() && RI.hasVectorRegisters(MRI.getRegClass(Src.getReg())))
       Src.setReg(readlaneVGPRToSGPR(Src.getReg(), MI, MRI));
     return;
   }
 
   // Legalize MIMG and MUBUF/MTBUF for shaders.
   //
   // Shaders only generate MUBUF/MTBUF instructions via intrinsics or via
   // scratch memory access. In both cases, the legalization never involves
   // conversion to the addr64 form.
   if (isMIMG(MI) ||
       (AMDGPU::isShader(MF.getFunction().getCallingConv()) &&
        (isMUBUF(MI) || isMTBUF(MI)))) {
     MachineOperand *SRsrc = getNamedOperand(MI, AMDGPU::OpName::srsrc);
-    if (SRsrc && !RI.isSGPRClass(MRI.getRegClass(SRsrc->getReg()))) {
-      unsigned SGPR = readlaneVGPRToSGPR(SRsrc->getReg(), MI, MRI);
-      SRsrc->setReg(SGPR);
-    }
+    if (SRsrc && !RI.isSGPRClass(MRI.getRegClass(SRsrc->getReg())))
+      loadSRsrcFromVGPR(*this, MI, *SRsrc, MDT);
 
     MachineOperand *SSamp = getNamedOperand(MI, AMDGPU::OpName::ssamp);
-    if (SSamp && !RI.isSGPRClass(MRI.getRegClass(SSamp->getReg()))) {
-      unsigned SGPR = readlaneVGPRToSGPR(SSamp->getReg(), MI, MRI);
-      SSamp->setReg(SGPR);
-    }
+    if (SSamp && !RI.isSGPRClass(MRI.getRegClass(SSamp->getReg())))
+      loadSRsrcFromVGPR(*this, MI, *SSamp, MDT);
+
     return;
   }
 
   // Legalize MUBUF* instructions.
   int RsrcIdx =
       AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::srsrc);
   if (RsrcIdx != -1) {
     // We have an MUBUF instruction