[llvm] [InferAddressSpaces] collect flat address expression from return value (PR #70610)

[Wenju He via llvm-commits]

https://github.com/wenju-he updated https://github.com/llvm/llvm-project/pull/70610

>From 5b18d81d915b923556047adb811f38d326ccf820 Mon Sep 17 00:00:00 2001
From: Wenju He <wenju.he at intel.com>
Date: Mon, 30 Oct 2023 08:17:47 +0800
Subject: [PATCH 1/2] [InferAddressSpaces] collect flat address expression from
 return value

If function return value's type is pointer, we can try to collect flat
address expression from it.
---
 llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp        | 4 ++++
 llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll | 7 +++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 2da521375c00161..28fe1b5e75327e6 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -523,6 +523,10 @@ InferAddressSpacesImpl::collectFlatAddressExpressions(Function &F) const {
     } else if (auto *I2P = dyn_cast<IntToPtrInst>(&I)) {
       if (isNoopPtrIntCastPair(cast<Operator>(I2P), *DL, TTI))
         PushPtrOperand(cast<Operator>(I2P->getOperand(0))->getOperand(0));
+    } else if (auto *RI = dyn_cast<ReturnInst>(&I)) {
+      if (auto *RV = RI->getReturnValue();
+          RV && RV->getType()->isPtrOrPtrVectorTy())
+        PushPtrOperand(RV);
     }
   }
 
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll
index 9495c5566b36c2a..150577e41818909 100644
--- a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll
@@ -4,10 +4,9 @@
 ;  this doesn't do something insane on non-canonical IR.
 
 ; CHECK-LABEL: @return_select_group_flat(
-; CHECK-NEXT: %cast0 = addrspacecast ptr addrspace(3) %group.ptr.0 to ptr
-; CHECK-NEXT: %cast1 = addrspacecast ptr addrspace(3) %group.ptr.1 to ptr
-; CHECK-NEXT: %select = select i1 %c, ptr %cast0, ptr %cast1
-; CHECK-NEXT: ret ptr %select
+; CHECK-NEXT: [[SELECT:%.*]] = select i1 %c, ptr addrspace(3) %group.ptr.0, ptr addrspace(3) %group.ptr.1
+; CHECK-NEXT: [[TMP1:%.*]] = addrspacecast ptr addrspace(3) [[SELECT]] to ptr
+; CHECK-NEXT: ret ptr [[TMP1]]
 define ptr @return_select_group_flat(i1 %c, ptr addrspace(3) %group.ptr.0, ptr addrspace(3) %group.ptr.1) #0 {
   %cast0 = addrspacecast ptr addrspace(3) %group.ptr.0 to ptr
   %cast1 = addrspacecast ptr addrspace(3) %group.ptr.1 to ptr

>From bca6a6664119b96c6736df131ebe71a7cece2b93 Mon Sep 17 00:00:00 2001
From: Wenju He <wenju.he at intel.com>
Date: Mon, 30 Oct 2023 16:01:52 +0800
Subject: [PATCH 2/2] add ret vector test, add signature test for select

---
 .../InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll        | 7 +++++++
 llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll | 9 +++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll
index a73cd517fd49d5a..b6713e96bed3ff5 100644
--- a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll
@@ -69,6 +69,13 @@ define ptr @noop_ptrint_pair_ce2() {
   ret ptr inttoptr (i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr)
 }
 
+; COMMON-LABEL: @noop_ptrint_pair_ce2_vec(
+; AMDGCN-NEXT: ret <2 x ptr> <ptr addrspacecast (ptr addrspace(1) @g to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>
+; NOTTI-NEXT: ret <2 x ptr> <ptr inttoptr (i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>
+define <2 x ptr> @noop_ptrint_pair_ce2_vec() {
+  ret <2 x ptr> <ptr inttoptr (i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>
+}
+
 ; COMMON-LABEL: @noop_ptrint_pair_ce3(
 ; AMDGCN-NEXT: %i = inttoptr i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr
 ; AMDGCN-NEXT: ret void
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll
index 150577e41818909..ab4e1277b81fdbd 100644
--- a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll
@@ -3,10 +3,11 @@
 ; Instcombine pulls the addrspacecast out of the select, make sure
 ;  this doesn't do something insane on non-canonical IR.
 
-; CHECK-LABEL: @return_select_group_flat(
-; CHECK-NEXT: [[SELECT:%.*]] = select i1 %c, ptr addrspace(3) %group.ptr.0, ptr addrspace(3) %group.ptr.1
-; CHECK-NEXT: [[TMP1:%.*]] = addrspacecast ptr addrspace(3) [[SELECT]] to ptr
-; CHECK-NEXT: ret ptr [[TMP1]]
+; CHECK-LABEL: define ptr @return_select_group_flat(
+; CHECK-SAME: i1 [[C:%.*]], ptr addrspace(3) [[GROUP_PTR_0:%.*]], ptr addrspace(3) [[GROUP_PTR_1:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[C]], ptr addrspace(3) [[GROUP_PTR_0]], ptr addrspace(3) [[GROUP_PTR_1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr addrspace(3) [[SELECT]] to ptr
+; CHECK-NEXT:    ret ptr [[TMP1]]
 define ptr @return_select_group_flat(i1 %c, ptr addrspace(3) %group.ptr.0, ptr addrspace(3) %group.ptr.1) #0 {
   %cast0 = addrspacecast ptr addrspace(3) %group.ptr.0 to ptr
   %cast1 = addrspacecast ptr addrspace(3) %group.ptr.1 to ptr