[AMDGPU] Prepare waterfall support for upstream changes (llvm#876)

Upstream change to restrict readfirstlane to non-m0 dst register causes
issues for waterfall support.

Add required changes before taking the upstream change. See
llvm#128851

Author: dstutt

llvm/lib/Target/AMDGPU/SIInsertWaterfall.cpp

@@ -123,14 +123,16 @@ static void readFirstLaneReg(MachineBasicBlock &MBB, MachineRegisterInfo *MRI,
   uint32_t RegSize = RI->getRegSizeInBits(*RFLRegRC) / 32;
   assert(RI->hasVGPRs(MRI->getRegClass(RFLSrcReg)) && "unexpected uniform operand for readfirstlane");
 
-  if (RegSize == 1)
+  if (RegSize == 1) {
+    MRI->constrainRegClass(RFLReg, &AMDGPU::SReg_32_XM0RegClass);
     BuildMI(MBB, I, DL, TII->get(AMDGPU::V_READFIRSTLANE_B32), RFLReg)
         .addReg(RFLSrcReg, getUndefRegState(RFLSrcOp.isUndef()),
                 RFLSrcOp.getSubReg());
+  }
   else {
     SmallVector<Register, 8> TRegs;
     for (unsigned i = 0; i < RegSize; ++i) {
-      Register TReg = MRI->createVirtualRegister(&AMDGPU::SGPR_32RegClass);
+      Register TReg = MRI->createVirtualRegister(&AMDGPU::SReg_32_XM0RegClass);
       BuildMI(MBB, I, DL, TII->get(AMDGPU::V_READFIRSTLANE_B32), TReg)
           .addReg(RFLSrcReg, 0, RI->getSubRegFromChannel(i));
       TRegs.push_back(TReg);
@@ -577,6 +579,9 @@ bool SIInsertWaterfall::processWaterfall(MachineBasicBlock &MBB) {
       CurrIdx.Index = TII->getNamedOperand(*(BeginMI), AMDGPU::OpName::idx);
       CurrIdx.IndexRC = RI->getRegClassForOperandReg(*MRI, *CurrIdx.Index);
       CurrIdx.IndexSRC = RI->getEquivalentSGPRClass(CurrIdx.IndexRC);
+      if (CurrIdx.IndexSRC == &AMDGPU::SGPR_32RegClass)
+        CurrIdx.IndexSRC = &AMDGPU::SReg_32_XM0RegClass;
+
       IndexList.push_back(CurrIdx);
 
       LLVM_DEBUG(if (RI->hasVGPRs(CurrIdx.IndexRC))

llvm/lib/Target/AMDGPU/SIInstructions.td

@@ -1142,7 +1142,7 @@ def SI_WATERFALL_BEGIN_V2 : SI_WATERFALL_BEGIN<VReg_64>;
 def SI_WATERFALL_BEGIN_V4 : SI_WATERFALL_BEGIN<VReg_128>;
 def SI_WATERFALL_BEGIN_V8 : SI_WATERFALL_BEGIN<VReg_256>;
 
-def SI_WATERFALL_READFIRSTLANE_V1 : SI_WATERFALL_READFIRSTLANE<SReg_32, VGPR_32>;
+def SI_WATERFALL_READFIRSTLANE_V1 : SI_WATERFALL_READFIRSTLANE<SReg_32_XM0, VGPR_32>;
 def SI_WATERFALL_READFIRSTLANE_V2 : SI_WATERFALL_READFIRSTLANE<SReg_64, VReg_64>;
 def SI_WATERFALL_READFIRSTLANE_V4 : SI_WATERFALL_READFIRSTLANE<SGPR_128, VReg_128>;
 def SI_WATERFALL_READFIRSTLANE_V8 : SI_WATERFALL_READFIRSTLANE<SReg_256, VReg_256>;

llvm/test/CodeGen/AMDGPU/si-insert-waterfall.mir

@@ -16,7 +16,7 @@ body:             |
     ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vgpr_32 = COPY $vgpr0
     ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vgpr_32 = COPY $vgpr1
     ; CHECK-NEXT: [[REG_SEQUENCE:%[0-9]+]]:sgpr_128 = REG_SEQUENCE [[COPY]], %subreg.sub0, [[COPY]], %subreg.sub1, [[COPY]], %subreg.sub2, [[COPY]], %subreg.sub3
-    ; CHECK-NEXT: [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 0
+    ; CHECK-NEXT: [[S_MOV_B32_:%[0-9]+]]:sreg_32_xm0 = S_MOV_B32 0
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: .1:
     ; CHECK-NEXT: successors: %bb.2(0x80000000)
@@ -28,7 +28,7 @@ body:             |
     ; CHECK-NEXT: successors: %bb.2(0x40000000), %bb.3(0x40000000)
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: [[PHI:%[0-9]+]]:sreg_32_xm0_xexec = PHI [[DEF]], %bb.1, %14, %bb.2
-    ; CHECK-NEXT: [[V_READFIRSTLANE_B32_:%[0-9]+]]:sgpr_32 = V_READFIRSTLANE_B32 [[COPY2]], implicit $exec
+    ; CHECK-NEXT: [[V_READFIRSTLANE_B32_:%[0-9]+]]:sreg_32_xm0 = V_READFIRSTLANE_B32 [[COPY2]], implicit $exec
     ; CHECK-NEXT: [[V_CMP_EQ_U32_e64_:%[0-9]+]]:sreg_32_xm0_xexec = V_CMP_EQ_U32_e64 [[V_READFIRSTLANE_B32_]], [[COPY2]], implicit $exec
     ; CHECK-NEXT: [[S_AND_SAVEEXEC_B32_:%[0-9]+]]:sreg_32_xm0_xexec = S_AND_SAVEEXEC_B32 killed [[V_CMP_EQ_U32_e64_]], implicit-def $exec, implicit-def $scc, implicit $exec
     ; CHECK-NEXT: [[REG_SEQUENCE1:%[0-9]+]]:sreg_64 = REG_SEQUENCE [[V_READFIRSTLANE_B32_]], %subreg.sub0, [[V_READFIRSTLANE_B32_]], %subreg.sub1
@@ -44,10 +44,10 @@ body:             |
   %1:vgpr_32 = COPY $vgpr0
   %2:vgpr_32 = COPY $vgpr1
   %3:sgpr_128 = REG_SEQUENCE %0:sgpr_32, %subreg.sub0, %0:sgpr_32, %subreg.sub1, %0:sgpr_32, %subreg.sub2, %0:sgpr_32, %subreg.sub3
-  %4:sreg_32 = S_MOV_B32 0
-  %5:sreg_32 = SI_WATERFALL_BEGIN_V1 killed %4:sreg_32, %2:vgpr_32
-  %6:sreg_32 = SI_WATERFALL_READFIRSTLANE_V1 %5:sreg_32, %2:vgpr_32
-  %7:sreg_64 = REG_SEQUENCE %6:sreg_32, %subreg.sub0, %6:sreg_32, %subreg.sub1
+  %4:sreg_32_xm0 = S_MOV_B32 0
+  %5:sreg_32 = SI_WATERFALL_BEGIN_V1 killed %4:sreg_32_xm0, %2:vgpr_32
+  %6:sreg_32_xm0 = SI_WATERFALL_READFIRSTLANE_V1 %5:sreg_32, %2:vgpr_32
+  %7:sreg_64 = REG_SEQUENCE %6:sreg_32_xm0, %subreg.sub0, %6:sreg_32_xm0, %subreg.sub1
   %8:sgpr_256 = S_LOAD_DWORDX8_IMM killed %7:sreg_64, 0, 0
   %9:vreg_128 = IMAGE_SAMPLE_V4_V1_gfx10 %1:vgpr_32, killed %8:sgpr_256, killed %3:sgpr_128, 15, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec :: (dereferenceable load (s128), addrspace 8)
   %10:vreg_128 = SI_WATERFALL_END_V4 %5:sreg_32, killed %9:vreg_128