AMDGPU/GlobalISel: Add waterfall lowering in regbanklegalize

Add rules for G_AMDGPU_BUFFER_LOAD and implement waterfall lowering
for divergent operands that must be sgpr.

Author: petar-avramovic

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.cpp

@@ -117,45 +117,72 @@ static LLT getReadAnyLaneSplitTy(LLT Ty) {
   return LLT::scalar(32);
 }
 
-static Register buildReadAnyLane(MachineIRBuilder &B, Register VgprSrc,
-                                 const RegisterBankInfo &RBI);
+using ReadLaneFnTy =
+    function_ref<MachineInstrBuilder(MachineIRBuilder &, Register, Register)>;
+
+static Register buildReadLane(MachineIRBuilder &, Register,
+                              const RegisterBankInfo &, ReadLaneFnTy);
 
 static void unmergeReadAnyLane(MachineIRBuilder &B,
                                SmallVectorImpl<Register> &SgprDstParts,
                                LLT UnmergeTy, Register VgprSrc,
-                               const RegisterBankInfo &RBI) {
+                               const RegisterBankInfo &RBI,
+                               ReadLaneFnTy BuildRL) {
   const RegisterBank *VgprRB = &RBI.getRegBank(AMDGPU::VGPRRegBankID);
   auto Unmerge = B.buildUnmerge({VgprRB, UnmergeTy}, VgprSrc);
   for (unsigned i = 0; i < Unmerge->getNumOperands() - 1; ++i) {
-    SgprDstParts.push_back(buildReadAnyLane(B, Unmerge.getReg(i), RBI));
+    SgprDstParts.push_back(buildReadLane(B, Unmerge.getReg(i), RBI, BuildRL));
   }
 }
 
-static Register buildReadAnyLane(MachineIRBuilder &B, Register VgprSrc,
-                                 const RegisterBankInfo &RBI) {
+static Register buildReadLane(MachineIRBuilder &B, Register VgprSrc,
+                              const RegisterBankInfo &RBI,
+                              ReadLaneFnTy BuildRL) {
   LLT Ty = B.getMRI()->getType(VgprSrc);
   const RegisterBank *SgprRB = &RBI.getRegBank(AMDGPU::SGPRRegBankID);
   if (Ty.getSizeInBits() == 32) {
-    return B.buildInstr(AMDGPU::G_AMDGPU_READANYLANE, {{SgprRB, Ty}}, {VgprSrc})
-        .getReg(0);
+    Register SgprDst = B.getMRI()->createVirtualRegister({SgprRB, Ty});
+    return BuildRL(B, SgprDst, VgprSrc).getReg(0);
   }
 
   SmallVector<Register, 8> SgprDstParts;
-  unmergeReadAnyLane(B, SgprDstParts, getReadAnyLaneSplitTy(Ty), VgprSrc, RBI);
+  unmergeReadAnyLane(B, SgprDstParts, getReadAnyLaneSplitTy(Ty), VgprSrc, RBI,
+                     BuildRL);
 
   return B.buildMergeLikeInstr({SgprRB, Ty}, SgprDstParts).getReg(0);
 }
 
-void AMDGPU::buildReadAnyLane(MachineIRBuilder &B, Register SgprDst,
-                              Register VgprSrc, const RegisterBankInfo &RBI) {
+static void buildReadLane(MachineIRBuilder &B, Register SgprDst,
+                          Register VgprSrc, const RegisterBankInfo &RBI,
+                          ReadLaneFnTy BuildReadLane) {
   LLT Ty = B.getMRI()->getType(VgprSrc);
   if (Ty.getSizeInBits() == 32) {
-    B.buildInstr(AMDGPU::G_AMDGPU_READANYLANE, {SgprDst}, {VgprSrc});
+    BuildReadLane(B, SgprDst, VgprSrc);
     return;
   }
 
   SmallVector<Register, 8> SgprDstParts;
-  unmergeReadAnyLane(B, SgprDstParts, getReadAnyLaneSplitTy(Ty), VgprSrc, RBI);
+  unmergeReadAnyLane(B, SgprDstParts, getReadAnyLaneSplitTy(Ty), VgprSrc, RBI,
+                     BuildReadLane);
 
   B.buildMergeLikeInstr(SgprDst, SgprDstParts).getReg(0);
 }
+
+void AMDGPU::buildReadAnyLane(MachineIRBuilder &B, Register SgprDst,
+                              Register VgprSrc, const RegisterBankInfo &RBI) {
+  return buildReadLane(
+      B, SgprDst, VgprSrc, RBI,
+      [](MachineIRBuilder &B, Register SgprDst, Register VgprSrc) {
+        return B.buildInstr(AMDGPU::G_AMDGPU_READANYLANE, {SgprDst}, {VgprSrc});
+      });
+}
+
+void AMDGPU::buildReadFirstLane(MachineIRBuilder &B, Register SgprDst,
+                                Register VgprSrc, const RegisterBankInfo &RBI) {
+  return buildReadLane(
+      B, SgprDst, VgprSrc, RBI,
+      [](MachineIRBuilder &B, Register SgprDst, Register VgprSrc) {
+        return B.buildIntrinsic(Intrinsic::amdgcn_readfirstlane, SgprDst)
+            .addReg(VgprSrc);
+      });
+}

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.h

@@ -51,6 +51,8 @@ class IntrinsicLaneMaskAnalyzer {
 
 void buildReadAnyLane(MachineIRBuilder &B, Register SgprDst, Register VgprSrc,
                       const RegisterBankInfo &RBI);
+void buildReadFirstLane(MachineIRBuilder &B, Register SgprDst, Register VgprSrc,
+                        const RegisterBankInfo &RBI);
 }
 }
 

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalize.cpp

@@ -182,9 +182,11 @@ class AMDGPURegBankLegalizeCombiner {
     if (UnMerge) {
       int Idx = UnMerge->findRegisterDefOperandIdx(Src, nullptr);
       auto *Merge = getOpcodeDef<GMergeLikeInstr>(UnMerge->getSourceReg(), MRI);
-      if (Merge) {
-        auto [RAL, RALSrc] =
-            tryMatch(Merge->getSourceReg(Idx), AMDGPU::G_AMDGPU_READANYLANE);
+      if (Merge && UnMerge->getNumDefs() == Merge->getNumSources()) {
+        Register SrcRegIdx = Merge->getSourceReg(Idx);
+        if (MRI.getType(Src) != MRI.getType(SrcRegIdx))
+          return {};
+        auto [RAL, RALSrc] = tryMatch(SrcRegIdx, AMDGPU::G_AMDGPU_READANYLANE);
         if (RAL)
           return RALSrc;
       }
@@ -203,7 +205,14 @@ class AMDGPURegBankLegalizeCombiner {
   bool tryEliminateReadAnyLane(MachineInstr &Copy) {
     Register Dst = Copy.getOperand(0).getReg();
     Register Src = Copy.getOperand(1).getReg();
-    if (!Src.isVirtual())
+
+    // Skip non-vgpr Dst
+    if ((Dst.isVirtual() && MRI.getRegBankOrNull(Dst) != VgprRB) ||
+        (Dst.isPhysical() && !TRI.isVGPR(MRI, Dst)))
+      return false;
+
+    // Skip physical source registers and source registers with register class
+    if (!Src.isVirtual() || MRI.getRegClassOrNull(Src))
       return false;
 
     Register RALDst = Src;