AMDGPU/GlobalISel: Add regbanklegalize rules for load and store (#153176)

Cover all the missing cases and add very detailed tests for each rule.
In summary:
- Flat and Scratch, addrspace(0) and addrspace(5), loads are always
  divergent.
- Global and Constant, addrspace(1) and addrspace(4), have real uniform
  loads, s_load, but require additional checks for align and flags in mmo.
  For not natural align or not uniform mmo do uniform-in-vgpr lowering.
- Private, addrspace(3), only has instructions for divergent load, for
  uniform do uniform-in-vgpr lowering.
- Store rules are simplified using Ptr32 and Ptr64.
  All operands need to be vgpr.

Some tests have code size regression since they use more sgpr instructions,
marked with FixMe comment to get back to later.

Author: petar-avramovic

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp

@@ -352,6 +352,32 @@ void RegBankLegalizeHelper::widenLoad(MachineInstr &MI, LLT WideTy,
   MI.eraseFromParent();
 }
 
+void RegBankLegalizeHelper::widenMMOToS32(GAnyLoad &MI) const {
+  Register Dst = MI.getDstReg();
+  Register Ptr = MI.getPointerReg();
+  MachineMemOperand &MMO = MI.getMMO();
+  unsigned MemSize = 8 * MMO.getSize().getValue();
+
+  MachineMemOperand *WideMMO = B.getMF().getMachineMemOperand(&MMO, 0, S32);
+
+  if (MI.getOpcode() == G_LOAD) {
+    B.buildLoad(Dst, Ptr, *WideMMO);
+  } else {
+    auto Load = B.buildLoad(SgprRB_S32, Ptr, *WideMMO);
+
+    if (MI.getOpcode() == G_ZEXTLOAD) {
+      APInt Mask = APInt::getLowBitsSet(S32.getSizeInBits(), MemSize);
+      auto MaskCst = B.buildConstant(SgprRB_S32, Mask);
+      B.buildAnd(Dst, Load, MaskCst);
+    } else {
+      assert(MI.getOpcode() == G_SEXTLOAD);
+      B.buildSExtInReg(Dst, Load, MemSize);
+    }
+  }
+
+  MI.eraseFromParent();
+}
+
 void RegBankLegalizeHelper::lowerVccExtToSel(MachineInstr &MI) {
   Register Dst = MI.getOperand(0).getReg();
   LLT Ty = MRI.getType(Dst);
@@ -744,6 +770,8 @@ void RegBankLegalizeHelper::lower(MachineInstr &MI,
     }
     break;
   }
+  case WidenMMOToS32:
+    return widenMMOToS32(cast<GAnyLoad>(MI));
   }
 
   if (!WaterfallSgprs.empty()) {
@@ -759,6 +787,7 @@ LLT RegBankLegalizeHelper::getTyFromID(RegBankLLTMappingApplyID ID) {
     return LLT::scalar(1);
   case Sgpr16:
   case Vgpr16:
+  case UniInVgprS16:
     return LLT::scalar(16);
   case Sgpr32:
   case Sgpr32_WF:
@@ -895,6 +924,7 @@ RegBankLegalizeHelper::getRegBankFromID(RegBankLLTMappingApplyID ID) {
   case SgprB256:
   case SgprB512:
   case UniInVcc:
+  case UniInVgprS16:
   case UniInVgprS32:
   case UniInVgprV2S16:
   case UniInVgprV4S32:
@@ -1015,6 +1045,18 @@ void RegBankLegalizeHelper::applyMappingDst(
       B.buildTrunc(Reg, CopyS32_Vcc);
       break;
     }
+    case UniInVgprS16: {
+      assert(Ty == getTyFromID(MethodIDs[OpIdx]));
+      assert(RB == SgprRB);
+      Register NewVgprDstS16 = MRI.createVirtualRegister({VgprRB, S16});
+      Register NewVgprDstS32 = MRI.createVirtualRegister({VgprRB, S32});
+      Register NewSgprDstS32 = MRI.createVirtualRegister({SgprRB, S32});
+      Op.setReg(NewVgprDstS16);
+      B.buildAnyExt(NewVgprDstS32, NewVgprDstS16);
+      buildReadAnyLane(B, NewSgprDstS32, NewVgprDstS32, RBI);
+      B.buildTrunc(Reg, NewSgprDstS32);
+      break;
+    }
     case UniInVgprS32:
     case UniInVgprV2S16:
     case UniInVgprV4S32: {

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.h

@@ -11,6 +11,7 @@
 
 #include "AMDGPURegBankLegalizeRules.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 
 namespace llvm {
@@ -107,6 +108,7 @@ class RegBankLegalizeHelper {
   void splitLoad(MachineInstr &MI, ArrayRef<LLT> LLTBreakdown,
                  LLT MergeTy = LLT());
   void widenLoad(MachineInstr &MI, LLT WideTy, LLT MergeTy = LLT());
+  void widenMMOToS32(GAnyLoad &MI) const;
 
   void lower(MachineInstr &MI, const RegBankLLTMapping &Mapping,
              SmallSet<Register, 4> &SgprWaterfallOperandRegs);

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp

@@ -176,6 +176,7 @@ enum RegBankLLTMappingApplyID {
 
   // Dst only modifiers: read-any-lane and truncs
   UniInVcc,
+  UniInVgprS16,
   UniInVgprS32,
   UniInVgprV2S16,
   UniInVgprV4S32,
@@ -221,6 +222,7 @@ enum LoweringMethodID {
   UniCstExt,
   SplitLoad,
   WidenLoad,
+  WidenMMOToS32
 };
 
 enum FastRulesTypes {

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h

@@ -1,7 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,GFX7 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=GCN,GFX8 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,GFX7 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=GCN,GFX8 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
 
 define i8 @atomic_load_flat_monotonic_i8(ptr %ptr) {
 ; GCN-LABEL: atomic_load_flat_monotonic_i8:

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_load_flat.ll

@@ -1,8 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=tahiti < %s | FileCheck -check-prefixes=GCN,GFX6 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,GFX7 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=GCN,GFX8 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=tahiti < %s | FileCheck -check-prefixes=GCN,GFX6 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,GFX7 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=GCN,GFX8 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
 
 define i8 @atomic_load_global_monotonic_i8(ptr addrspace(1) %ptr) {
 ; GFX6-LABEL: atomic_load_global_monotonic_i8:

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_load_global.ll

@@ -1,5 +1,5 @@
-; RUN: llc -global-isel -global-isel-abort=0 -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,CI %s
-; RUN: llc -global-isel -global-isel-abort=0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
+; RUN: llc -global-isel -new-reg-bank-select -global-isel-abort=0 -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,CI %s
+; RUN: llc -global-isel -new-reg-bank-select -global-isel-abort=0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
 
 ; GCN-LABEL: {{^}}atomic_load_monotonic_i8:
 ; GCN: s_waitcnt

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_load_local.ll

@@ -1,7 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,GFX7 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=GCN,GFX8 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,GFX7 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=GCN,GFX8 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
 
 ; TODO: Merge with atomic_load_local.ll
 

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_load_local_2.ll

@@ -1,5 +1,5 @@
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,CI %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=kaveri < %s | FileCheck -check-prefixes=GCN,CI %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,GFX9 %s
 
 ; GCN-LABEL: {{^}}atomic_store_monotonic_i8:
 ; GCN: s_waitcnt

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_store_local.ll

@@ -1,21 +1,21 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -global-isel -mtriple=amdgcn -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,WAVE64 %s
-; RUN: llc -global-isel -mtriple=amdgcn -mcpu=gfx1031 < %s | FileCheck -check-prefixes=GCN,WAVE32 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,WAVE64 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn -mcpu=gfx1031 < %s | FileCheck -check-prefixes=GCN,WAVE32 %s
 
 ; End to end tests for scalar vs. vector boolean legalization strategies.
 
 define amdgpu_ps float @select_vgpr_sgpr_trunc_cond(i32 inreg %a, i32 %b, i32 %c) {
 ; WAVE64-LABEL: select_vgpr_sgpr_trunc_cond:
 ; WAVE64:       ; %bb.0:
-; WAVE64-NEXT:    s_and_b32 s0, 1, s0
-; WAVE64-NEXT:    v_cmp_ne_u32_e64 vcc, 0, s0
+; WAVE64-NEXT:    s_cmp_lg_u32 s0, 0
+; WAVE64-NEXT:    s_cselect_b64 vcc, exec, 0
 ; WAVE64-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
 ; WAVE64-NEXT:    ; return to shader part epilog
 ;
 ; WAVE32-LABEL: select_vgpr_sgpr_trunc_cond:
 ; WAVE32:       ; %bb.0:
-; WAVE32-NEXT:    s_and_b32 s0, 1, s0
-; WAVE32-NEXT:    v_cmp_ne_u32_e64 vcc_lo, 0, s0
+; WAVE32-NEXT:    s_cmp_lg_u32 s0, 0
+; WAVE32-NEXT:    s_cselect_b32 vcc_lo, exec_lo, 0
 ; WAVE32-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
 ; WAVE32-NEXT:    ; return to shader part epilog
   %cc = trunc i32 %a to i1
@@ -28,16 +28,16 @@ define amdgpu_ps float @select_vgpr_sgpr_trunc_and_cond(i32 inreg %a.0, i32 inre
 ; WAVE64-LABEL: select_vgpr_sgpr_trunc_and_cond:
 ; WAVE64:       ; %bb.0:
 ; WAVE64-NEXT:    s_and_b32 s0, s0, s1
-; WAVE64-NEXT:    s_and_b32 s0, 1, s0
-; WAVE64-NEXT:    v_cmp_ne_u32_e64 vcc, 0, s0
+; WAVE64-NEXT:    s_cmp_lg_u32 s0, 0
+; WAVE64-NEXT:    s_cselect_b64 vcc, exec, 0
 ; WAVE64-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
 ; WAVE64-NEXT:    ; return to shader part epilog
 ;
 ; WAVE32-LABEL: select_vgpr_sgpr_trunc_and_cond:
 ; WAVE32:       ; %bb.0:
 ; WAVE32-NEXT:    s_and_b32 s0, s0, s1
-; WAVE32-NEXT:    s_and_b32 s0, 1, s0
-; WAVE32-NEXT:    v_cmp_ne_u32_e64 vcc_lo, 0, s0
+; WAVE32-NEXT:    s_cmp_lg_u32 s0, 0
+; WAVE32-NEXT:    s_cselect_b32 vcc_lo, exec_lo, 0
 ; WAVE32-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
 ; WAVE32-NEXT:    ; return to shader part epilog
   %cc.0 = trunc i32 %a.0 to i1