[AMDGPU] Optimizing unnecessary copies for REG_SEQUENCE PHI operand. Also fixes rocBLAS error

[alex-t]

Summary

This change optimizes out unnecessary copies that appear in the case when the PHI node input register is defined by the REG_SEQUENCE.
The REG_SEQUENCE not necessarily is located in the PHI's predecessor block.

rocBLAS error is really caused by the subregisters tracking bug. PostRA Machine Sink cannot discover that sreg_32 writes are defining superreg that is used as source operand in subsequent COPY. Machine Sink then sinks the superreg COPY to the "THEN" block of the branch but the copies of sreg32 that deine superreg COPY source operand - to the "ELSE" block. As a result "THEN" block receives uninitialized register.
This issue is going to be addressed in separate change later on.

The optimization in current change removes the IR pattern that triggers the Machine Sink bug.
Since it need to be done anyway to remove unnecessary copies, it goes first.

[vpykhtin]

Looks good, but the test would be nice to have.

[arsenm]

Definitely needs a test. I'm not sure I understand the fix here. This is some kind of workaround? I would be happier with the subreg fix happening first. Removing unnecessary copies should not be relevant to the correctness fix?

llvm/lib/Target/AMDGPU/SIFixSGPRCopies.cpp

275

Factoring this out is a separate patch

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

4274

I think the legalization done here should only touch the local instruction, and not be scanning through copy chains. This seems like it's trying to avoid a failure of the PeepholeOptimizer?

4290–4294

Needs comments

4303

Doesn't need to be a reference

[alex-t]

In D71089#1772921, @arsenm wrote:

Definitely needs a test. I'm not sure I understand the fix here. This is some kind of workaround? I would be happier with the subreg fix happening first. Removing unnecessary copies should not be relevant to the correctness fix?

Agree. The reason to try submitting this first is that change in lib/CodeGen/MachineSink.cpp that fixes the issue will likely take a long time to be approved.
This issue with this https://reviews.llvm.org/D68635 patch breaks one of the rocBLAS kernels correctness because of sinking subreg definition out from the superreg use.
This in order creates several duplicated P1 bugs.

Removing unnecessary copies just covers the IR pattern that drives PostRA Machine Sink mad.

If I could to speedup the PostRA Machine Sink fix review somehow, I prefer to submit it first and work on the copy elimination patch refactoring to make it more elegant.

The PostRA Machine Sink fix is here: https://reviews.llvm.org/D71132

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

4274

Initially I had intention to generalize your code that handles single COPY user of REG_SEQUENCE (SIFixSGPRCopies.cpp). Then I realized that the pattern I'm chasing to eliminate is only related to the PHIs.
Namely, it is the single-use chain of copies that connects REG_SEQUENCE that produce the value to the PHI that consumes it. The chain is possibly going through the several BBs.
That's why I opted for the backward search from the PHI up to the REG_SEQUENCE .
I changed the walk in the runOnMachineFunction main loop to post order to ensure that PHI is processed first, before the REG_SEQUENCE.
I agree that generalizing foldVGPRCopyIntoRegSequence would be more neat.

[arsenm]

Is this still relevant?

llvm/lib/Target/AMDGPU/SIFixSGPRCopies.cpp

 #include "AMDGPUSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "SIInstrInfo.h"
 #include "SIRegisterInfo.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
static bool foldVGPRCopyIntoRegSequence(MachineInstr &MI,
 
   if (tryChangeVGPRtoSGPRinCopy(CopyUse, TRI, TII))
     return true;
 
   // TODO: Could have multiple extracts?
   unsigned SubReg = CopyUse.getOperand(1).getSubReg();
   if (SubReg != AMDGPU::NoSubRegister)
     return false;
-
-  MRI.setRegClass(DstReg, DstRC);
-
-  // SGPRx = ...
-  // SGPRy = REG_SEQUENCE SGPRx, sub0 ...
-  // VGPRz = COPY SGPRy
-
-  // =>
-  // VGPRx = COPY SGPRx
-  // VGPRz = REG_SEQUENCE VGPRx, sub0
-
-  MI.getOperand(0).setReg(CopyUse.getOperand(0).getReg());
-  bool IsAGPR = TRI->hasAGPRs(DstRC);
-
-  for (unsigned I = 1, N = MI.getNumOperands(); I != N; I += 2) {
-    Register SrcReg = MI.getOperand(I).getReg();
-    unsigned SrcSubReg = MI.getOperand(I).getSubReg();
-
-    const TargetRegisterClass *SrcRC = MRI.getRegClass(SrcReg);
-    assert(TRI->isSGPRClass(SrcRC) &&
-           "Expected SGPR REG_SEQUENCE to only have SGPR inputs");
-
-    SrcRC = TRI->getSubRegClass(SrcRC, SrcSubReg);
-    const TargetRegisterClass *NewSrcRC = TRI->getEquivalentVGPRClass(SrcRC);
-
-    Register TmpReg = MRI.createVirtualRegister(NewSrcRC);
-
-    BuildMI(*MI.getParent(), &MI, MI.getDebugLoc(), TII->get(AMDGPU::COPY),
-            TmpReg)
-        .add(MI.getOperand(I));
-
-    if (IsAGPR) {
-      const TargetRegisterClass *NewSrcRC = TRI->getEquivalentAGPRClass(SrcRC);
-      Register TmpAReg = MRI.createVirtualRegister(NewSrcRC);
-      unsigned Opc = NewSrcRC == &AMDGPU::AGPR_32RegClass ?
-        AMDGPU::V_ACCVGPR_WRITE_B32 : AMDGPU::COPY;
-      BuildMI(*MI.getParent(), &MI, MI.getDebugLoc(), TII->get(Opc),
-            TmpAReg)
-        .addReg(TmpReg, RegState::Kill);
-      TmpReg = TmpAReg;
-    }
-
-    MI.getOperand(I).setReg(TmpReg);
-  }
+  TII->convertRegSequenceToVGPR(MI, CopyUse.getOperand(0).getReg(), DstRC);
 
   CopyUse.eraseFromParent();
   return true;
 }
 
 static bool isSafeToFoldImmIntoCopy(const MachineInstr *Copy,
                                     const MachineInstr *MoveImm,
                                     const SIInstrInfo *TII,

 bool SIFixSGPRCopies::runOnMachineFunction(MachineFunction &MF) {
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   MRI = &MF.getRegInfo();
   TRI = ST.getRegisterInfo();
   TII = ST.getInstrInfo();
   MDT = &getAnalysis<MachineDominatorTree>();
 
   SmallVector<MachineInstr *, 16> Worklist;
+  for (auto MBB : post_order<MachineFunction*>(&MF)) {
 
-  for (MachineFunction::iterator BI = MF.begin(), BE = MF.end();
-                                                  BI != BE; ++BI) {
-    MachineBasicBlock &MBB = *BI;
-    for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end();
+    for (MachineBasicBlock::iterator I = MBB->begin(), E = MBB->end();
          I != E; ++I) {
       MachineInstr &MI = *I;
 
       switch (MI.getOpcode()) {
       default:
         continue;
       case AMDGPU::COPY:
       case AMDGPU::WQM:
       case AMDGPU::SOFT_WQM:
       case AMDGPU::WWM: {
         Register DstReg = MI.getOperand(0).getReg();
 
         const TargetRegisterClass *SrcRC, *DstRC;
         std::tie(SrcRC, DstRC) = getCopyRegClasses(MI, *TRI, *MRI);
 
         if (!Register::isVirtualRegister(DstReg)) {
           // If the destination register is a physical register there isn't
           // really much we can do to fix this.
           // Some special instructions use M0 as an input. Some even only use
           // the first lane. Insert a readfirstlane and hope for the best.
           if (DstReg == AMDGPU::M0 && TRI->hasVectorRegisters(SrcRC)) {
             Register TmpReg
               = MRI->createVirtualRegister(&AMDGPU::SReg_32_XM0RegClass);
 
-            BuildMI(MBB, MI, MI.getDebugLoc(),
+            BuildMI(*MBB, MI, MI.getDebugLoc(),
                     TII->get(AMDGPU::V_READFIRSTLANE_B32), TmpReg)
               .add(MI.getOperand(1));
             MI.getOperand(1).setReg(TmpReg);
           }
 
           continue;
         }
 

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
 
   void legalizeOperandsSMRD(MachineRegisterInfo &MRI, MachineInstr &MI) const;
 
   void legalizeGenericOperand(MachineBasicBlock &InsertMBB,
                               MachineBasicBlock::iterator I,
                               const TargetRegisterClass *DstRC,
                               MachineOperand &Op, MachineRegisterInfo &MRI,
                               const DebugLoc &DL) const;
+  // SGPRx = ...
+  // SGPRy = REG_SEQUENCE SGPRx, sub0 ...
+  // VGPRz = COPY SGPRy
+  // =>
+  // VGPRx = COPY SGPRx
+  // VGPRz = REG_SEQUENCE VGPRx, sub0
+  void convertRegSequenceToVGPR(MachineInstr &MI, Register DstReg,
+                                const TargetRegisterClass *DstRC) const;
 
   /// Legalize all operands in this instruction.  This function may create new
   /// instructions and control-flow around \p MI.  If present, \p MDT is
   /// updated.
   void legalizeOperands(MachineInstr &MI,
                         MachineDominatorTree *MDT = nullptr) const;
 
   /// Replace this instruction's opcode with the equivalent VALU

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

void SIInstrInfo::legalizeOperandsSMRD(MachineRegisterInfo &MRI,
   }
   MachineOperand *SOff = getNamedOperand(MI, AMDGPU::OpName::soff);
   if (SOff && !RI.isSGPRClass(MRI.getRegClass(SOff->getReg()))) {
     unsigned SGPR = readlaneVGPRToSGPR(SOff->getReg(), MI, MRI);
     SOff->setReg(SGPR);
   }
 }
 
 void SIInstrInfo::legalizeGenericOperand(MachineBasicBlock &InsertMBB,
                                          MachineBasicBlock::iterator I,
                                          const TargetRegisterClass *DstRC,
                                          MachineOperand &Op,
                                          MachineRegisterInfo &MRI,
                                          const DebugLoc &DL) const {
   Register OpReg = Op.getReg();
   unsigned OpSubReg = Op.getSubReg();
 
   const TargetRegisterClass *OpRC = RI.getSubClassWithSubReg(
       RI.getRegClassForReg(MRI, OpReg), OpSubReg);
 
   // Check if operand is already the correct register class.
   if (DstRC == OpRC)
     return;
 
+  MachineInstr *Def = MRI.getVRegDef(OpReg);
+  if (Op.getParent()->isPHI() && !RI.isSGPRClass(DstRC) &&
+      RI.isSGPRClass(OpRC)) {
+    SmallVector<MachineInstr*, 4> CopiesToDelete;
+    while (Def && Def->isCopy() && MRI.hasOneUse(Def->getOperand(0).getReg())) {
+      CopiesToDelete.push_back(Def);
+      Def = MRI.getUniqueVRegDef(Def->getOperand(1).getReg());
+    }
+    if (Def && Def->isRegSequence()) {
+      Register Dst = Def->getOperand(0).getReg();
+      if (MRI.hasOneUse(Dst)) {
+        convertRegSequenceToVGPR(*Def,
+          OpReg, DstRC);
+        for (auto& Copy : CopiesToDelete)
+          Copy->eraseFromParent();
+        return;
+      }
+    }
+  }
+
   Register DstReg = MRI.createVirtualRegister(DstRC);
   MachineInstr *Copy =
       BuildMI(InsertMBB, I, DL, get(AMDGPU::COPY), DstReg).add(Op);
 
   Op.setReg(DstReg);
   Op.setSubReg(0);
 
-  MachineInstr *Def = MRI.getVRegDef(OpReg);
   if (!Def)
     return;
 
   // Try to eliminate the copy if it is copying an immediate value.
   if (Def->isMoveImmediate() && DstRC != &AMDGPU::VReg_1RegClass)
     FoldImmediate(*Copy, *Def, OpReg, &MRI);
 
   bool ImpDef = Def->isImplicitDef();
   while (!ImpDef && Def && Def->isCopy()) {
     if (Def->getOperand(1).getReg().isPhysical())
       break;
     Def = MRI.getUniqueVRegDef(Def->getOperand(1).getReg());
     ImpDef = Def && Def->isImplicitDef();
   }
   if (!RI.isSGPRClass(DstRC) && !Copy->readsRegister(AMDGPU::EXEC, &RI) &&
       !ImpDef)
     Copy->addOperand(MachineOperand::CreateReg(AMDGPU::EXEC, false, true));
 }
 
+void SIInstrInfo::convertRegSequenceToVGPR(MachineInstr &MI, Register DstReg,
+                                           const TargetRegisterClass *DstRC) const {
+  assert(MI.isRegSequence());
+  MachineRegisterInfo& MRI = MI.getParent()->getParent()->getRegInfo();
+
+  MRI.setRegClass(DstReg, DstRC);
+
+  MI.getOperand(0).setReg(DstReg);
+  bool IsAGPR = RI.hasAGPRs(DstRC);
+
+  for (unsigned I = 1, N = MI.getNumOperands(); I != N; I += 2) {
+    Register SrcReg = MI.getOperand(I).getReg();
+    unsigned SrcSubReg = MI.getOperand(I).getSubReg();
+
+    const TargetRegisterClass* SrcRC = MRI.getRegClass(SrcReg);
+
+    SrcRC = RI.getSubRegClass(SrcRC, SrcSubReg);
+    const TargetRegisterClass* NewSrcRC = RI.getEquivalentVGPRClass(SrcRC);
+
+    Register TmpReg = MRI.createVirtualRegister(NewSrcRC);
+
+    BuildMI(*MI.getParent(), &MI, MI.getDebugLoc(), get(AMDGPU::COPY),
+      TmpReg)
+      .add(MI.getOperand(I));
+
+    if (IsAGPR) {
+      const TargetRegisterClass* NewSrcRC = RI.getEquivalentAGPRClass(SrcRC);
+      Register TmpAReg = MRI.createVirtualRegister(NewSrcRC);
+      unsigned Opc = NewSrcRC == &AMDGPU::AGPR_32RegClass ?
+        AMDGPU::V_ACCVGPR_WRITE_B32 : AMDGPU::COPY;
+      BuildMI(*MI.getParent(), &MI, MI.getDebugLoc(), get(Opc),
+        TmpAReg)
+        .addReg(TmpReg, RegState::Kill);
+      TmpReg = TmpAReg;
+    }
+
+    MI.getOperand(I).setReg(TmpReg);
+  }
+}
+
 // 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
 // iteration, in the worst case we execute 64 (once per lane).
 static void
 emitLoadSRsrcFromVGPRLoop(const SIInstrInfo &TII, MachineRegisterInfo &MRI,
                           MachineBasicBlock &OrigBB, MachineBasicBlock &LoopBB,
                           const DebugLoc &DL, MachineOperand &Rsrc) {
   MachineFunction &MF = *OrigBB.getParent();