[SROA] Enhance SROA to handle `addrspacecast`ed allocas

[hliao]

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.

[arsenm]

This test case already works on trunk?

[arsenm]

Nevermind, there's an assert with a debug build. LGTM

llvm/lib/Transforms/Scalar/SROA.cpp

static Value *getAdjustedPtr(IRBuilderTy &IRB, const DataLayout &DL, Value *Ptr,
   // Remember any i8 pointer we come across to re-use if we need to do a raw
   // byte offset.
   Value *Int8Ptr = nullptr;
   APInt Int8PtrOffset(Offset.getBitWidth(), 0);
 
   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)) {
       APInt GEPOffset(Offset.getBitWidth(), 0);
       if (!GEP->accumulateConstantOffset(DL, GEPOffset))
         break;
       Offset += GEPOffset;
       Ptr = GEP->getPointerOperand();

llvm/test/Transforms/SROA/addrspacecast.ll

;
   %p.0.c = select i1 undef, i64* %c, i64* %c
   %asc = addrspacecast i64* %p.0.c to i64 addrspace(1)*
 
   %cond.in = select i1 undef, i64 addrspace(1)* %asc, i64 addrspace(1)* @gv
   %cond = load i64, i64 addrspace(1)* %cond.in, align 8
   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"}
 !3 = !{!4, !4, i64 0, i64 1}
 !4 = !{!2, i64 1, !"type_3"}