[SROA] Fix crash due to bad bitcast

[loladiro]

Summary

As shown in the test case, SROA was crashing when trying to split
stores (to the alloca) of loads (from anywhere), because it assumed
the pointer operand to the loads and stores had to have the same
address space. This isn't the case. Make sure to use the correct
pointer type for both the load and the store.

[arsenm]

test/Transforms/SROA/address-spaces.ll

88 ↗

(On Diff #96916)

Needs check lines

[loladiro]

Add some CHECK lines to the test case

[loladiro]

Bump. Does this look good?

[loladiro]

Bump. Can somebody review this, so we can get this in. @arsenm maybe?

[yaxunl]

I would like to review this since I encounter the same issue.

[loladiro]

Sure, see if you think this is the right way to do it. If so, mark it as accepted, and I'll go ahead and commit it. I feel like we've given other people enough time to weigh in ;).

[yaxunl]

LGTM. Thanks!

llvm/trunk/lib/Transforms/Scalar/SROA.cpp

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 *LoadPartPtrTy = PartTy->getPointerTo(LI->getPointerAddressSpace());
+      auto *StorePartPtrTy = PartTy->getPointerTo(SI->getPointerAddressSpace());
 
       // 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() + "."),
+                           LoadPartPtrTy, 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,
                                 APInt(DL.getPointerSizeInBits(), PartOffset),
-                                PartPtrTy, StoreBasePtr->getName() + "."),
+                                StorePartPtrTy, StoreBasePtr->getName() + "."),
           getAdjustedAlignment(SI, PartOffset, DL), /*IsVolatile*/ false);
 
       // Now build a new slice for the alloca.
       NewSlices.push_back(
           Slice(BaseOffset + PartOffset, BaseOffset + PartOffset + PartSize,
                 &PStore->getOperandUse(PStore->getPointerOperandIndex()),
                 /*IsSplittable*/ false));
       DEBUG(dbgs() << "    new slice [" << NewSlices.back().beginOffset()

llvm/trunk/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
 }
+
+; 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
+   %l = load i64, i64 addrspace(1)* %p
+   store i64 %l, i64* %b
+   %f1 = load float, float* %b.gep1
+   %f2 = load float, float* %b.gep2
+   %ret = fadd float %f1, %f2
+   ret float %ret
+}