AMDGPU: Fix assertion on load of vector of pointers (#110436)

Fix InferAddressSpaces asserting on a load of a vector of flat
pointers.

Fixes #110433

Author: arsenm

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -799,12 +799,11 @@ bool AMDGPUTargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
 
 unsigned AMDGPUTargetMachine::getAssumedAddrSpace(const Value *V) const {
   const auto *LD = dyn_cast<LoadInst>(V);
-  if (!LD)
+  if (!LD) // TODO: Handle invariant load like constant.
     return AMDGPUAS::UNKNOWN_ADDRESS_SPACE;
 
   // It must be a generic pointer loaded.
-  assert(V->getType()->isPointerTy() &&
-         V->getType()->getPointerAddressSpace() == AMDGPUAS::FLAT_ADDRESS);
+  assert(V->getType()->getPointerAddressSpace() == AMDGPUAS::FLAT_ADDRESS);
 
   const auto *Ptr = LD->getPointerOperand();
   if (Ptr->getType()->getPointerAddressSpace() != AMDGPUAS::CONSTANT_ADDRESS)

llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue110433.ll

@@ -0,0 +1,45 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=infer-address-spaces %s | FileCheck %s
+
+define <8 x i1> @load_vector_of_flat_ptr_from_constant(ptr addrspace(4) %ptr) {
+; CHECK-LABEL: define <8 x i1> @load_vector_of_flat_ptr_from_constant(
+; CHECK-SAME: ptr addrspace(4) [[PTR:%.*]]) {
+; CHECK-NEXT:    [[LD:%.*]] = load <8 x ptr>, ptr addrspace(4) [[PTR]], align 128
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast <8 x ptr> [[LD]] to <8 x ptr addrspace(1)>
+; CHECK-NEXT:    [[TMP2:%.*]] = addrspacecast <8 x ptr addrspace(1)> [[TMP1]] to <8 x ptr>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <8 x ptr> [[TMP2]], zeroinitializer
+; CHECK-NEXT:    ret <8 x i1> [[CMP]]
+;
+  %ld = load <8 x ptr>, ptr addrspace(4) %ptr, align 128
+  %cmp = icmp eq <8 x ptr> %ld, zeroinitializer
+  ret <8 x i1> %cmp
+}
+
+define <8 x i1> @load_vector_of_flat_ptr_from_global(ptr addrspace(1) %ptr) {
+; CHECK-LABEL: define <8 x i1> @load_vector_of_flat_ptr_from_global(
+; CHECK-SAME: ptr addrspace(1) [[PTR:%.*]]) {
+; CHECK-NEXT:    [[LD:%.*]] = load <8 x ptr>, ptr addrspace(1) [[PTR]], align 128
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <8 x ptr> [[LD]], zeroinitializer
+; CHECK-NEXT:    ret <8 x i1> [[CMP]]
+;
+  %ld = load <8 x ptr>, ptr addrspace(1) %ptr, align 128
+  %cmp = icmp eq <8 x ptr> %ld, zeroinitializer
+  ret <8 x i1> %cmp
+}
+
+define <8 x i1> @load_vector_of_flat_ptr_from_global_invariant(ptr addrspace(1) %ptr) {
+; CHECK-LABEL: define <8 x i1> @load_vector_of_flat_ptr_from_global_invariant(
+; CHECK-SAME: ptr addrspace(1) [[PTR:%.*]]) {
+; CHECK-NEXT:    [[LD:%.*]] = load <8 x ptr>, ptr addrspace(1) [[PTR]], align 128, !invariant [[META0:![0-9]+]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <8 x ptr> [[LD]], zeroinitializer
+; CHECK-NEXT:    ret <8 x i1> [[CMP]]
+;
+  %ld = load <8 x ptr>, ptr addrspace(1) %ptr, align 128, !invariant !0
+  %cmp = icmp eq <8 x ptr> %ld, zeroinitializer
+  ret <8 x i1> %cmp
+}
+
+!0 = !{}
+;.
+; CHECK: [[META0]] = !{}
+;.