AMDGPU/GlobalISel: add RegBankLegalize rules for select (#132384)

Uniform condition S1 is AnyExtended to S32 and high bits are cleaned using
AND with 1. Divergent S1 uses VCC.
Using B32/B64 rules to cover scalars vector and pointer types.
Divergent B64 is split to S32.

Author: petar-avramovic

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp

@@ -23,6 +23,7 @@
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineUniformityAnalysis.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/Support/AMDGPUAddrSpace.h"
 
 #define DEBUG_TYPE "amdgpu-regbanklegalize"
 
@@ -286,6 +287,25 @@ void RegBankLegalizeHelper::lowerSplitTo32(MachineInstr &MI) {
   MI.eraseFromParent();
 }
 
+void RegBankLegalizeHelper::lowerSplitTo32Select(MachineInstr &MI) {
+  Register Dst = MI.getOperand(0).getReg();
+  LLT DstTy = MRI.getType(Dst);
+  assert(DstTy == V4S16 || DstTy == V2S32 || DstTy == S64 ||
+         (DstTy.isPointer() && DstTy.getSizeInBits() == 64));
+  LLT Ty = DstTy == V4S16 ? V2S16 : S32;
+  auto Op2 = B.buildUnmerge({VgprRB, Ty}, MI.getOperand(2).getReg());
+  auto Op3 = B.buildUnmerge({VgprRB, Ty}, MI.getOperand(3).getReg());
+  Register Cond = MI.getOperand(1).getReg();
+  auto Flags = MI.getFlags();
+  auto Lo =
+      B.buildSelect({VgprRB, Ty}, Cond, Op2.getReg(0), Op3.getReg(0), Flags);
+  auto Hi =
+      B.buildSelect({VgprRB, Ty}, Cond, Op2.getReg(1), Op3.getReg(1), Flags);
+
+  B.buildMergeLikeInstr(Dst, {Lo, Hi});
+  MI.eraseFromParent();
+}
+
 void RegBankLegalizeHelper::lower(MachineInstr &MI,
                                   const RegBankLLTMapping &Mapping,
                                   SmallSet<Register, 4> &WaterfallSgprs) {
@@ -372,6 +392,8 @@ void RegBankLegalizeHelper::lower(MachineInstr &MI,
     return lowerS_BFE(MI);
   case SplitTo32:
     return lowerSplitTo32(MI);
+  case SplitTo32Select:
+    return lowerSplitTo32Select(MI);
   case SplitLoad: {
     LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
     unsigned Size = DstTy.getSizeInBits();
@@ -485,7 +507,8 @@ LLT RegBankLegalizeHelper::getBTyFromID(RegBankLLTMappingApplyID ID, LLT Ty) {
   case UniInVgprB64:
     if (Ty == LLT::scalar(64) || Ty == LLT::fixed_vector(2, 32) ||
         Ty == LLT::fixed_vector(4, 16) || Ty == LLT::pointer(0, 64) ||
-        Ty == LLT::pointer(1, 64) || Ty == LLT::pointer(4, 64))
+        Ty == LLT::pointer(1, 64) || Ty == LLT::pointer(4, 64) ||
+        (Ty.isPointer() && Ty.getAddressSpace() > AMDGPUAS::MAX_AMDGPU_ADDRESS))
       return Ty;
     return LLT();
   case SgprB96:

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.h

@@ -114,6 +114,7 @@ class RegBankLegalizeHelper {
   void lowerV_BFE(MachineInstr &MI);
   void lowerS_BFE(MachineInstr &MI);
   void lowerSplitTo32(MachineInstr &MI);
+  void lowerSplitTo32Select(MachineInstr &MI);
 };
 
 } // end namespace AMDGPU

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp

@@ -198,7 +198,8 @@ UniformityLLTOpPredicateID LLTToBId(LLT Ty) {
     return B32;
   if (Ty == LLT::scalar(64) || Ty == LLT::fixed_vector(2, 32) ||
       Ty == LLT::fixed_vector(4, 16) || Ty == LLT::pointer(1, 64) ||
-      Ty == LLT::pointer(4, 64))
+      Ty == LLT::pointer(4, 64) ||
+      (Ty.isPointer() && Ty.getAddressSpace() > AMDGPUAS::MAX_AMDGPU_ADDRESS))
     return B64;
   if (Ty == LLT::fixed_vector(3, 32))
     return B96;
@@ -485,8 +486,12 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
   addRulesForGOpcs({G_BR}).Any({{_}, {{}, {None}}});
 
   addRulesForGOpcs({G_SELECT}, StandardB)
+      .Any({{DivS16}, {{Vgpr16}, {Vcc, Vgpr16, Vgpr16}}})
+      .Any({{UniS16}, {{Sgpr16}, {Sgpr32AExtBoolInReg, Sgpr16, Sgpr16}}})
       .Div(B32, {{VgprB32}, {Vcc, VgprB32, VgprB32}})
-      .Uni(B32, {{SgprB32}, {Sgpr32AExtBoolInReg, SgprB32, SgprB32}});
+      .Uni(B32, {{SgprB32}, {Sgpr32AExtBoolInReg, SgprB32, SgprB32}})
+      .Div(B64, {{VgprB64}, {Vcc, VgprB64, VgprB64}, SplitTo32Select})
+      .Uni(B64, {{SgprB64}, {Sgpr32AExtBoolInReg, SgprB64, SgprB64}});
 
   addRulesForGOpcs({G_ANYEXT})
       .Any({{UniS16, S1}, {{None}, {None}}}) // should be combined away

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h

@@ -177,6 +177,7 @@ enum LoweringMethodID {
   V_BFE,
   VgprToVccCopy,
   SplitTo32,
+  SplitTo32Select,
   Ext32To64,
   UniCstExt,
   SplitLoad,