[SROA] Fix for PR25873: Pre-splitting should always preserve the partition structure.

[andreadb]

Summary

Hi,

This is an attempt to fix PR25873.
When splitting the store of a load, the pre splitting algorithm should check if the store uses a slice of the current alloca. If so, then the algorithm should update the slice information to preserve the partition structure. With this patch, the slice information related to the original store is used to create new slices for the split stores. The original store is eventually marked as dead and its associated slice is "killed".

If we don't update the alloca slices, then we risk to wrongly rewrite the partition of the alloca. I have posted a detailed description of the problem found in bug 25873. The reproducible from PR25873 has been added as a new test case for the existing test SROA/basictest.ll.

Example:

%tmpData = alloca %struct.STest;
...
[0,8):   %3 = load i64, i64* %1, align 8
[8,16):  store i64 %5, i64* %2, align 8

Both load and store use a slice of the %tmpData alloca (slice information is between brackets). In this example, the load is a candidate for pre splitting, and the store is the only user of the load.
Before this patch, the pre splitting algorithm would have rewritten the load/store pair as follows:

[0, 4):   %3 = load i32, i32* %.sroa_cast, align 8
[4, 8):   %4 = load i32, i32* %.sroa_cast1, align 4
[0, 8):   %3 = load i64, i64* %1, align 8   (dead)

      :  store i32 %3, i32* %.sroa_cast3, align 8
      :  store i32 %4, i32* %.sroa_cast5, align 4
[8,16):   store i64 %5, i64* %2, align 8    (dead)

The problem is that no slice information is built for the two new split stores. When the alloca is re-analyzed by SROA, two new problematic partitions are created: [8, 12) [12, 16). Those new partitions would overlap with the already consolidated partition [8, 16).

With this patch, we check if the store uses a slice of the alloca. If so, then we build new slices for the split stores and we mark the previous slice as dead. This is done to ensure that we choose the right partitioning of the alloca for the rewriting stage.
After this patch, the load and store are rewritten as:

[0, 4):   %3 = load i32, i32* %.sroa_cast, align 8
[4, 8):   %4 = load i32, i32* %.sroa_cast1, align 4
[0, 8):   %3 = load i64, i64* %1, align 8   (dead)

[8, 12):  store i32 %3, i32* %.sroa_cast3, align 8
[12,16):  store i32 %4, i32* %.sroa_cast5, align 4
       :   store i64 %5, i64* %2, align 8    (dead)

Please let me know what you think.

Thanks,
Andrea

[• rafael]

Wrong

[rob.lougher]

Looks to be a typo - it's PR25873.

[andreadb]

In D16057#323936, @rob.lougher wrote:

Looks to be a typo - it's PR25873.

Yes, it was the wrong PR. Sorry for the mistake..
I fixed the title and summary (s/PR25837/PR25873).

Again, sorry for the confusion caused.

-Andrea

[andreadb]

Ping.

[andreadb]

Ping^2.

Patch rebased.

[andreadb]

Ping.

[chandlerc]

Sorry its taken so long to cycle back here.

lib/Transforms/Scalar/SROA.cpp

3619–3626

This test is pretty expensive. I think you want to instead build a map from stores to slices above when we are scanning the alloca slices.

3681

You can re-order things to handle the smaller of these cases as an early 'continue' path I think.

[andreadb]

In D16057#341997, @chandlerc wrote:

Sorry its taken so long to cycle back here.

No problem at all. Thanks for the review!

I will address your comments and post an updated patch.

-Andrea

[andreadb]

Patch updated based on Chandler's review comments.

I have added a map from stores to slices. That map is populated when we initially scan through the slices of the alloca in 'presplitLoadsAndStores'.
I have also simplified the code adding an early continue as suggested.

Please let me know if this is ok to commit.

Thanks!
Andrea

[andreadb]

Ping.

[andreadb]

ping.

[andreadb]

Ping^3

[andreadb]

Ping! (patch rebased).

Can somebody please review this patch?

Thanks,
Andrea

[andreadb]

Ping.

[chandlerc]

I'm so, so sorry about this.

Not only for taking so long to get to it, but finally, I understand why I kept being confused when I looked at this patch.

I *COMPLETELY* messed up when I said that this was the correct fix. =[ The root cause of the bug was actually failing to bail out of pre-splitting earlier. There is a simple explanation, and it all stemmed from a previous patch of mine. I just missed a case.

[ =[ =[ =[ =[

I feel really terrible about how much work you've done because I misunderstood the root cause. So sorry about that.

Anyways, the root cause is also *super easy* to fix. I've gone ahead and landed the fix in r263121 using the excellent test case here.

Again, so sorry. =[ =[ =[

[andreadb]

In D16057#371909, @chandlerc wrote:

I'm so, so sorry about this.

Not only for taking so long to get to it, but finally, I understand why I kept being confused when I looked at this patch.

I *COMPLETELY* messed up when I said that this was the correct fix. =[ The root cause of the bug was actually failing to bail out of pre-splitting earlier. There is a simple explanation, and it all stemmed from a previous patch of mine. I just missed a case.

[ =[ =[ =[ =[

I feel really terrible about how much work you've done because I misunderstood the root cause. So sorry about that.

Anyways, the root cause is also *super easy* to fix. I've gone ahead and landed the fix in r263121 using the excellent test case here.

Again, so sorry. =[ =[ =[

Hey Chandler,

Don't worry! I knew that you have been very busy lately.
For what is worth, overall I don't feel like it was a waste time; at least now I know a bit more about SROA ;-).
Also, we weren't blocked by this issue as internally we had a conservative fix in place.

Anyway, I am very happy that the bug is fixed now.

Thanks,
Andrea

[andreadb]

Fixed at revision 263121.

lib/Transforms/Scalar/SROA.cpp

bool SROA::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
   // decreased by a frontend which just concatenated all of its local allocas
   // into one big flat alloca. But defeating such patterns is exactly the job
   // SROA is tasked with! Sadly, to not have this discrepancy we would have
   // change store pre-splitting to actually force pre-splitting of the load
   // that feeds it *and all stores*. That makes pre-splitting much harder, but
   // maybe it would make it more principled?
   SmallPtrSet<LoadInst *, 8> UnsplittableLoads;
 
+  // Track alloca slices used by store instructions.
+  SmallDenseMap<const StoreInst *, Slice *> Store2SliceMap;
+
   DEBUG(dbgs() << "  Searching for candidate loads and stores\n");
   for (auto &P : AS.partitions()) {
     for (Slice &S : P) {
       Instruction *I = cast<Instruction>(S.getUse()->getUser());
+
+      // Populate the Store2Slice map.
+      if (auto *SI = dyn_cast<StoreInst>(I))
+        Store2SliceMap.insert(std::make_pair(SI, &S));
+
       if (!S.isSplittable() ||S.endOffset() <= P.endOffset()) {
         // If this was a load we have to track that it can't participate in any
         // pre-splitting!
         if (auto *LI = dyn_cast<LoadInst>(I))
           UnsplittableLoads.insert(LI);
         continue;
       }
       assert(P.endOffset() > S.beginOffset() &&
for (User *LU : LI->users()) {
         DeferredStores = true;
         DEBUG(dbgs() << "    Deferred splitting of store: " << *SI << "\n");
         continue;
       }
 
       Value *StoreBasePtr = SI->getPointerOperand();
       IRB.SetInsertPoint(SI);
 
+      // We need to check if this store uses a slice of the alloca. If so, then
+      // we have to remember to build a new slice for each split store, and
+      // kill the original slice. Otherwise we risk to leave the partition
+      // structure in an inconsistent state.
+      auto Store2SliceMapI = Store2SliceMap.find(SI);
+      Slice *StoreSlice = nullptr;
+      if (Store2SliceMapI != Store2SliceMap.end())
+          StoreSlice = Store2SliceMapI->second;
+
       DEBUG(dbgs() << "    Splitting store of load: " << *SI << "\n");
 
       for (int Idx = 0, Size = SplitLoads.size(); Idx < Size; ++Idx) {
         LoadInst *PLoad = SplitLoads[Idx];
         uint64_t PartOffset = Idx == 0 ? 0 : Offsets.Splits[Idx - 1];
         auto *PartPtrTy =
             PLoad->getType()->getPointerTo(SI->getPointerAddressSpace());
 
         StoreInst *PStore = IRB.CreateAlignedStore(
             PLoad, getAdjustedPtr(IRB, DL, StoreBasePtr,
                                   APInt(DL.getPointerSizeInBits(), PartOffset),
                                   PartPtrTy, StoreBasePtr->getName() + "."),
             getAdjustedAlignment(SI, PartOffset, DL), /*IsVolatile*/ false);
         (void)PStore;
         DEBUG(dbgs() << "      +" << PartOffset << ":" << *PStore << "\n");
+
+        // Check if we have to build alloca slices for this split store.
+        if (StoreSlice) {
+          uint64_t BaseOffset = StoreSlice->beginOffset();
+          NewSlices.push_back(
+            Slice(BaseOffset + PartOffset, BaseOffset + PartOffset + PartSize,
+                  &PStore->getOperandUse(PStore->getPointerOperandIndex()),
+                  /*IsSplittable*/ false));
+          DEBUG(dbgs() << "    new slice [" << NewSlices.back().beginOffset()
+                       << ", " << NewSlices.back().endOffset()
+                       << "): " << *PStore << "\n");
+        }
+      }
+
+      // Mark the original store as dead.
+      DeadInsts.insert(SI);
+
+      // Check if we have to kill the slice associated to the original store.
+      // For consistency, we need to kill the original slice only if we have
+      // built new slices for the split stores.
+      if (StoreSlice) {
+        StoreSlice->kill();
+        continue;
       }
 
       // We want to immediately iterate on any allocas impacted by splitting
       // this store, and we have to track any promotable alloca (indicated by
       // a direct store) as needing to be resplit because it is no longer
       // promotable.
       if (AllocaInst *OtherAI = dyn_cast<AllocaInst>(StoreBasePtr)) {
         ResplitPromotableAllocas.insert(OtherAI);
         Worklist.insert(OtherAI);
       } else if (AllocaInst *OtherAI = dyn_cast<AllocaInst>(
                      StoreBasePtr->stripInBoundsOffsets())) {
         Worklist.insert(OtherAI);
       }
-
-      // Mark the original store as dead.
-      DeadInsts.insert(SI);
     }
 
     // Save the split loads if there are deferred stores among the users.
     if (DeferredStores)
       SplitLoadsMap.insert(std::make_pair(LI, std::move(SplitLoads)));
 
     // Mark the original load as dead and kill the original slice.
     DeadInsts.insert(LI);
     Offsets.S->kill();
   }
 

test/Transforms/SROA/basictest.ll

entry:
   %gep1 = getelementptr inbounds [3 x i8], [3 x i8]* %alloca, i64 0, i64 1
   %bc1 = bitcast i8* %gep1 to i16*
   store i16 0, i16* %bc1
   %gep2 = getelementptr inbounds [3 x i8], [3 x i8]* %alloca, i64 0, i64 2
   %bc2 = bitcast i8* %gep2 to i16*
   %load = load i16, i16* %bc2
   ret i16 %load
 }
+
+%struct.STest = type { %struct.SPos, %struct.SPos }
+%struct.SPos = type { float, float }
+
+define void @PR25873(%struct.STest* %outData) {
+; CHECK-LABEL: @PR25873(
+; CHECK: store i32 1123418112
+; CHECK: store i32 1139015680
+; CHECK: store i32 1123418112
+; CHECK: store i32 1139015680
+entry:
+  %tmpData = alloca %struct.STest, align 8
+  %0 = bitcast %struct.STest* %tmpData to i8*
+  call void @llvm.lifetime.start(i64 16, i8* %0)
+  %x = getelementptr inbounds %struct.STest, %struct.STest* %tmpData, i64 0, i32 0, i32 0
+  store float 1.230000e+02, float* %x, align 8
+  %y = getelementptr inbounds %struct.STest, %struct.STest* %tmpData, i64 0, i32 0, i32 1
+  store float 4.560000e+02, float* %y, align 4
+  %m_posB = getelementptr inbounds %struct.STest, %struct.STest* %tmpData, i64 0, i32 1
+  %1 = bitcast %struct.STest* %tmpData to i64*
+  %2 = bitcast %struct.SPos* %m_posB to i64*
+  %3 = load i64, i64* %1, align 8
+  store i64 %3, i64* %2, align 8
+  %4 = bitcast %struct.STest* %outData to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* %4, i8* %0, i64 16, i32 4, i1 false)
+  call void @llvm.lifetime.end(i64 16, i8* %0)
+  ret void
+}
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i32, i1) nounwind