[InstCombine] Insert a bitcast to enable merging similar store insts

[gandhi21299]

Summary

Given two Store instructions with equivalent pointer operands,
they could be merged into their common successor basic block if
the value operand of one is bitcasted to match the type of the
other. This patch only allows a bitcast between non-aggregate
primitive types with matching bitwidths.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

1633

Check for isBitOrNoopPointerCastable() instead, and please add a test that requires inserting ptrtoint or inttoptr.

Also this is missing a check for hasSameSpecialState(). Please add a test with for example different alignments.

1641

You should only insert the bitcast when actually doing the transform. We should not inserting if the transform is aborted later.

[gandhi21299]

• Replaced bitwidth equivalence check with isBitOrNoopPointerCastable()
• Inserted hasSameSpecialState() check
• Added a test with different alignments
• Added ptrtoint test
• Created an AMDGPU-specific test for inttoptr. This test appears to add an iteration to instruction combining until it reaches a fixed point. That is, instcombine will assume that it will run indefinitely when instcombine-infinite-loop-threshold is set to 3 whereas the test passes when I change the threshold value to 4.

[arsenm]

llvm/test/Transforms/InstCombine/merging-multiple-stores-into-successor.ll

221

Can't these merge using the minimum of the alignments?

[arsenm]

llvm/test/Transforms/InstCombine/merging-multiple-stores-into-successor.ll

243

This one is undefined because the alloca has default align 2

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

1625

These isSingleValueType / is VectorTy checks should not be necessary, isBitOrNoopPointerCastable will already do all the necessary type checks.

1682

There is no need to rewrite to a new temporary store instruction. You can insert the bitcast directly at the phi node operand.

You also need to use the IRBuilder to create the cast, otherwise it will not be queued for reprocessing. In that case you also don't need the InsertBitcast flag, because the IRBuilder will omit the bitcast if it is not needed.

If you do everything correctly, then your AMDGPU problem should go away as well -- the instcombine-infinite-loop-threshold flag exists specifically to detect these worklist management bugs.

llvm/test/Transforms/InstCombine/merging-multiple-stores-into-successor.ll

221

It's possible, but please not in this patch.

[gandhi21299]

• Use IRBuilder to create cast
• Eliminate InsertBitcast flag
• Fixed alignment test
• Removed redundant checks
• Eliminated temporary store

@nikic The issue with infinite-loop-threshold continues to persist.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

1614

Doesn't look like returning StoreInst * is needed/useful anymore -- return bool instead?

1648

This break got dropped. I think this is the reason for the test regressions and the crashes in pre-merge checks.

Took me a while to spot because I usually only look on the right side of the diff...

[gandhi21299]

• Brought back the break statement which was removed accidentally
• OtherStoreIsMergable returns a bool now, instead of a StoreInst *
• Corrected tests

[gandhi21299]

• rebase and re-run premerge checks

[nikic]

This basically looks good to me, but I'll have to check what is going on with instcombine-infinite-loop-threshold here.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

1618–1619

Can drop this check -- it's a subset of the one below.

1623–1626

[gandhi21299]

• Simplified OtherStoreIsMergeable as requested

The reason for another iteration of instcombine for @inttoptr_merge(..) is that mergeStoreIntoSuccessor() generates a case where a store/load pair is merged into an inttoptr which is later merged with the PHI inserted by mergeStoreIntoSuccessor() in the first iteration. I am not sure if implementing these cases in mergeStoreIntoSuccessor() is viable since it will make the code redundant.

[gandhi21299]

• clang-format and rebase

[foad]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

1653–1657

This can be simplified as suggested, since if we find a non-mergable store, the mayWriteToMemory test just below will do the return false.

Also I think moving the cast-to-StoreInst inside the helper function would make the patch simpler overall, i.e. something like:

auto IsMergeableStore = [&](Instruction *OtherStore) -> bool ...

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

1653–1657

We still need the StoreInst to assign to OtherStore, so I don't think moving the cast into the helper would work.

[nikic]

In D150900#4359334, @gandhi21299 wrote:

The reason for another iteration of instcombine for @inttoptr_merge(..) is that mergeStoreIntoSuccessor() generates a case where a store/load pair is merged into an inttoptr which is later merged with the PHI inserted by mergeStoreIntoSuccessor() in the first iteration. I am not sure if implementing these cases in mergeStoreIntoSuccessor() is viable since it will make the code redundant.

I've looked into this as well, and I think this would mostly be solved by D75362 (we visit instructions in the wrong order, in a way that is relevant here). It's okay to ignore this.

It took me a while to understand why AMDGPU is relevant here. The difference is that with default data layout i64 only has 4 byte alignment, so we first need to promote the alignment to 8 before the transform can be performed. For AMDGPU the alignment is 8 to start with. You can make this test AMDGPU independent by just explicitly specifying the alignment on the stores, instead of using natural alignment.

[gandhi21299]

• Simplified code
• Made inttoptr_merge() target independent by inserting align 8 to store instructions and support from D75362 as mentioned

[nikic]

LGTM

[gandhi21299]

Many thanks for the code review. I will push this patch to main.

[gandhi21299]

• Removed the inclusion of 'Instructions.h'

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

if (StoreBB == DestBB || OtherBB == DestBB)
     return false;
 
   // Verify that the other block ends in a branch and is not otherwise empty.
   BasicBlock::iterator BBI(OtherBB->getTerminator());
   BranchInst *OtherBr = dyn_cast<BranchInst>(BBI);
   if (!OtherBr || BBI == OtherBB->begin())
     return false;
 
+  auto OtherStoreIsMergeable = [&](StoreInst *OtherStore) -> bool {
+    if (!OtherStore ||
+        OtherStore->getPointerOperand() != SI.getPointerOperand())
+      return false;
+
+    auto *SIVTy = SI.getValueOperand()->getType();
+    auto *OSVTy = OtherStore->getValueOperand()->getType();
+    return CastInst::isBitOrNoopPointerCastable(OSVTy, SIVTy, DL) &&
+           SI.hasSameSpecialState(OtherStore);
+  };
+
   // If the other block ends in an unconditional branch, check for the 'if then
   // else' case. There is an instruction before the branch.
   StoreInst *OtherStore = nullptr;
   if (OtherBr->isUnconditional()) {
     --BBI;
     // Skip over debugging info and pseudo probes.
     while (BBI->isDebugOrPseudoInst() ||
            (isa<BitCastInst>(BBI) && BBI->getType()->isPointerTy())) {
       if (BBI==OtherBB->begin())
         return false;
       --BBI;
     }
     // If this isn't a store, isn't a store to the same location, or is not the
     // right kind of store, bail out.
     OtherStore = dyn_cast<StoreInst>(BBI);
-    if (!OtherStore || OtherStore->getOperand(1) != SI.getOperand(1) ||
-        !SI.isSameOperationAs(OtherStore))
+    if (!OtherStoreIsMergeable(OtherStore))
       return false;
   } else {
     // Otherwise, the other block ended with a conditional branch. If one of the
     // destinations is StoreBB, then we have the if/then case.
     if (OtherBr->getSuccessor(0) != StoreBB &&
         OtherBr->getSuccessor(1) != StoreBB)
       return false;
 
     // Okay, we know that OtherBr now goes to Dest and StoreBB, so this is an
     // if/then triangle. See if there is a store to the same ptr as SI that
     // lives in OtherBB.
     for (;; --BBI) {
       // Check to see if we find the matching store.
-      if ((OtherStore = dyn_cast<StoreInst>(BBI))) {
-        if (OtherStore->getOperand(1) != SI.getOperand(1) ||
-            !SI.isSameOperationAs(OtherStore))
-          return false;
+      OtherStore = dyn_cast<StoreInst>(BBI);
+      if (OtherStoreIsMergeable(OtherStore))
         break;
-      }
+
       // If we find something that may be using or overwriting the stored
       // value, or if we run out of instructions, we can't do the transform.
       if (BBI->mayReadFromMemory() || BBI->mayThrow() ||
           BBI->mayWriteToMemory() || BBI == OtherBB->begin())
         return false;
     }
 
     // In order to eliminate the store in OtherBr, we have to make sure nothing
     // reads or overwrites the stored value in StoreBB.
     for (BasicBlock::iterator I = StoreBB->begin(); &*I != &SI; ++I) {
       // FIXME: This should really be AA driven.
       if (I->mayReadFromMemory() || I->mayThrow() || I->mayWriteToMemory())
         return false;
     }
   }
 
   // Insert a PHI node now if we need it.
-  Value *MergedVal = OtherStore->getOperand(0);
+  Value *MergedVal = OtherStore->getValueOperand();
   // The debug locations of the original instructions might differ. Merge them.
   DebugLoc MergedLoc = DILocation::getMergedLocation(SI.getDebugLoc(),
                                                      OtherStore->getDebugLoc());
-  if (MergedVal != SI.getOperand(0)) {
-    PHINode *PN = PHINode::Create(MergedVal->getType(), 2, "storemerge");
-    PN->addIncoming(SI.getOperand(0), SI.getParent());
-    PN->addIncoming(OtherStore->getOperand(0), OtherBB);
+  if (MergedVal != SI.getValueOperand()) {
+    PHINode *PN =
+        PHINode::Create(SI.getValueOperand()->getType(), 2, "storemerge");
+    PN->addIncoming(SI.getValueOperand(), SI.getParent());
+    Builder.SetInsertPoint(OtherStore);
+    PN->addIncoming(Builder.CreateBitOrPointerCast(MergedVal, PN->getType()),
+                    OtherBB);
     MergedVal = InsertNewInstBefore(PN, DestBB->front());
     PN->setDebugLoc(MergedLoc);
   }
 
   // Advance to a place where it is safe to insert the new store and insert it.
   BBI = DestBB->getFirstInsertionPt();
   StoreInst *NewSI =
       new StoreInst(MergedVal, SI.getOperand(1), SI.isVolatile(), SI.getAlign(),

llvm/test/Transforms/InstCombine/merging-multiple-stores-into-successor.ll

bb10:                                             ; preds = %bb
   store i32 %i11, ptr getelementptr inbounds ([0 x i32], ptr @arr_2, i64 0, i64 1), align 4
   store i16 %i4, ptr getelementptr inbounds ([0 x i16], ptr @arr_4, i64 0, i64 1), align 2
   store i32 %i8, ptr getelementptr inbounds ([8 x i32], ptr @arr_3, i64 0, i64 1), align 4
   br label %bb12
 
 bb12:                                             ; preds = %bb10, %bb9
   ret void
 }
+
+define half @diff_types_same_width_merge(i1 %cond, half %a, i16 %b) {
+; CHECK-LABEL: @diff_types_same_width_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB0:%.*]], label [[BB1:%.*]]
+; CHECK:       BB0:
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       BB1:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i16 [[B:%.*]] to half
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[STOREMERGE:%.*]] = phi half [ [[TMP0]], [[BB1]] ], [ [[A:%.*]], [[BB0]] ]
+; CHECK-NEXT:    ret half [[STOREMERGE]]
+;
+entry:
+  %alloca = alloca half
+  br i1 %cond, label %BB0, label %BB1
+BB0:
+  store half %a, ptr %alloca
+  br label %sink
+BB1:
+  store i16 %b, ptr %alloca
+  br label %sink
+sink:
+  %val = load half, ptr %alloca
+  ret half %val
+}
+
+define i32 @diff_types_diff_width_no_merge(i1 %cond, i32 %a, i64 %b) {
+; CHECK-LABEL: @diff_types_diff_width_no_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       A:
+; CHECK-NEXT:    store i32 [[A:%.*]], ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       B:
+; CHECK-NEXT:    store i64 [[B:%.*]], ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    ret i32 [[VAL]]
+;
+entry:
+  %alloca = alloca i64
+  br i1 %cond, label %A, label %B
+A:
+  store i32 %a, ptr %alloca
+  br label %sink
+B:
+  store i64 %b, ptr %alloca
+  br label %sink
+sink:
+  %val = load i32, ptr %alloca
+  ret i32 %val
+}
+
+define <4 x i32> @vec_no_merge(i1 %cond, <2 x i32> %a, <4 x i32> %b) {
+; CHECK-LABEL: @vec_no_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i64, align 16
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       A:
+; CHECK-NEXT:    store <2 x i32> [[A:%.*]], ptr [[ALLOCA]], align 16
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       B:
+; CHECK-NEXT:    store <4 x i32> [[B:%.*]], ptr [[ALLOCA]], align 16
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[VAL:%.*]] = load <4 x i32>, ptr [[ALLOCA]], align 16
+; CHECK-NEXT:    ret <4 x i32> [[VAL]]
+;
+entry:
+  %alloca = alloca i64
+  br i1 %cond, label %A, label %B
+A:
+  store <2 x i32> %a, ptr %alloca
+  br label %sink
+B:
+  store <4 x i32> %b, ptr %alloca
+  br label %sink
+sink:
+  %val = load <4 x i32>, ptr %alloca
+  ret <4 x i32> %val
+}
+
+%struct.half = type { half };
+
+define %struct.half @one_elem_struct_merge(i1 %cond, %struct.half %a, half %b) {
+; CHECK-LABEL: @one_elem_struct_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB0:%.*]], label [[BB1:%.*]]
+; CHECK:       BB0:
+; CHECK-NEXT:    [[TMP0:%.*]] = extractvalue [[STRUCT_HALF:%.*]] [[A:%.*]], 0
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       BB1:
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[STOREMERGE:%.*]] = phi half [ [[TMP0]], [[BB0]] ], [ [[B:%.*]], [[BB1]] ]
+; CHECK-NEXT:    [[VAL1:%.*]] = insertvalue [[STRUCT_HALF]] poison, half [[STOREMERGE]], 0
+; CHECK-NEXT:    ret [[STRUCT_HALF]] [[VAL1]]
+;
+entry:
+  %alloca = alloca i64
+  br i1 %cond, label %BB0, label %BB1
+BB0:
+  store %struct.half %a, ptr %alloca
+  br label %sink
+BB1:
+  store half %b, ptr %alloca
+  br label %sink
+sink:
+  %val = load %struct.half, ptr %alloca
+  ret %struct.half %val
+}
+
+%struct.tup = type { half, i32 };
+
+define %struct.tup @multi_elem_struct_no_merge(i1 %cond, %struct.tup %a, half %b) {
+; CHECK-LABEL: @multi_elem_struct_no_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       A:
+; CHECK-NEXT:    store [[STRUCT_TUP:%.*]] [[A:%.*]], ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       B:
+; CHECK-NEXT:    store half [[B:%.*]], ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[VAL:%.*]] = load [[STRUCT_TUP]], ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    ret [[STRUCT_TUP]] [[VAL]]
+;
+entry:
+  %alloca = alloca i64
+  br i1 %cond, label %A, label %B
+A:
+  store %struct.tup %a, ptr %alloca
+  br label %sink
+B:
+  store half %b, ptr %alloca
+  br label %sink
+sink:
+  %val = load %struct.tup, ptr %alloca
+  ret %struct.tup %val
+}
+
+define i16 @same_types_diff_align_no_merge(i1 %cond, i16 %a, i16 %b) {
+; CHECK-LABEL: @same_types_diff_align_no_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i16, align 4
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB0:%.*]], label [[BB1:%.*]]
+; CHECK:       BB0:
+; CHECK-NEXT:    store i16 [[A:%.*]], ptr [[ALLOCA]], align 8
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       BB1:
+; CHECK-NEXT:    store i16 [[B:%.*]], ptr [[ALLOCA]], align 4
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[VAL:%.*]] = load i16, ptr [[ALLOCA]], align 4
+; CHECK-NEXT:    ret i16 [[VAL]]
+;
+entry:
+  %alloca = alloca i16, align 4
+  br i1 %cond, label %BB0, label %BB1
+BB0:
+  store i16 %a, ptr %alloca, align 8
+  br label %sink
+BB1:
+  store i16 %b, ptr %alloca, align 4
+  br label %sink
+sink:
+  %val = load i16, ptr %alloca
+  ret i16 %val
+}
+
+define i64 @ptrtoint_merge(i1 %cond, i64 %a, ptr %b) {
+; CHECK-LABEL: @ptrtoint_merge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB0:%.*]], label [[BB1:%.*]]
+; CHECK:       BB0:
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       BB1:
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint ptr [[B:%.*]] to i64
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[STOREMERGE:%.*]] = phi i64 [ [[A:%.*]], [[BB0]] ], [ [[TMP0]], [[BB1]] ]
+; CHECK-NEXT:    ret i64 [[STOREMERGE]]
+;
+entry:
+  %alloca = alloca ptr
+  br i1 %cond, label %BB0, label %BB1
+BB0:
+  store i64 %a, ptr %alloca
+  br label %sink
+BB1:
+  store ptr %b, ptr %alloca
+  br label %sink
+sink:
+  %val = load i64, ptr %alloca
+  ret i64 %val
+}
+
+define ptr @inttoptr_merge(i1 %cond, i64 %a, ptr %b) {
+; CHECK-LABEL: define ptr @inttoptr_merge
+; CHECK-SAME: (i1 [[COND:%.*]], i64 [[A:%.*]], ptr [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND]], label [[BB0:%.*]], label [[BB1:%.*]]
+; CHECK:       BB0:
+; CHECK-NEXT:    [[TMP0:%.*]] = inttoptr i64 [[A]] to ptr
+; CHECK-NEXT:    br label [[SINK:%.*]]
+; CHECK:       BB1:
+; CHECK-NEXT:    br label [[SINK]]
+; CHECK:       sink:
+; CHECK-NEXT:    [[STOREMERGE:%.*]] = phi ptr [ [[B]], [[BB1]] ], [ [[TMP0]], [[BB0]] ]
+; CHECK-NEXT:    ret ptr [[STOREMERGE]]
+;
+entry:
+  %alloca = alloca ptr
+  br i1 %cond, label %BB0, label %BB1
+BB0:
+  store i64 %a, ptr %alloca, align 8
+  br label %sink
+BB1:
+  store ptr %b, ptr %alloca, align 8
+  br label %sink
+sink:
+  %val = load ptr, ptr %alloca
+  ret ptr %val
+}