AMDGPU: Refactor indirect vector lowering

[arsenm]

Summary

Allow inserting multiple instructions in the expanded loop.

[• tstellarAMD]

LGTM.

[arsenm]

r283177

lib/Target/AMDGPU/SIISelLowering.cpp

MachineBasicBlock *SITargetLowering::splitKillBlock(MachineInstr &MI,
   return SplitBB;
 }
 
 // Do a v_movrels_b32 or v_movreld_b32 for each unique value of \p IdxReg in the
 // wavefront. If the value is uniform and just happens to be in a VGPR, this
 // will only do one iteration. In the worst case, this will loop 64 times.
 //
 // TODO: Just use v_readlane_b32 if we know the VGPR has a uniform value.
-static void emitLoadM0FromVGPRLoop(const SIInstrInfo *TII,
-                                   MachineRegisterInfo &MRI,
-                                   MachineBasicBlock &OrigBB,
-                                   MachineBasicBlock &LoopBB,
-                                   const DebugLoc &DL,
-                                   MachineInstr *MovRel,
-                                   const MachineOperand &IdxReg,
-                                   unsigned InitReg,
-                                   unsigned ResultReg,
-                                   unsigned PhiReg,
-                                   unsigned InitSaveExecReg,
-                                   int Offset) {
+static MachineBasicBlock::iterator emitLoadM0FromVGPRLoop(
+  const SIInstrInfo *TII,
   MachineBasicBlock::iterator I = LoopBB.begin();
 
   unsigned PhiExec = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);
   unsigned NewExec = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);
   unsigned CurrentIdxReg = MRI.createVirtualRegister(&AMDGPU::SGPR_32RegClass);
   unsigned CondReg = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);
 
   BuildMI(LoopBB, I, DL, TII->get(TargetOpcode::PHI), PhiReg)
static MachineBasicBlock::iterator emitLoadM0FromVGPRLoop(
   }
 
   // Update EXEC, save the original EXEC value to VCC.
   BuildMI(LoopBB, I, DL, TII->get(AMDGPU::S_AND_SAVEEXEC_B64), NewExec)
     .addReg(CondReg, RegState::Kill);
 
   MRI.setSimpleHint(NewExec, CondReg);
 
-  // Do the actual move.
-  LoopBB.insert(I, MovRel);
-
   // Update EXEC, switch all done bits to 0 and all todo bits to 1.
+  MachineInstr *InsertPt =
     BuildMI(LoopBB, I, DL, TII->get(AMDGPU::S_XOR_B64), AMDGPU::EXEC)
     .addReg(AMDGPU::EXEC)
     .addReg(NewExec);
 
   // XXX - s_xor_b64 sets scc to 1 if the result is nonzero, so can we use
   // s_cbranch_scc0?
 
   // Loop back to V_READFIRSTLANE_B32 if there are still variants to cover.
   BuildMI(LoopBB, I, DL, TII->get(AMDGPU::S_CBRANCH_EXECNZ))
     .addMBB(&LoopBB);
+
+  return InsertPt->getIterator();
 }
 
 // This has slightly sub-optimal regalloc when the source vector is killed by
 // the read. The register allocator does not understand that the kill is
 // per-workitem, so is kept alive for the whole loop so we end up not re-using a
 // subregister from it, using 1 more VGPR than necessary. This was saved when
 // this was expanded after register allocation.
-static MachineBasicBlock *loadM0FromVGPR(const SIInstrInfo *TII,
+static MachineBasicBlock::iterator loadM0FromVGPR(const SIInstrInfo *TII,
                                                   MachineBasicBlock &MBB,
                                                   MachineInstr &MI,
-                                         MachineInstr *MovRel,
                                                   unsigned InitResultReg,
                                                   unsigned PhiReg,
                                                   int Offset) {
   MachineFunction *MF = MBB.getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
   const DebugLoc &DL = MI.getDebugLoc();
   MachineBasicBlock::iterator I(&MI);
 
   unsigned DstReg = MI.getOperand(0).getReg();
   unsigned SaveExec = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);
   unsigned TmpExec = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);
static MachineBasicBlock::iterator loadM0FromVGPR(const SIInstrInfo *TII,
   // Move the rest of the block into a new block.
   RemainderBB->transferSuccessorsAndUpdatePHIs(&MBB);
   RemainderBB->splice(RemainderBB->begin(), &MBB, I, MBB.end());
 
   MBB.addSuccessor(LoopBB);
 
   const MachineOperand *Idx = TII->getNamedOperand(MI, AMDGPU::OpName::idx);
 
-  emitLoadM0FromVGPRLoop(TII, MRI, MBB, *LoopBB, DL, MovRel, *Idx,
-                         InitResultReg, DstReg, PhiReg, TmpExec, Offset);
+  auto InsPt = emitLoadM0FromVGPRLoop(TII, MRI, MBB, *LoopBB, DL, *Idx,
+                                      InitResultReg, DstReg, PhiReg, TmpExec,
+                                      Offset);
 
   MachineBasicBlock::iterator First = RemainderBB->begin();
   BuildMI(*RemainderBB, First, DL, TII->get(AMDGPU::S_MOV_B64), AMDGPU::EXEC)
     .addReg(SaveExec);
 
   MI.eraseFromParent();
 
-  return RemainderBB;
+  return InsPt;
 }
 
 // Returns subreg index, offset
 static std::pair<unsigned, int>
 computeIndirectRegAndOffset(const SIRegisterInfo &TRI,
                             const TargetRegisterClass *SuperRC,
                             unsigned VecReg,
                             int Offset) {
static bool setM0ToIndexFromSGPR(const SIInstrInfo *TII,
   }
 
   return true;
 }
 
 // Control flow needs to be inserted if indexing with a VGPR.
 static MachineBasicBlock *emitIndirectSrc(MachineInstr &MI,
                                           MachineBasicBlock &MBB,
-                                          const SIInstrInfo *TII) {
+                                          const SISubtarget &ST) {
+  const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo &TRI = TII->getRegisterInfo();
   MachineFunction *MF = MBB.getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
 
   unsigned Dst = MI.getOperand(0).getReg();
   const MachineOperand *SrcVec = TII->getNamedOperand(MI, AMDGPU::OpName::src);
   int Offset = TII->getNamedOperand(MI, AMDGPU::OpName::offset)->getImm();
 
static MachineBasicBlock *emitIndirectSrc(MachineInstr &MI,
   const DebugLoc &DL = MI.getDebugLoc();
   MachineBasicBlock::iterator I(&MI);
 
   unsigned PhiReg = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
   unsigned InitReg = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
 
   BuildMI(MBB, I, DL, TII->get(TargetOpcode::IMPLICIT_DEF), InitReg);
 
-  MachineInstr *MovRel =
-    BuildMI(*MF, DL, TII->get(AMDGPU::V_MOVRELS_B32_e32), Dst)
+
+  auto InsPt = loadM0FromVGPR(TII, MBB, MI, InitReg, PhiReg, Offset);
+
+  BuildMI(*InsPt->getParent(), InsPt, DL,
+          TII->get(AMDGPU::V_MOVRELS_B32_e32), Dst)
     .addReg(SrcVec->getReg(), RegState::Undef, SubReg)
     .addReg(SrcVec->getReg(), RegState::Implicit);
 
-  return loadM0FromVGPR(TII, MBB, MI, MovRel, InitReg, PhiReg, Offset);
+  return InsPt->getParent();
 }
 
 static MachineBasicBlock *emitIndirectDst(MachineInstr &MI,
                                           MachineBasicBlock &MBB,
-                                          const SIInstrInfo *TII) {
+                                          const SISubtarget &ST) {
+  const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo &TRI = TII->getRegisterInfo();
   MachineFunction *MF = MBB.getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
 
   unsigned Dst = MI.getOperand(0).getReg();
   const MachineOperand *SrcVec = TII->getNamedOperand(MI, AMDGPU::OpName::src);
   const MachineOperand *Idx = TII->getNamedOperand(MI, AMDGPU::OpName::idx);
   const MachineOperand *Val = TII->getNamedOperand(MI, AMDGPU::OpName::val);
static MachineBasicBlock *emitIndirectDst(MachineInstr &MI,
   }
 
   if (Val->isReg())
     MRI.clearKillFlags(Val->getReg());
 
   const DebugLoc &DL = MI.getDebugLoc();
   unsigned PhiReg = MRI.createVirtualRegister(VecRC);
 
+  auto InsPt = loadM0FromVGPR(TII, MBB, MI, SrcVec->getReg(), PhiReg, Offset);
+
   // vdst is not actually read and just provides the base register index.
   MachineInstr *MovRel =
-    BuildMI(*MF, DL, MovRelDesc)
+    BuildMI(*InsPt->getParent(), InsPt, DL, MovRelDesc)
     .addReg(PhiReg, RegState::Undef, SubReg) // vdst
     .addOperand(*Val)
     .addReg(Dst, RegState::ImplicitDefine)
     .addReg(PhiReg, RegState::Implicit);
 
   const int ImpDefIdx = MovRelDesc.getNumOperands() +
     MovRelDesc.getNumImplicitUses();
   const int ImpUseIdx = ImpDefIdx + 1;
 
   MovRel->tieOperands(ImpDefIdx, ImpUseIdx);
 
-  return loadM0FromVGPR(TII, MBB, MI, MovRel,
-                        SrcVec->getReg(), PhiReg, Offset);
+  return InsPt->getParent();
 }
 
 MachineBasicBlock *SITargetLowering::EmitInstrWithCustomInserter(
   MachineInstr &MI, MachineBasicBlock *BB) const {
   switch (MI.getOpcode()) {
   case AMDGPU::SI_INIT_M0: {
     const SIInstrInfo *TII = getSubtarget()->getInstrInfo();
     BuildMI(*BB, MI.getIterator(), MI.getDebugLoc(),
case AMDGPU::GET_GROUPSTATICSIZE: {
     MI.eraseFromParent();
     return BB;
   }
   case AMDGPU::SI_INDIRECT_SRC_V1:
   case AMDGPU::SI_INDIRECT_SRC_V2:
   case AMDGPU::SI_INDIRECT_SRC_V4:
   case AMDGPU::SI_INDIRECT_SRC_V8:
   case AMDGPU::SI_INDIRECT_SRC_V16:
-    return emitIndirectSrc(MI, *BB, getSubtarget()->getInstrInfo());
+    return emitIndirectSrc(MI, *BB, *getSubtarget());
   case AMDGPU::SI_INDIRECT_DST_V1:
   case AMDGPU::SI_INDIRECT_DST_V2:
   case AMDGPU::SI_INDIRECT_DST_V4:
   case AMDGPU::SI_INDIRECT_DST_V8:
   case AMDGPU::SI_INDIRECT_DST_V16:
-    return emitIndirectDst(MI, *BB, getSubtarget()->getInstrInfo());
+    return emitIndirectDst(MI, *BB, *getSubtarget());
   case AMDGPU::SI_KILL:
     return splitKillBlock(MI, BB);
   case AMDGPU::V_CNDMASK_B64_PSEUDO: {
     MachineRegisterInfo &MRI = BB->getParent()->getRegInfo();
     const SIInstrInfo *TII = getSubtarget()->getInstrInfo();
 
     unsigned Dst = MI.getOperand(0).getReg();
     unsigned Src0 = MI.getOperand(1).getReg();