[AMDGPU] Insert waterfall loops for divergent calls

[Flakebi]

Summary

Extend loadSRsrcFromVGPR to allow moving a range of instructions into
the loop. The call instruction is preceded by copies into physical
registers which should be part of the waterfall loop, as the registers
can be overwritten by the call.

[madhur13490]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

4897

Dump() is not required.

5125

Should this block be executed for AGPRs too? If this is meant only for VGPRs then !SGPR is not correct.

[Flakebi]

Remove debug dumps

[Flakebi]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

4897

Thanks, I forgot to remove them.

5125

I don’t know how AGPRs work and documentation seems to be scarce, the check for calls is the same as for image operations above.
It seems like AGPRs can be copied to VGPRs but not to SGPRs (https://reviews.llvm.org/D68358#change-cnuV0NuhUhzL), so I think calling a function pointer that is stored in AGPRs should copy them to VGPRs and insert a waterfall loop.

[madhur13490]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

5125

My main concern is that this check is over relaxing and allows all non-SGPR classes. This also includes any future register class we may add. AGPR is just an example. What about enabling it just for VGPR for now and put a TODO?

[nhaehnle]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

5125

Dest is the destination we're calling, i.e. the function pointer. The point of the logic, AFAIU, is that if the function pointer is non-uniform, we need to do something about that. AGPRs are non-uniform... so I'd say !isSGPR is correct.

[Flakebi]

Fix return value handling, also need to copy COPYs following the call.

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

for (unsigned Idx = 0; Idx < NumSubRegs; Idx += 2) {
     Register CurReg = MRI.createVirtualRegister(&AMDGPU::SGPR_64RegClass);
     BuildMI(LoopBB, I, DL, TII.get(AMDGPU::REG_SEQUENCE), CurReg)
             .addReg(CurRegLo)
             .addImm(AMDGPU::sub0)
             .addReg(CurRegHi)
             .addImm(AMDGPU::sub1);
 
     Register NewCondReg = MRI.createVirtualRegister(BoolXExecRC);
+    auto Cmp =
         BuildMI(LoopBB, I, DL, TII.get(AMDGPU::V_CMP_EQ_U64_e64), NewCondReg)
-           .addReg(CurReg)
-           .addReg(VRsrc, VRsrcUndef, TRI->getSubRegFromChannel(Idx, 2));
+            .addReg(CurReg);
+    if (NumSubRegs <= 2)
+      Cmp.addReg(VRsrc);
+    else
+      Cmp.addReg(VRsrc, VRsrcUndef, TRI->getSubRegFromChannel(Idx, 2));
 
     // Combine the comparision results with AND.
     if (CondReg == AMDGPU::NoRegister) // First.
       CondReg = NewCondReg;
     else { // If not the first, we create an AND.
       Register AndReg = MRI.createVirtualRegister(BoolXExecRC);
       BuildMI(LoopBB, I, DL, TII.get(AndOpc), AndReg)
               .addReg(CondReg)
BuildMI(LoopBB, I, DL, TII.get(XorTermOpc), 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) {
+                              MachineOperand &Rsrc, MachineDominatorTree *MDT,
+                              MachineBasicBlock::iterator Begin = nullptr) {
   MachineBasicBlock &MBB = *MI.getParent();
   MachineFunction &MF = *MBB.getParent();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
   MachineRegisterInfo &MRI = MF.getRegInfo();
-  MachineBasicBlock::iterator I(&MI);
+  if (!Begin.isValid())
+    Begin = &MI;
   const DebugLoc &DL = MI.getDebugLoc();
   unsigned Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
   unsigned MovExecOpc = ST.isWave32() ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64;
   const auto *BoolXExecRC = TRI->getRegClass(AMDGPU::SReg_1_XEXECRegClassID);
 
   Register SaveExec = MRI.createVirtualRegister(BoolXExecRC);
 
   // Save the EXEC mask
-  BuildMI(MBB, I, DL, TII.get(MovExecOpc), SaveExec).addReg(Exec);
+  BuildMI(MBB, Begin, DL, TII.get(MovExecOpc), SaveExec).addReg(Exec);
 
   // Killed uses in the instruction we are waterfalling around will be
   // incorrect due to the added control-flow.
   for (auto &MO : MI.uses()) {
     if (MO.isReg() && MO.isUse()) {
       MRI.clearKillFlags(MO.getReg());
     }
   }
 
   // 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 *RemainderBB = MF.CreateMachineBasicBlock();
   MachineFunction::iterator MBBI(MBB);
   ++MBBI;
 
   MF.insert(MBBI, LoopBB);
   MF.insert(MBBI, RemainderBB);
 
   LoopBB->addSuccessor(LoopBB);
   LoopBB->addSuccessor(RemainderBB);
 
-  // Move MI to the LoopBB, and the remainder of the block to RemainderBB.
-  MachineBasicBlock::iterator J = I++;
+  // Move Begin to MI to the LoopBB, and the remainder of the block to
+  // RemainderBB.
+  MachineBasicBlock::iterator I(&MI);
+  ++I;
   RemainderBB->transferSuccessorsAndUpdatePHIs(&MBB);
   RemainderBB->splice(RemainderBB->begin(), &MBB, I, MBB.end());
-  LoopBB->splice(LoopBB->begin(), &MBB, J);
+  LoopBB->splice(LoopBB->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.
   if (MDT) {
if (isMIMG(MI) ||
 
     MachineOperand *SSamp = getNamedOperand(MI, AMDGPU::OpName::ssamp);
     if (SSamp && !RI.isSGPRClass(MRI.getRegClass(SSamp->getReg())))
       loadSRsrcFromVGPR(*this, MI, *SSamp, MDT);
 
     return;
   }
 
+  // Legalize SI_CALL
+  if (MI.getOpcode() == AMDGPU::SI_CALL_ISEL) {
+    MachineOperand *Dest = &MI.getOperand(0);
+    if (!RI.isSGPRClass(MRI.getRegClass(Dest->getReg()))) {
+      // Also move the copies to physical registers into the loop block
+      MachineBasicBlock &MBB = *MI.getParent();
+      MachineBasicBlock::iterator SplitPoint(&MI);
+      --SplitPoint;
+      while (SplitPoint != MBB.begin() && SplitPoint->isCopy() &&
+             SplitPoint->getOperand(0).isReg() &&
+             MI.readsRegister(SplitPoint->getOperand(0).getReg()))
+        --SplitPoint;
+      if (!SplitPoint->isCopy())
+        ++SplitPoint;
+
+      loadSRsrcFromVGPR(*this, MI, *Dest, MDT, SplitPoint);
+    }
+  }
+
   // Legalize MUBUF* instructions.
   int RsrcIdx =
       AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::srsrc);
   if (RsrcIdx != -1) {
     // We have an MUBUF instruction
     MachineOperand *Rsrc = &MI.getOperand(RsrcIdx);
     unsigned RsrcRC = get(MI.getOpcode()).OpInfo[RsrcIdx].RegClass;
     if (RI.getCommonSubClass(MRI.getRegClass(Rsrc->getReg()),

llvm/test/CodeGen/AMDGPU/indirect-call.ll

 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    s_swappc_b64 s[30:31], s[18:19]
 ; GCN-NEXT:    s_endpgm
   %fptr = load void(i32)*, void(i32)* addrspace(4)* @gv.fptr1
   call void %fptr(i32 123)
   ret void
 }
 
-; FIXME
-; define void @test_indirect_call_vgpr_ptr(void()* %fptr) {
-;   call void %fptr()
-;   ret void
-; }
+define void @test_indirect_call_vgpr_ptr(void()* %fptr) {
+; GCN-LABEL: test_indirect_call_vgpr_ptr:
+; GCN:       ; %bb.0:
+; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GCN-NEXT:    s_or_saveexec_b64 s[16:17], -1
+; GCN-NEXT:    buffer_store_dword v43, off, s[0:3], s32 offset:12 ; 4-byte Folded Spill
+; GCN-NEXT:    s_mov_b64 exec, s[16:17]
+; GCN-NEXT:    v_writelane_b32 v43, s33, 17
+; GCN-NEXT:    s_mov_b32 s33, s32
+; GCN-NEXT:    s_add_u32 s32, s32, 0x800
+; GCN-NEXT:    buffer_store_dword v40, off, s[0:3], s33 offset:8 ; 4-byte Folded Spill
+; GCN-NEXT:    buffer_store_dword v41, off, s[0:3], s33 offset:4 ; 4-byte Folded Spill
+; GCN-NEXT:    buffer_store_dword v42, off, s[0:3], s33 ; 4-byte Folded Spill
+; GCN-NEXT:    v_writelane_b32 v43, s34, 0
+; GCN-NEXT:    v_writelane_b32 v43, s35, 1
+; GCN-NEXT:    v_writelane_b32 v43, s36, 2
+; GCN-NEXT:    v_writelane_b32 v43, s38, 3
+; GCN-NEXT:    v_writelane_b32 v43, s39, 4
+; GCN-NEXT:    v_writelane_b32 v43, s40, 5
+; GCN-NEXT:    v_writelane_b32 v43, s41, 6
+; GCN-NEXT:    v_writelane_b32 v43, s42, 7
+; GCN-NEXT:    v_writelane_b32 v43, s43, 8
+; GCN-NEXT:    v_writelane_b32 v43, s44, 9
+; GCN-NEXT:    v_writelane_b32 v43, s45, 10
+; GCN-NEXT:    v_writelane_b32 v43, s46, 11
+; GCN-NEXT:    v_writelane_b32 v43, s47, 12
+; GCN-NEXT:    v_writelane_b32 v43, s48, 13
+; GCN-NEXT:    v_writelane_b32 v43, s49, 14
+; GCN-NEXT:    v_writelane_b32 v43, s30, 15
+; GCN-NEXT:    v_writelane_b32 v43, s31, 16
+; GCN-NEXT:    v_mov_b32_e32 v40, v31
+; GCN-NEXT:    s_mov_b32 s34, s14
+; GCN-NEXT:    s_mov_b32 s35, s13
+; GCN-NEXT:    s_mov_b32 s36, s12
+; GCN-NEXT:    s_mov_b64 s[38:39], s[10:11]
+; GCN-NEXT:    s_mov_b64 s[40:41], s[8:9]
+; GCN-NEXT:    s_mov_b64 s[42:43], s[6:7]
+; GCN-NEXT:    s_mov_b64 s[44:45], s[4:5]
+; GCN-NEXT:    v_mov_b32_e32 v42, v1
+; GCN-NEXT:    v_mov_b32_e32 v41, v0
+; GCN-NEXT:    s_mov_b64 s[46:47], exec
+; GCN-NEXT:  BB2_1: ; =>This Inner Loop Header: Depth=1
+; GCN-NEXT:    v_readfirstlane_b32 s16, v41
+; GCN-NEXT:    v_readfirstlane_b32 s17, v42
+; GCN-NEXT:    v_cmp_eq_u64_e32 vcc, s[16:17], v[41:42]
+; GCN-NEXT:    s_and_saveexec_b64 s[48:49], vcc
+; GCN-NEXT:    s_mov_b64 s[4:5], s[44:45]
+; GCN-NEXT:    s_mov_b64 s[6:7], s[42:43]
+; GCN-NEXT:    s_mov_b64 s[8:9], s[40:41]
+; GCN-NEXT:    s_mov_b64 s[10:11], s[38:39]
+; GCN-NEXT:    s_mov_b32 s12, s36
+; GCN-NEXT:    s_mov_b32 s13, s35
+; GCN-NEXT:    s_mov_b32 s14, s34
+; GCN-NEXT:    v_mov_b32_e32 v31, v40
+; GCN-NEXT:    s_swappc_b64 s[30:31], s[16:17]
+; GCN-NEXT:    s_xor_b64 exec, exec, s[48:49]
+; GCN-NEXT:    s_cbranch_execnz BB2_1
+; GCN-NEXT:  ; %bb.2:
+; GCN-NEXT:    s_mov_b64 exec, s[46:47]
+; GCN-NEXT:    v_readlane_b32 s4, v43, 15
+; GCN-NEXT:    v_readlane_b32 s5, v43, 16
+; GCN-NEXT:    v_readlane_b32 s49, v43, 14
+; GCN-NEXT:    v_readlane_b32 s48, v43, 13
+; GCN-NEXT:    v_readlane_b32 s47, v43, 12
+; GCN-NEXT:    v_readlane_b32 s46, v43, 11
+; GCN-NEXT:    v_readlane_b32 s45, v43, 10
+; GCN-NEXT:    v_readlane_b32 s44, v43, 9
+; GCN-NEXT:    v_readlane_b32 s43, v43, 8
+; GCN-NEXT:    v_readlane_b32 s42, v43, 7
+; GCN-NEXT:    v_readlane_b32 s41, v43, 6
+; GCN-NEXT:    v_readlane_b32 s40, v43, 5
+; GCN-NEXT:    v_readlane_b32 s39, v43, 4
+; GCN-NEXT:    v_readlane_b32 s38, v43, 3
+; GCN-NEXT:    v_readlane_b32 s36, v43, 2
+; GCN-NEXT:    v_readlane_b32 s35, v43, 1
+; GCN-NEXT:    v_readlane_b32 s34, v43, 0
+; GCN-NEXT:    buffer_load_dword v42, off, s[0:3], s33 ; 4-byte Folded Reload
+; GCN-NEXT:    buffer_load_dword v41, off, s[0:3], s33 offset:4 ; 4-byte Folded Reload
+; GCN-NEXT:    buffer_load_dword v40, off, s[0:3], s33 offset:8 ; 4-byte Folded Reload
+; GCN-NEXT:    s_sub_u32 s32, s32, 0x800
+; GCN-NEXT:    v_readlane_b32 s33, v43, 17
+; GCN-NEXT:    s_or_saveexec_b64 s[6:7], -1
+; GCN-NEXT:    buffer_load_dword v43, off, s[0:3], s32 offset:12 ; 4-byte Folded Reload
+; GCN-NEXT:    s_mov_b64 exec, s[6:7]
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    s_setpc_b64 s[4:5]
+  call void %fptr()
+  ret void
+}
 
-; define void @test_indirect_call_vgpr_ptr_arg(void(i32)* %fptr) {
-;   call void %fptr(i32 123)
-;   ret void
-; }
+define void @test_indirect_call_vgpr_ptr_arg(void(i32)* %fptr) {
+; GCN-LABEL: test_indirect_call_vgpr_ptr_arg:
+; GCN:       ; %bb.0:
+; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GCN-NEXT:    s_or_saveexec_b64 s[16:17], -1
+; GCN-NEXT:    buffer_store_dword v44, off, s[0:3], s32 offset:16 ; 4-byte Folded Spill
+; GCN-NEXT:    s_mov_b64 exec, s[16:17]
+; GCN-NEXT:    v_writelane_b32 v44, s33, 21
+; GCN-NEXT:    s_mov_b32 s33, s32
+; GCN-NEXT:    s_add_u32 s32, s32, 0x800
+; GCN-NEXT:    buffer_store_dword v40, off, s[0:3], s33 offset:12 ; 4-byte Folded Spill
+; GCN-NEXT:    buffer_store_dword v41, off, s[0:3], s33 offset:8 ; 4-byte Folded Spill
+; GCN-NEXT:    buffer_store_dword v42, off, s[0:3], s33 offset:4 ; 4-byte Folded Spill
+; GCN-NEXT:    buffer_store_dword v43, off, s[0:3], s33 ; 4-byte Folded Spill
+; GCN-NEXT:    v_writelane_b32 v44, s34, 0
+; GCN-NEXT:    v_writelane_b32 v44, s35, 1
+; GCN-NEXT:    v_writelane_b32 v44, s36, 2
+; GCN-NEXT:    v_writelane_b32 v44, s38, 3
+; GCN-NEXT:    v_writelane_b32 v44, s39, 4
+; GCN-NEXT:    v_writelane_b32 v44, s40, 5
+; GCN-NEXT:    v_writelane_b32 v44, s41, 6
+; GCN-NEXT:    v_writelane_b32 v44, s42, 7
+; GCN-NEXT:    v_writelane_b32 v44, s43, 8
+; GCN-NEXT:    v_writelane_b32 v44, s44, 9
+; GCN-NEXT:    v_writelane_b32 v44, s45, 10
+; GCN-NEXT:    v_writelane_b32 v44, s46, 11
+; GCN-NEXT:    v_writelane_b32 v44, s47, 12
+; GCN-NEXT:    v_writelane_b32 v44, s48, 13
+; GCN-NEXT:    v_writelane_b32 v44, s49, 14
+; GCN-NEXT:    v_writelane_b32 v44, s50, 15
+; GCN-NEXT:    v_writelane_b32 v44, s51, 16
+; GCN-NEXT:    v_writelane_b32 v44, s52, 17
+; GCN-NEXT:    v_writelane_b32 v44, s53, 18
+; GCN-NEXT:    v_writelane_b32 v44, s30, 19
+; GCN-NEXT:    v_writelane_b32 v44, s31, 20
+; GCN-NEXT:    v_mov_b32_e32 v40, v31
+; GCN-NEXT:    s_mov_b32 s34, s14
+; GCN-NEXT:    s_mov_b32 s35, s13
+; GCN-NEXT:    s_mov_b32 s36, s12
+; GCN-NEXT:    s_mov_b64 s[38:39], s[10:11]
+; GCN-NEXT:    s_mov_b64 s[40:41], s[8:9]
+; GCN-NEXT:    s_mov_b64 s[42:43], s[6:7]
+; GCN-NEXT:    s_mov_b64 s[44:45], s[4:5]
+; GCN-NEXT:    v_mov_b32_e32 v42, v1
+; GCN-NEXT:    v_mov_b32_e32 v41, v0
+; GCN-NEXT:    s_mov_b64 s[50:51], s[2:3]
+; GCN-NEXT:    s_mov_b64 s[48:49], s[0:1]
+; GCN-NEXT:    v_mov_b32_e32 v43, 0x7b
+; GCN-NEXT:    s_mov_b64 s[46:47], exec
+; GCN-NEXT:  BB3_1: ; =>This Inner Loop Header: Depth=1
+; GCN-NEXT:    v_readfirstlane_b32 s16, v41
+; GCN-NEXT:    v_readfirstlane_b32 s17, v42
+; GCN-NEXT:    v_cmp_eq_u64_e32 vcc, s[16:17], v[41:42]
+; GCN-NEXT:    s_and_saveexec_b64 s[52:53], vcc
+; GCN-NEXT:    s_mov_b64 s[4:5], s[44:45]
+; GCN-NEXT:    s_mov_b64 s[6:7], s[42:43]
+; GCN-NEXT:    s_mov_b64 s[8:9], s[40:41]
+; GCN-NEXT:    s_mov_b64 s[10:11], s[38:39]
+; GCN-NEXT:    s_mov_b32 s12, s36
+; GCN-NEXT:    s_mov_b32 s13, s35
+; GCN-NEXT:    s_mov_b32 s14, s34
+; GCN-NEXT:    v_mov_b32_e32 v31, v40
+; GCN-NEXT:    s_mov_b64 s[0:1], s[48:49]
+; GCN-NEXT:    s_mov_b64 s[2:3], s[50:51]
+; GCN-NEXT:    v_mov_b32_e32 v0, v43
+; GCN-NEXT:    s_swappc_b64 s[30:31], s[16:17]
+; GCN-NEXT:    s_xor_b64 exec, exec, s[52:53]
+; GCN-NEXT:    s_cbranch_execnz BB3_1
+; GCN-NEXT:  ; %bb.2:
+; GCN-NEXT:    s_mov_b64 exec, s[46:47]
+; GCN-NEXT:    v_readlane_b32 s4, v44, 19
+; GCN-NEXT:    v_readlane_b32 s5, v44, 20
+; GCN-NEXT:    v_readlane_b32 s53, v44, 18
+; GCN-NEXT:    v_readlane_b32 s52, v44, 17
+; GCN-NEXT:    v_readlane_b32 s51, v44, 16
+; GCN-NEXT:    v_readlane_b32 s50, v44, 15
+; GCN-NEXT:    v_readlane_b32 s49, v44, 14
+; GCN-NEXT:    v_readlane_b32 s48, v44, 13
+; GCN-NEXT:    v_readlane_b32 s47, v44, 12
+; GCN-NEXT:    v_readlane_b32 s46, v44, 11
+; GCN-NEXT:    v_readlane_b32 s45, v44, 10
+; GCN-NEXT:    v_readlane_b32 s44, v44, 9
+; GCN-NEXT:    v_readlane_b32 s43, v44, 8
+; GCN-NEXT:    v_readlane_b32 s42, v44, 7
+; GCN-NEXT:    v_readlane_b32 s41, v44, 6
+; GCN-NEXT:    v_readlane_b32 s40, v44, 5
+; GCN-NEXT:    v_readlane_b32 s39, v44, 4
+; GCN-NEXT:    v_readlane_b32 s38, v44, 3
+; GCN-NEXT:    v_readlane_b32 s36, v44, 2
+; GCN-NEXT:    v_readlane_b32 s35, v44, 1
+; GCN-NEXT:    v_readlane_b32 s34, v44, 0
+; GCN-NEXT:    buffer_load_dword v43, off, s[0:3], s33 ; 4-byte Folded Reload
+; GCN-NEXT:    buffer_load_dword v42, off, s[0:3], s33 offset:4 ; 4-byte Folded Reload
+; GCN-NEXT:    buffer_load_dword v41, off, s[0:3], s33 offset:8 ; 4-byte Folded Reload
+; GCN-NEXT:    buffer_load_dword v40, off, s[0:3], s33 offset:12 ; 4-byte Folded Reload
+; GCN-NEXT:    s_sub_u32 s32, s32, 0x800
+; GCN-NEXT:    v_readlane_b32 s33, v44, 21
+; GCN-NEXT:    s_or_saveexec_b64 s[6:7], -1
+; GCN-NEXT:    buffer_load_dword v44, off, s[0:3], s32 offset:16 ; 4-byte Folded Reload
+; GCN-NEXT:    s_mov_b64 exec, s[6:7]
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    s_setpc_b64 s[4:5]
+  call void %fptr(i32 123)
+  ret void
+}