InferAddressSpaces: Fix assert on inferred source for inttoptr/ptrtoint

If we had some source value we could infer an address space from that
went through a ptrtoint/inttoptr pair, this would fail since bitcast
can't change the address space.

Fixes issue 53665.

(cherry picked from commit 52fbb786a638ecc7349641b45b62a5abafffdf75)

2 files changed, 61 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 8f5933b..ddc747a 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -655,10 +655,13 @@ Value *InferAddressSpacesImpl::cloneInstructionWithNewAddressSpace(
   case Instruction::IntToPtr: {
     assert(isNoopPtrIntCastPair(cast<Operator>(I), *DL, TTI));
     Value *Src = cast<Operator>(I->getOperand(0))->getOperand(0);
-    assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
-    if (Src->getType() != NewPtrType)
-      return new BitCastInst(Src, NewPtrType);
-    return Src;
+    if (Src->getType() == NewPtrType)
+      return Src;
+
+    // If we had a no-op inttoptr/ptrtoint pair, we may still have inferred a
+    // source address space from a generic pointer source need to insert a cast
+    // back.
+    return CastInst::CreatePointerBitCastOrAddrSpaceCast(Src, NewPtrType);
   }
   default:
     llvm_unreachable("Unexpected opcode");
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue53665.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue53665.ll
new file mode 100644
index 0000000..fcc1f56
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/issue53665.ll
@@ -0,0 +1,54 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -infer-address-spaces -o - %s | FileCheck %s
+; https://github.com/llvm/llvm-project/issues/53665
+
+define i32 @addrspacecast_ptrtoint_inttoptr(i8 addrspace(1)* %arg) {
+; CHECK-LABEL: @addrspacecast_ptrtoint_inttoptr(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8 addrspace(1)* [[ARG:%.*]] to i32 addrspace(1)*
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32 addrspace(1)* [[TMP0]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+bb:
+  %asc = addrspacecast i8 addrspace(1)* %arg to i8*
+  %p2i = ptrtoint i8* %asc to i64
+  %i2p = inttoptr i64 %p2i to i32*
+  %load = load i32, i32* %i2p
+  ret i32 %load
+}
+
+define i32 @assumed_ptrtoint_inttoptr(i8* %arg) {
+bb:
+  %is.priv = call i1 @llvm.amdgcn.is.private(i8* %arg)
+  %not.is.priv = xor i1 %is.priv, -1
+  %is.shared = call i1 @llvm.amdgcn.is.shared(i8* %arg)
+  %not.is.shared = xor i1 %is.shared, -1
+  %and = and i1 %not.is.priv, %not.is.shared
+  tail call void @llvm.assume(i1 %and)
+  %p2i = ptrtoint i8* %arg to i64
+  %i2p = inttoptr i64 %p2i to i32*
+  %load = load i32, i32* %i2p
+  ret i32 %load
+}
+
+define i32 @addrspacecast_ptrtoint_inttptr_nontrivial(i8 addrspace(3)* %arg) {
+; CHECK-LABEL: @addrspacecast_ptrtoint_inttptr_nontrivial(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8 addrspace(3)* [[ARG:%.*]] to i32 addrspace(3)*
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32 addrspace(3)* [[TMP0]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+bb:
+  %asc = addrspacecast i8 addrspace(3)* %arg to i8*
+  %p2i = ptrtoint i8* %asc to i64
+  %i2p = inttoptr i64 %p2i to i32*
+  %load = load i32, i32* %i2p
+  ret i32 %load
+}
+
+declare void @llvm.assume(i1 noundef) #0
+declare i1 @llvm.amdgcn.is.shared(i8* nocapture) #1
+declare i1 @llvm.amdgcn.is.private(i8* nocapture) #1
+
+attributes #0 = { inaccessiblememonly nofree nosync nounwind willreturn }
+attributes #1 = { nounwind readnone speculatable willreturn }