[InferAddressSpaces] collect flat address expression from return value

[wenju-he]

If function return value's type is pointer, we can try to collect flat
address expression from it.
This PR also fixes noop_ptrint_pair_ce2 in noop-ptrint-pair.ll in #70611

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

@llvm/pr-subscribers-llvm-transforms

Author: Wenju He (wenju-he)

Changes

If function return value's type is pointer, we can try to collect flat
address expression from it.

---

Full diff: https://github.com/llvm/llvm-project/pull/70610.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp (+4)
• (modified) llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll (+3-4)

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

[wenju-he]

@arsenm could you please help to review? I don't have permission to add you as reviewer.

[arsenm]

Can you also add a test with a vector of pointers

[wenju-he]

I added noop_ptrint_pair_ce2_vec to noop-ptrint-pair.ll
I find that currently isAddressExpression isn't handling ConstantAggregate so vector of pointers is not inferred. The change at line 73 of bca6a66 is just because currently constant is replaced globally, i.e. it is replaced while handling another function.

[arsenm]

I think this also should check the signature of the function; need to make sure you aren't mutating it in a way that breaks uses

[wenju-he]

done