SROA: Use correct address space when splitting loads (PR26154)

[hans]

Summary

escha reported this bug and proposed patch some time ago. Let's get it landed.

[chandlerc]

This test case should be cleaned up and merged into SROA/address-spaces.ll. Also, if I delete the 'getPointerAddressSpace()` call on line 3701, no test fails, so we clearly need better tests for the address space manipulation code in this file, or the call on line 3701 is actually dead and it should be *replaced* with the load-based address space in this patch.

[hans]

In D32608#741986, @chandlerc wrote:

This test case should be cleaned up and merged into SROA/address-spaces.ll.

Done.

While doing that I ran into problems because the data layout in that file specifies 16 bits for pointers in address space 1, which means the GEPs into that address space need to use i16 for the offset. I've tried to fix that for the GEPs involed here.

Also, if I delete the 'getPointerAddressSpace()` call on line 3701, no test fails, so we clearly need better tests for the address space manipulation code in this file,

:-/

or the call on line 3701 is actually dead and it should be *replaced* with the load-based address space in this patch.

No, using the address space from the load for the store will not pass the verifier in this test.

I'm not very familiar with SROA, but this patch makes sense to me: if we cannot find an already split load, we create a new one, and we do it in the same way as we would originally have done up around line 3575, i.e. using the load instruction's address space.

[hans]

Cleaning up and moving test into address-spaces.ll.

lib/Transforms/Scalar/SROA.cpp

for (LoadInst *LI : Loads) {
 
     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());
       LoadInst *PLoad = IRB.CreateAlignedLoad(
           getAdjustedPtr(IRB, DL, BasePtr,
-                         APInt(DL.getPointerSizeInBits(), PartOffset),
+                         APInt(DL.getPointerSizeInBits(
+                                   PartPtrTy->getPointerAddressSpace()),
+                               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);
for (StoreInst *SI : Stores) {
     } else {
       DEBUG(dbgs() << "          of 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(SI->getPointerAddressSpace());
+      auto *PartPtrTy = PartTy->getPointerTo();
 
       // Either lookup a split load or create one.
       LoadInst *PLoad;
       if (SplitLoads) {
         PLoad = (*SplitLoads)[Idx];
       } else {
         IRB.SetInsertPoint(LI);
         PLoad = IRB.CreateAlignedLoad(
-            getAdjustedPtr(IRB, DL, LoadBasePtr,
-                           APInt(DL.getPointerSizeInBits(), PartOffset),
-                           PartPtrTy, LoadBasePtr->getName() + "."),
+            getAdjustedPtr(
+                IRB, DL, LoadBasePtr,
+                APInt(DL.getPointerSizeInBits(LI->getPointerAddressSpace()),
+                      PartOffset),
+                PartTy->getPointerTo(LI->getPointerAddressSpace()),
+                LoadBasePtr->getName() + "."),
             getAdjustedAlignment(LI, PartOffset, DL), /*IsVolatile*/ false,
             LI->getName());
       }
 
       // And store this partition.
       IRB.SetInsertPoint(SI);
       StoreInst *PStore = IRB.CreateAlignedStore(
           PLoad, getAdjustedPtr(IRB, DL, StoreBasePtr,

test/Transforms/SROA/address-spaces.ll

 ; 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
   ret void
 }
+
+%pr26154.S = type { <4 x i32>, <3 x i8> }
+
+; Just check that we don't crash due to using the wrong address space when
+; splitting the load or store.
+; CHECK-LABEL: pr26154
+; CHECK: ret i8
+
+define i8 @pr26154(%pr26154.S addrspace(1)* %arg) {
+bb:
+  %a = alloca %pr26154.S, align 16
+  %b = alloca %pr26154.S, align 16
+
+  %arg_ptr = getelementptr inbounds %pr26154.S, %pr26154.S addrspace(1)* %arg, i16 0, i32 1
+  %arg_bc = bitcast <3 x i8> addrspace(1)* %arg_ptr to i32 addrspace(1)*
+  %arg_val = load i32, i32 addrspace(1)* %arg_bc, align 4
+
+  %a_ptr = getelementptr inbounds %pr26154.S, %pr26154.S* %a, i64 0, i32 1
+  %a_bc = bitcast <3 x i8>* %a_ptr to i32*
+  store i32 %arg_val, i32* %a_bc, align 16
+
+  %amem = bitcast %pr26154.S* %a to i8*
+  %bmem = bitcast %pr26154.S* %b to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %bmem, i8* %amem, i32 20, i32 16, i1 false)
+
+  %b_ptr = getelementptr %pr26154.S, %pr26154.S* %b, i64 0, i32 1
+  %b_val = load <3 x i8>, <3 x i8>* %b_ptr, align 16
+  %x = extractelement <3 x i8> %b_val, i32 0
+  ret i8 %x
+}