[llvm] AMDGPU: Fix assertion on load of vector of pointers (PR #110436)

[Matt Arsenault via llvm-commits]

https://github.com/arsenm updated https://github.com/llvm/llvm-project/pull/110436

>From 3867866b9492d0248aaa59626aca1005e6399374 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Sun, 29 Sep 2024 23:49:58 +0400
Subject: [PATCH] AMDGPU: Fix assertion on load of vector of pointers

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

Fixes #110433
---
 .../lib/Target/AMDGPU/AMDGPUTargetMachine.cpp |  5 +--
 .../InferAddressSpaces/AMDGPU/issue110433.ll  | 45 +++++++++++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue110433.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index abd50748f2cc05..ce3d70332d0a67 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/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)
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue110433.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue110433.ll
new file mode 100644
index 00000000000000..4297d1c27639e5
--- /dev/null
+++ b/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]] = !{}
+;.