From 3709f6b81ea8e94f6a6b26d015d9216c9435edce Mon Sep 17 00:00:00 2001
From: Michael Liao <michael.hliao@gmail.com>
Date: Tue, 18 Jun 2019 17:58:49 +0000
Subject: [PATCH] [SROA] Enhance SROA to handle `addrspacecast`ed allocas

Summary:
- After `addrspacecast` is allowed to be eliminated in SROA, the
  adjusting of storage pointer (from `alloca) needs to handle the
  potential different address spaces between the storage pointer (from
  alloca) and the pointer being used.

Reviewers: arsenm

Subscribers: wdng, hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D63501

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@363711 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Scalar/SROA.cpp        |  6 ++++++
 test/Transforms/SROA/addrspacecast.ll | 15 +++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/lib/Transforms/Scalar/SROA.cpp b/lib/Transforms/Scalar/SROA.cpp
index 6454f1ea1eb..627199a27e6 100644
--- a/lib/Transforms/Scalar/SROA.cpp
+++ b/lib/Transforms/Scalar/SROA.cpp
@@ -1589,6 +1589,12 @@ static Value *getAdjustedPtr(IRBuilderTy &IRB, const DataLayout &DL, Value *Ptr,
   PointerType *TargetPtrTy = cast<PointerType>(PointerTy);
   Type *TargetTy = TargetPtrTy->getElementType();
 
+  // As `addrspacecast` is , `Ptr` (the storage pointer) may have different
+  // address space from the expected `PointerTy` (the pointer to be used).
+  // Adjust the pointer type based the original storage pointer.
+  auto AS = cast<PointerType>(Ptr->getType())->getAddressSpace();
+  PointerTy = PointerTy->getPointerTo(AS);
+
   do {
     // First fold any existing GEPs into the offset.
     while (GEPOperator *GEP = dyn_cast<GEPOperator>(Ptr)) {
diff --git a/test/Transforms/SROA/addrspacecast.ll b/test/Transforms/SROA/addrspacecast.ll
index 5edef89e8cd..d2a1cef81ae 100644
--- a/test/Transforms/SROA/addrspacecast.ll
+++ b/test/Transforms/SROA/addrspacecast.ll
@@ -299,6 +299,21 @@ define void @select_addrspacecast_gv(i1 %a, i1 %b) {
   ret void
 }
 
+; CHECK-LABEL: @select_addrspacecast_i8(
+; CHECK: [[SEL:%.*]] = select i1 undef, i8 undef, i8 undef
+; CHECK-NEXT: ret i8 [[SEL]]
+define i8 @select_addrspacecast_i8() {
+  %a = alloca i8
+  %b = alloca i8
+
+  %a.ptr = addrspacecast i8* %a to i8 addrspace(1)*
+  %b.ptr = addrspacecast i8* %b to i8 addrspace(1)*
+
+  %ptr = select i1 undef, i8 addrspace(1)* %a.ptr, i8 addrspace(1)* %b.ptr
+  %ret = load i8, i8 addrspace(1)* %ptr
+  ret i8 %ret
+}
+
 !0 = !{!1, !1, i64 0, i64 1}
 !1 = !{!2, i64 1, !"type_0"}
 !2 = !{!"root"}
-- 
2.40.0