AMDGPU/GlobalISel: Fix RegBankSelect for GEP.

[arsenm]

Summary

This is basically a pointer typed add, so shouldn't be any different.
This was assuming everything was an SGPR, which is not true.

Also cleanup legality for GEP. I don't seem to be seeing the problem
the hack marking s64 as a legal pointer type the comment mentions.

[arsenm]

ping

[tstellar]

LGTM.

[arsenm]

r354067

lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST,
   const LLT GlobalPtr = GetAddrSpacePtr(AMDGPUAS::GLOBAL_ADDRESS);
   const LLT ConstantPtr = GetAddrSpacePtr(AMDGPUAS::CONSTANT_ADDRESS);
   const LLT LocalPtr = GetAddrSpacePtr(AMDGPUAS::LOCAL_ADDRESS);
   const LLT FlatPtr = GetAddrSpacePtr(AMDGPUAS::FLAT_ADDRESS);
   const LLT PrivatePtr = GetAddrSpacePtr(AMDGPUAS::PRIVATE_ADDRESS);
 
   const LLT CodePtr = FlatPtr;
 
-  const LLT AddrSpaces[] = {
-    GlobalPtr,
-    ConstantPtr,
-    LocalPtr,
-    FlatPtr,
-    PrivatePtr
+  const std::initializer_list<LLT> AddrSpaces64 = {
+    GlobalPtr, ConstantPtr, FlatPtr
+  };
+
+  const std::initializer_list<LLT> AddrSpaces32 = {
+    LocalPtr, PrivatePtr
   };
 
   setAction({G_BRCOND, S1}, Legal);
 
   getActionDefinitionsBuilder({G_ADD, G_SUB, G_MUL, G_UMULH, G_SMULH})
     .legalFor({S32})
     .scalarize(0);
 
AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST,
 
   getActionDefinitionsBuilder({G_FPTOSI, G_FPTOUI})
     .legalFor({{S32, S32}, {S32, S64}})
     .scalarize(0);
 
   getActionDefinitionsBuilder({G_INTRINSIC_TRUNC, G_INTRINSIC_ROUND})
     .legalFor({S32, S64});
 
-  for (LLT PtrTy : AddrSpaces) {
-    LLT IdxTy = LLT::scalar(PtrTy.getSizeInBits());
-    setAction({G_GEP, PtrTy}, Legal);
-    setAction({G_GEP, 1, IdxTy}, Legal);
-  }
 
-  // FIXME: When RegBankSelect inserts copies, it will only create new registers
-  // with scalar types. This means we can end up with G_LOAD/G_STORE/G_GEP
-  // instruction with scalar types for their pointer operands. In assert builds,
-  // the instruction selector will assert if it sees a generic instruction which
-  // isn't legal, so we need to tell it that scalar types are legal for pointer
-  // operands
-  setAction({G_GEP, S64}, Legal);
+  getActionDefinitionsBuilder(G_GEP)
+    .legalForCartesianProduct(AddrSpaces64, {S64})
+    .legalForCartesianProduct(AddrSpaces32, {S32})
+    .scalarize(0);
 
   setAction({G_BLOCK_ADDR, CodePtr}, Legal);
 
   getActionDefinitionsBuilder(G_ICMP)
     .legalForCartesianProduct(
       {S1}, {S32, S64, GlobalPtr, LocalPtr, ConstantPtr, PrivatePtr, FlatPtr})
     .legalFor({{S1, S32}, {S1, S64}})
     .widenScalarToNextPow2(1)
auto greaterThan = [](unsigned TypeIdx0, unsigned TypeIdx1) {
     return [=](const LegalityQuery &Query) {
       return Query.Types[TypeIdx0].getSizeInBits() >
              Query.Types[TypeIdx1].getSizeInBits();
     };
   };
 
   getActionDefinitionsBuilder(G_INTTOPTR)
     // List the common cases
-    .legalForCartesianProduct({GlobalPtr, ConstantPtr, FlatPtr}, {S64})
-    .legalForCartesianProduct({LocalPtr, PrivatePtr}, {S32})
+    .legalForCartesianProduct(AddrSpaces64, {S64})
+    .legalForCartesianProduct(AddrSpaces32, {S32})
     .scalarize(0)
     // Accept any address space as long as the size matches
     .legalIf(sameSize(0, 1))
     .widenScalarIf(smallerThan(1, 0),
       [](const LegalityQuery &Query) {
         return std::make_pair(1, LLT::scalar(Query.Types[0].getSizeInBits()));
       })
     .narrowScalarIf(greaterThan(1, 0),
       [](const LegalityQuery &Query) {
         return std::make_pair(1, LLT::scalar(Query.Types[0].getSizeInBits()));
       });
 
   getActionDefinitionsBuilder(G_PTRTOINT)
     // List the common cases
-    .legalForCartesianProduct({GlobalPtr, ConstantPtr, FlatPtr}, {S64})
-    .legalForCartesianProduct({LocalPtr, PrivatePtr}, {S32})
+    .legalForCartesianProduct(AddrSpaces64, {S64})
+    .legalForCartesianProduct(AddrSpaces32, {S32})
     .scalarize(0)
     // Accept any address space as long as the size matches
     .legalIf(sameSize(0, 1))
     .widenScalarIf(smallerThan(0, 1),
       [](const LegalityQuery &Query) {
         return std::make_pair(0, LLT::scalar(Query.Types[1].getSizeInBits()));
       })
     .narrowScalarIf(

lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

if (Size == 64) {
       }
 
       break;
     }
 
     LLVM_FALLTHROUGH;
   }
 
+  case AMDGPU::G_GEP:
   case AMDGPU::G_ADD:
   case AMDGPU::G_SUB:
   case AMDGPU::G_MUL:
   case AMDGPU::G_SHL:
   case AMDGPU::G_LSHR:
   case AMDGPU::G_ASHR:
   case AMDGPU::G_UADDO:
   case AMDGPU::G_SADDO:
case AMDGPU::G_FCMP: {
     unsigned Size = MRI.getType(MI.getOperand(2).getReg()).getSizeInBits();
     unsigned Op2Bank = getRegBankID(MI.getOperand(2).getReg(), MRI, *TRI);
     OpdsMapping[0] = AMDGPU::getValueMapping(AMDGPU::VCCRegBankID, 1);
     OpdsMapping[1] = nullptr; // Predicate Operand.
     OpdsMapping[2] = AMDGPU::getValueMapping(Op2Bank, Size);
     OpdsMapping[3] = AMDGPU::getValueMapping(AMDGPU::VGPRRegBankID, Size);
     break;
   }
-  case AMDGPU::G_GEP: {
-    for (unsigned i = 0, e = MI.getNumOperands(); i != e; ++i) {
-      if (!MI.getOperand(i).isReg())
-        continue;
-
-      unsigned Size = MRI.getType(MI.getOperand(i).getReg()).getSizeInBits();
-      OpdsMapping[i] = AMDGPU::getValueMapping(AMDGPU::SGPRRegBankID, Size);
-    }
-    break;
-  }
   case AMDGPU::G_STORE: {
     assert(MI.getOperand(0).isReg());
     unsigned Size = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
     // FIXME: We need to specify a different reg bank once scalar stores
     // are supported.
     const ValueMapping *ValMapping =
         AMDGPU::getValueMapping(AMDGPU::VGPRRegBankID, Size);
     // FIXME: Depending on the type of store, the pointer could be in

test/CodeGen/AMDGPU/GlobalISel/regbankselect-gep.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -run-pass=regbankselect %s -verify-machineinstrs -o - -regbankselect-fast | FileCheck %s
+# RUN: llc -march=amdgcn -run-pass=regbankselect %s -verify-machineinstrs -o - -regbankselect-greedy | FileCheck %s
+
+---
+name: gep_p1_s_k
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0_sgpr1
+
+    ; CHECK-LABEL: name: gep_p1_s_k
+    ; CHECK: [[COPY:%[0-9]+]]:sgpr(p1) = COPY $sgpr0_sgpr1
+    ; CHECK: [[C:%[0-9]+]]:sgpr(s64) = G_CONSTANT i64 1
+    ; CHECK: [[GEP:%[0-9]+]]:sgpr(p1) = G_GEP [[COPY]], [[C]](s64)
+    %0:_(p1) = COPY $sgpr0_sgpr1
+    %1:_(s64) = G_CONSTANT i64 1
+    %2:_(p1) = G_GEP %0, %1
+...
+
+---
+name: gep_p1_s_s
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0_sgpr1, $sgpr2_sgpr3
+
+    ; CHECK-LABEL: name: gep_p1_s_s
+    ; CHECK: [[COPY:%[0-9]+]]:sgpr(p1) = COPY $sgpr0_sgpr1
+    ; CHECK: [[COPY1:%[0-9]+]]:sgpr(s64) = COPY $sgpr2_sgpr3
+    ; CHECK: [[GEP:%[0-9]+]]:sgpr(p1) = G_GEP [[COPY]], [[COPY1]](s64)
+    %0:_(p1) = COPY $sgpr0_sgpr1
+    %1:_(s64) = COPY $sgpr2_sgpr3
+    %2:_(p1) = G_GEP %0, %1
+...
+
+---
+name: gep_p1_v_k
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+
+    ; CHECK-LABEL: name: gep_p1_v_k
+    ; CHECK: [[COPY:%[0-9]+]]:vgpr(p1) = COPY $vgpr0_vgpr1
+    ; CHECK: [[C:%[0-9]+]]:sgpr(s64) = G_CONSTANT i64 1
+    ; CHECK: [[COPY1:%[0-9]+]]:vgpr(s64) = COPY [[C]](s64)
+    ; CHECK: [[GEP:%[0-9]+]]:vgpr(p1) = G_GEP [[COPY]], [[COPY1]](s64)
+    %0:_(p1) = COPY $vgpr0_vgpr1
+    %1:_(s64) = G_CONSTANT i64 1
+    %2:_(p1) = G_GEP %0, %1
+...
+
+---
+name: gep_p1_v_s
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1, $sgpr0_sgpr1
+
+    ; CHECK-LABEL: name: gep_p1_v_s
+    ; CHECK: [[COPY:%[0-9]+]]:vgpr(p1) = COPY $vgpr0_vgpr1
+    ; CHECK: [[COPY1:%[0-9]+]]:sgpr(s64) = COPY $sgpr0_sgpr1
+    ; CHECK: [[COPY2:%[0-9]+]]:vgpr(s64) = COPY [[COPY1]](s64)
+    ; CHECK: [[GEP:%[0-9]+]]:vgpr(p1) = G_GEP [[COPY]], [[COPY2]](s64)
+    %0:_(p1) = COPY $vgpr0_vgpr1
+    %1:_(s64) = COPY $sgpr0_sgpr1
+    %2:_(p1) = G_GEP %0, %1
+...
+
+---
+name: gep_p1_v_v
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1, $vgpr2_vgpr3
+
+    ; CHECK-LABEL: name: gep_p1_v_v
+    ; CHECK: [[COPY:%[0-9]+]]:vgpr(p1) = COPY $vgpr0_vgpr1
+    ; CHECK: [[COPY1:%[0-9]+]]:vgpr(s64) = COPY $vgpr2_vgpr3
+    ; CHECK: [[GEP:%[0-9]+]]:vgpr(p1) = G_GEP [[COPY]], [[COPY1]](s64)
+    %0:_(p1) = COPY $vgpr0_vgpr1
+    %1:_(s64) = COPY $vgpr2_vgpr3
+    %2:_(p1) = G_GEP %0, %1
+...