[SROA] Fix APInt size when alloca address space is not 0

[yaxunl]

Summary

SROA assumes alloca address space is 0, which causes assertion. This patch fixes that.

[arsenm]

You can add the test to test/SROA/alloca-address-space.ll

[yaxunl]

remove test non-zero-alloca-addr-space.ll and add test to address-spaces.ll

[yaxunl]

Ping

[arsenm]

test/Transforms/SROA/address-spaces.ll

2

I think this test should be left as is. Versions of this test with the non-0 alloca addrspace belong in a separate test file

2

I meant to merge it with the test that already changes the address space alloca, not this one

[yaxunl]

move the test to another one.

[arsenm]

LGTM

lib/Transforms/Scalar/SROA.cpp

     IRB.SetInsertPoint(LI);
 
     DEBUG(dbgs() << "  Splitting load: " << *LI << "\n");
 
     uint64_t PartOffset = 0, PartSize = Offsets.Splits.front();
     int Idx = 0, Size = Offsets.Splits.size();
     for (;;) {
       auto *PartTy = Type::getIntNTy(Ty->getContext(), PartSize * 8);
-      auto *PartPtrTy = PartTy->getPointerTo(LI->getPointerAddressSpace());
+      auto AS = LI->getPointerAddressSpace();
+      auto *PartPtrTy = PartTy->getPointerTo(AS);
       LoadInst *PLoad = IRB.CreateAlignedLoad(
           getAdjustedPtr(IRB, DL, BasePtr,
-                         APInt(DL.getPointerSizeInBits(), PartOffset),
+                         APInt(DL.getPointerSizeInBits(AS), PartOffset),
                          PartPtrTy, BasePtr->getName() + "."),
           getAdjustedAlignment(LI, PartOffset, DL), /*IsVolatile*/ false,
           LI->getName());
       PLoad->copyMetadata(*LI, LLVMContext::MD_mem_parallel_loop_access); 
 
       // Append this load onto the list of split loads so we can find it later
       // to rewrite the stores.
       SplitLoads.push_back(PLoad);

test/Transforms/SROA/address-spaces.ll

 ; RUN: opt < %s -sroa -S | FileCheck %s
-target datalayout = "e-p:64:64:64-p1:16:16:16-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-n8:16:32:64"
+target datalayout = "e-p:64:64:64-p1:16:16:16-p5:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-n8:16:32:64-A5"
 
-declare void @llvm.memcpy.p0i8.p0i8.i32(i8* nocapture, i8* nocapture readonly, i32, i32, i1)
-declare void @llvm.memcpy.p1i8.p0i8.i32(i8 addrspace(1)* nocapture, i8* nocapture readonly, i32, i32, i1)
-declare void @llvm.memcpy.p0i8.p1i8.i32(i8* nocapture, i8 addrspace(1)* nocapture readonly, i32, i32, i1)
+declare void @llvm.memcpy.p0i8.p5i8.i32(i8* nocapture, i8 addrspace(5)* nocapture readonly, i32, i32, i1)
+declare void @llvm.memcpy.p1i8.p5i8.i32(i8 addrspace(1)* nocapture, i8 addrspace(5)* nocapture readonly, i32, i32, i1)
+declare void @llvm.memcpy.p5i8.p0i8.i32(i8 addrspace(5)* nocapture, i8* nocapture readonly, i32, i32, i1)
+declare void @llvm.memcpy.p5i8.p1i8.i32(i8 addrspace(5)* nocapture, i8 addrspace(1)* nocapture readonly, i32, i32, i1)
 declare void @llvm.memcpy.p1i8.p1i8.i32(i8 addrspace(1)* nocapture, i8 addrspace(1)* nocapture readonly, i32, i32, i1)
 
 
 ; Make sure an illegal bitcast isn't introduced
 define void @test_address_space_1_1(<2 x i64> addrspace(1)* %a, i16 addrspace(1)* %b) {
 ; CHECK-LABEL: @test_address_space_1_1(
 ; CHECK: load <2 x i64>, <2 x i64> addrspace(1)* %a, align 2
 ; CHECK: store <2 x i64> {{.*}}, <2 x i64> addrspace(1)* {{.*}}, align 2
 ; CHECK: ret void
-  %aa = alloca <2 x i64>, align 16
+  %aa = alloca <2 x i64>, align 16, addrspace(5)
   %aptr = bitcast <2 x i64> addrspace(1)* %a to i8 addrspace(1)*
-  %aaptr = bitcast <2 x i64>* %aa to i8*
-  call void @llvm.memcpy.p0i8.p1i8.i32(i8* %aaptr, i8 addrspace(1)* %aptr, i32 16, i32 2, i1 false)
+  %aaptr = bitcast <2 x i64> addrspace(5)* %aa to i8 addrspace(5)*
+  call void @llvm.memcpy.p5i8.p1i8.i32(i8 addrspace(5)* %aaptr, i8 addrspace(1)* %aptr, i32 16, i32 2, i1 false)
   %bptr = bitcast i16 addrspace(1)* %b to i8 addrspace(1)*
-  call void @llvm.memcpy.p1i8.p0i8.i32(i8 addrspace(1)* %bptr, i8* %aaptr, i32 16, i32 2, i1 false)
+  call void @llvm.memcpy.p1i8.p5i8.i32(i8 addrspace(1)* %bptr, i8 addrspace(5)* %aaptr, i32 16, i32 2, i1 false)
   ret void
 }
 
 define void @test_address_space_1_0(<2 x i64> addrspace(1)* %a, i16* %b) {
 ; CHECK-LABEL: @test_address_space_1_0(
 ; CHECK: load <2 x i64>, <2 x i64> addrspace(1)* %a, align 2
 ; CHECK: store <2 x i64> {{.*}}, <2 x i64>* {{.*}}, align 2
 ; CHECK: ret void
-  %aa = alloca <2 x i64>, align 16
+  %aa = alloca <2 x i64>, align 16, addrspace(5)
   %aptr = bitcast <2 x i64> addrspace(1)* %a to i8 addrspace(1)*
-  %aaptr = bitcast <2 x i64>* %aa to i8*
-  call void @llvm.memcpy.p0i8.p1i8.i32(i8* %aaptr, i8 addrspace(1)* %aptr, i32 16, i32 2, i1 false)
+  %aaptr = bitcast <2 x i64> addrspace(5)* %aa to i8 addrspace(5)*
+  call void @llvm.memcpy.p5i8.p1i8.i32(i8 addrspace(5)* %aaptr, i8 addrspace(1)* %aptr, i32 16, i32 2, i1 false)
   %bptr = bitcast i16* %b to i8*
-  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %bptr, i8* %aaptr, i32 16, i32 2, i1 false)
+  call void @llvm.memcpy.p0i8.p5i8.i32(i8* %bptr, i8 addrspace(5)* %aaptr, i32 16, i32 2, i1 false)
   ret void
 }
 
 define void @test_address_space_0_1(<2 x i64>* %a, i16 addrspace(1)* %b) {
 ; CHECK-LABEL: @test_address_space_0_1(
 ; CHECK: load <2 x i64>, <2 x i64>* %a, align 2
 ; CHECK: store <2 x i64> {{.*}}, <2 x i64> addrspace(1)* {{.*}}, align 2
 ; CHECK: ret void
-  %aa = alloca <2 x i64>, align 16
+  %aa = alloca <2 x i64>, align 16, addrspace(5)
   %aptr = bitcast <2 x i64>* %a to i8*
-  %aaptr = bitcast <2 x i64>* %aa to i8*
-  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %aaptr, i8* %aptr, i32 16, i32 2, i1 false)
+  %aaptr = bitcast <2 x i64> addrspace(5)* %aa to i8 addrspace(5)*
+  call void @llvm.memcpy.p5i8.p0i8.i32(i8 addrspace(5)* %aaptr, i8* %aptr, i32 16, i32 2, i1 false)
   %bptr = bitcast i16 addrspace(1)* %b to i8 addrspace(1)*
-  call void @llvm.memcpy.p1i8.p0i8.i32(i8 addrspace(1)* %bptr, i8* %aaptr, i32 16, i32 2, i1 false)
+  call void @llvm.memcpy.p1i8.p5i8.i32(i8 addrspace(1)* %bptr, i8 addrspace(5)* %aaptr, i32 16, i32 2, i1 false)
   ret void
 }
 
 %struct.struct_test_27.0.13 = type { i32, float, i64, i8, [4 x i32] }
 
 ; Function Attrs: nounwind
 define void @copy_struct([5 x i64] %in.coerce) {
 ; CHECK-LABEL: @copy_struct(
 ; CHECK-NOT: memcpy
 for.end:
-  %in = alloca %struct.struct_test_27.0.13, align 8
-  %0 = bitcast %struct.struct_test_27.0.13* %in to [5 x i64]*
-  store [5 x i64] %in.coerce, [5 x i64]* %0, align 8
-  %scevgep9 = getelementptr %struct.struct_test_27.0.13, %struct.struct_test_27.0.13* %in, i32 0, i32 4, i32 0
-  %scevgep910 = bitcast i32* %scevgep9 to i8*
-  call void @llvm.memcpy.p1i8.p0i8.i32(i8 addrspace(1)* undef, i8* %scevgep910, i32 16, i32 4, i1 false)
+  %in = alloca %struct.struct_test_27.0.13, align 8, addrspace(5)
+  %0 = bitcast %struct.struct_test_27.0.13 addrspace(5)* %in to [5 x i64] addrspace(5)*
+  store [5 x i64] %in.coerce, [5 x i64] addrspace(5)* %0, align 8
+  %scevgep9 = getelementptr %struct.struct_test_27.0.13, %struct.struct_test_27.0.13 addrspace(5)* %in, i32 0, i32 4, i32 0
+  %scevgep910 = bitcast i32 addrspace(5)* %scevgep9 to i8 addrspace(5)*
+  call void @llvm.memcpy.p1i8.p5i8.i32(i8 addrspace(1)* undef, i8 addrspace(5)* %scevgep910, i32 16, i32 4, i1 false)
   ret void
 }
  
 %union.anon = type { i32* }
 
 @g = common global i32 0, align 4
 @l = common addrspace(3) global i32 0, align 4
 
 ; Make sure an illegal bitcast isn't introduced
 define void @pr27557() {
 ; CHECK-LABEL: @pr27557(
-; CHECK: %[[CAST:.*]] = bitcast i32** {{.*}} to i32 addrspace(3)**
-; CHECK: store i32 addrspace(3)* @l, i32 addrspace(3)** %[[CAST]]
-  %1 = alloca %union.anon, align 8
-  %2 = bitcast %union.anon* %1 to i32**
-  store i32* @g, i32** %2, align 8
-  %3 = bitcast %union.anon* %1 to i32 addrspace(3)**
-  store i32 addrspace(3)* @l, i32 addrspace(3)** %3, align 8
+; CHECK: %[[CAST:.*]] = bitcast i32* addrspace(5)* {{.*}} to i32 addrspace(3)* addrspace(5)*
+; CHECK: store i32 addrspace(3)* @l, i32 addrspace(3)* addrspace(5)* %[[CAST]]
+  %1 = alloca %union.anon, align 8, addrspace(5)
+  %2 = bitcast %union.anon addrspace(5)* %1 to i32* addrspace(5)*
+  store i32* @g, i32* addrspace(5)* %2, align 8
+  %3 = bitcast %union.anon addrspace(5)* %1 to i32 addrspace(3)* addrspace(5)*
+  store i32 addrspace(3)* @l, i32 addrspace(3)* addrspace(5)* %3, align 8
   ret void
 }
 
 ; Make sure pre-splitting doesn't try to introduce an illegal bitcast
 define float @presplit(i64 addrspace(1)* %p) {
 entry:
 ; CHECK-LABEL: @presplit(
 ; CHECK: %[[CAST:.*]] = bitcast i64 addrspace(1)* {{.*}} to i32 addrspace(1)*
 ; CHECK: load i32, i32 addrspace(1)* %[[CAST]]
-   %b = alloca i64
-   %b.cast = bitcast i64* %b to [2 x float]*
-   %b.gep1 = getelementptr [2 x float], [2 x float]* %b.cast, i32 0, i32 0
-   %b.gep2 = getelementptr [2 x float], [2 x float]* %b.cast, i32 0, i32 1
+   %b = alloca i64, align 8, addrspace(5)
+   %b.cast = bitcast i64 addrspace(5)* %b to [2 x float] addrspace(5)*
+   %b.gep1 = getelementptr [2 x float], [2 x float] addrspace(5)* %b.cast, i32 0, i32 0
+   %b.gep2 = getelementptr [2 x float], [2 x float] addrspace(5)* %b.cast, i32 0, i32 1
    %l = load i64, i64 addrspace(1)* %p
-   store i64 %l, i64* %b
-   %f1 = load float, float* %b.gep1
-   %f2 = load float, float* %b.gep2
+   store i64 %l, i64 addrspace(5)* %b
+   %f1 = load float, float addrspace(5)* %b.gep1
+   %f2 = load float, float addrspace(5)* %b.gep2
    %ret = fadd float %f1, %f2
    ret float %ret
 }
 
 ; Test load from and store to non-zero address space.
 define void @test_load_store_diff_addr_space([2 x float] addrspace(1)* %complex1, [2 x float] addrspace(1)* %complex2) {
 ; CHECK-LABEL: @test_load_store_diff_addr_space
 ; CHECK-NOT: alloca
 ; CHECK: load i32, i32 addrspace(1)*
 ; CHECK: load i32, i32 addrspace(1)*
 ; CHECK: store i32 %{{.*}}, i32 addrspace(1)*
 ; CHECK: store i32 %{{.*}}, i32 addrspace(1)*
-  %a = alloca i64
-  %a.cast = bitcast i64* %a to [2 x float]*
-  %a.gep1 = getelementptr [2 x float], [2 x float]* %a.cast, i32 0, i32 0
-  %a.gep2 = getelementptr [2 x float], [2 x float]* %a.cast, i32 0, i32 1
+  %a0 = alloca [2 x i64], align 8, addrspace(5)
+  %a = getelementptr [2 x i64], [2 x i64] addrspace(5)* %a0, i32 0, i32 0
+  %a.cast = bitcast i64 addrspace(5)* %a to [2 x float] addrspace(5)*
+  %a.gep1 = getelementptr [2 x float], [2 x float] addrspace(5)* %a.cast, i32 0, i32 0
+  %a.gep2 = getelementptr [2 x float], [2 x float] addrspace(5)* %a.cast, i32 0, i32 1
   %complex1.gep = getelementptr [2 x float], [2 x float] addrspace(1)* %complex1, i32 0, i32 0
   %p1 = bitcast float addrspace(1)* %complex1.gep to i64 addrspace(1)*
   %v1 = load i64, i64 addrspace(1)* %p1
-  store i64 %v1, i64* %a
-  %f1 = load float, float* %a.gep1
-  %f2 = load float, float* %a.gep2
+  store i64 %v1, i64 addrspace(5)* %a
+  %f1 = load float, float addrspace(5)* %a.gep1
+  %f2 = load float, float addrspace(5)* %a.gep2
   %sum = fadd float %f1, %f2
-  store float %sum, float* %a.gep1
-  store float %sum, float* %a.gep2
-  %v2 = load i64, i64* %a
+  store float %sum, float addrspace(5)* %a.gep1
+  store float %sum, float addrspace(5)* %a.gep2
+  %v2 = load i64, i64 addrspace(5)* %a
   %complex2.gep = getelementptr [2 x float], [2 x float] addrspace(1)* %complex2, i32 0, i32 0
   %p2 = bitcast float addrspace(1)* %complex2.gep to i64 addrspace(1)*
   store i64 %v2, i64 addrspace(1)* %p2
   ret void
 }