AMDGPU/GlobalISel: Fix crash on wide constant load with VGPR pointer

[arsenm]

Summary

This was ignoring the register bank of the input pointer, and
isUniformMMO seems overly aggressive.

This will now conservatively assume a VGPR in cases where the incoming
bank hasn't been determined yet (i.e. is from a loop phi).

[nhaehnle]

test/CodeGen/AMDGPU/GlobalISel/legalize-extract.mir

1118–1131 ↗

(On Diff #223669)

This change seems unrelated.

[arsenm]

Fix wrong test

[nhaehnle]

I do have one question. Apart from that it LGTM.

lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

326–327

You mean because isUniformMMO returns true if the MMO doesn't have a pointer? There's a comment in that function which justifies that (though I'm not sure whether that comment is correct).

[arsenm]

lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

326–327

The comment there isn't entirely wrong, but also isn't entirely correct. There can also be null without a PSV. I don't think this would ever happen in a real program, and is more a MIR semantics question.

[arsenm]

r374255

lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

case Intrinsic::amdgcn_s_sendmsghalt: {
     const std::array<unsigned, 1> RegSrcOpIdx = { { 2 } };
     return addMappingFromTable<1>(MI, MRI, RegSrcOpIdx, makeArrayRef(Table));
   }
   default:
     return RegisterBankInfo::getInstrAlternativeMappings(MI);
   }
 }
 
+// FIXME: Returns uniform if there's no source value information. This is
+// probably wrong.
 static bool isInstrUniformNonExtLoadAlign4(const MachineInstr &MI) {
   if (!MI.hasOneMemOperand())
     return false;
 
   const MachineMemOperand *MMO = *MI.memoperands_begin();
   return MMO->getSize() >= 4 && MMO->getAlignment() >= 4 &&
          AMDGPUInstrInfo::isUniformMMO(MMO);
 }
bool AMDGPURegisterBankInfo::applyMappingWideLoad(MachineInstr &MI,
   // 128-bit loads are supported for all instruction types.
   if (LoadSize <= MaxNonSmrdLoadSize)
     return false;
 
   SmallVector<unsigned, 16> DefRegs(OpdMapper.getVRegs(0));
   SmallVector<unsigned, 1> SrcRegs(OpdMapper.getVRegs(1));
 
   // If the pointer is an SGPR, we have nothing to do.
-  if (SrcRegs.empty())
+  if (SrcRegs.empty()) {
+    Register PtrReg = MI.getOperand(1).getReg();
+    const RegisterBank *PtrBank = getRegBank(PtrReg, MRI, *TRI);
+    if (PtrBank == &AMDGPU::SGPRRegBank)
       return false;
+    SrcRegs.push_back(PtrReg);
+  }
 
   assert(LoadSize % MaxNonSmrdLoadSize == 0);
 
   // We want to get the repair instruction now, because it will help us
   // determine which instruction the legalizer inserts that will also
   // write to DstReg.
   MachineInstr *RepairInst = getOtherVRegDef(MRI, DstReg, MI);
 
AMDGPURegisterBankInfo::getImageMapping(const MachineRegisterInfo &MRI,
   return getInstructionMapping(1, 1, getOperandsMapping(OpdsMapping), NumOps);
 }
 
 const RegisterBankInfo::InstructionMapping &
 AMDGPURegisterBankInfo::getInstrMappingForLoad(const MachineInstr &MI) const {
 
   const MachineFunction &MF = *MI.getParent()->getParent();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
-  SmallVector<const ValueMapping*, 8> OpdsMapping(MI.getNumOperands());
+  SmallVector<const ValueMapping*, 2> OpdsMapping(2);
   unsigned Size = getSizeInBits(MI.getOperand(0).getReg(), MRI, *TRI);
   LLT LoadTy = MRI.getType(MI.getOperand(0).getReg());
   Register PtrReg = MI.getOperand(1).getReg();
   LLT PtrTy = MRI.getType(PtrReg);
   unsigned AS = PtrTy.getAddressSpace();
   unsigned PtrSize = PtrTy.getSizeInBits();
 
   const ValueMapping *ValMapping;
   const ValueMapping *PtrMapping;
 
-  if ((AS != AMDGPUAS::LOCAL_ADDRESS && AS != AMDGPUAS::REGION_ADDRESS &&
+  const RegisterBank *PtrBank = getRegBank(PtrReg, MRI, *TRI);
+
+  if (PtrBank == &AMDGPU::SGPRRegBank &&
+      (AS != AMDGPUAS::LOCAL_ADDRESS && AS != AMDGPUAS::REGION_ADDRESS &&
        AS != AMDGPUAS::PRIVATE_ADDRESS) &&
       isInstrUniformNonExtLoadAlign4(MI)) {
     // We have a uniform instruction so we want to use an SMRD load
     ValMapping = AMDGPU::getValueMapping(AMDGPU::SGPRRegBankID, Size);
     PtrMapping = AMDGPU::getValueMapping(AMDGPU::SGPRRegBankID, PtrSize);
   } else {
     ValMapping = AMDGPU::getValueMappingLoadSGPROnly(AMDGPU::VGPRRegBankID, LoadTy);
     PtrMapping = AMDGPU::getValueMapping(AMDGPU::VGPRRegBankID, PtrSize);

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

--- |
   define amdgpu_kernel void @extload_constant_i8_to_i32_uniform() { ret void }
   define amdgpu_kernel void @extload_global_i8_to_i32_uniform() { ret void }
   define amdgpu_kernel void @extload_constant_i16_to_i32_uniform() { ret void }
   define amdgpu_kernel void @extload_global_i16_to_i32_uniform() { ret void }
   define amdgpu_kernel void @load_constant_i32_uniform_align4() {ret void}
   define amdgpu_kernel void @load_constant_i32_uniform_align2() {ret void}
   define amdgpu_kernel void @load_constant_i32_uniform_align1() {ret void}
   define amdgpu_kernel void @load_private_uniform_sgpr_i32() {ret void}
+  define amdgpu_kernel void @load_constant_v8i32_vgpr_crash() { ret void }
+  define amdgpu_kernel void @load_constant_v8i32_vgpr_crash_loop_phi() { ret void }
 
   declare i32 @llvm.amdgcn.workitem.id.x() #0
   attributes #0 = { nounwind readnone }
 ...
 
 ---
 name: load_global_v8i32_non_uniform
 legalized: true
bb.0:
 
     ; CHECK-LABEL: name: load_private_uniform_sgpr_i32
     ; CHECK: [[COPY:%[0-9]+]]:sgpr(p5) = COPY $sgpr0
     ; CHECK: [[COPY1:%[0-9]+]]:vgpr(p5) = COPY [[COPY]](p5)
     ; CHECK: [[LOAD:%[0-9]+]]:vgpr(s32) = G_LOAD [[COPY1]](p5) :: (load 4, addrspace 5)
     %0:_(p5) = COPY $sgpr0
     %1:_(s32) = G_LOAD %0 :: (load 4, addrspace 5, align 4)
 ...
+
+---
+name: load_constant_v8i32_vgpr_crash
+legalized: true
+tracksRegLiveness: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+
+    ; CHECK-LABEL: name: load_constant_v8i32_vgpr_crash
+    ; CHECK: %0:vgpr(p4) = COPY $vgpr0_vgpr1
+    ; CHECK: vgpr(<4 x s32>) = G_LOAD %0(p4)
+    ; CHECK: vgpr(<4 x s32>) = G_LOAD
+    ; CHECK: G_CONCAT_VECTORS
+    %0:_(p4) = COPY $vgpr0_vgpr1
+    %1:_(<8 x s32>) = G_LOAD %0 :: (load 32, addrspace 4)
+...
+
+---
+name: load_constant_v8i32_vgpr_crash_loop_phi
+legalized: true
+tracksRegLiveness: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0_sgpr1, $sgpr2_sgpr3
+
+    ; CHECK-LABEL: name: load_constant_v8i32_vgpr_crash_loop_phi
+    ; CHECK: G_PHI
+    ; CHECK: vgpr(<4 x s32>) = G_LOAD
+    ; CHECK: vgpr(<4 x s32>) = G_LOAD
+    ; CHECK: G_CONCAT_VECTORS
+
+    %0:_(p4) = COPY $sgpr0_sgpr1
+    %1:_(p4) = COPY $sgpr2_sgpr3
+    G_BR %bb.1
+
+  bb.1:
+    %2:_(p4) = G_PHI %0, %bb.0, %4, %bb.1
+    %3:_(<8 x s32>) = G_LOAD %2 :: (load 32, addrspace 4)
+    %4:_(p4) = COPY %1
+    G_BR %bb.1
+...