[DAG] Improve store merge candidate pruning.

[niravd]

Summary

Remove non-consecutive stores from candidate search as they can never
be merged and may prevent us from finding subsequent merges.

[spatel]

The limitations of store merging aren't clear to me. Can you add a code comment about why this is necessary? I would have assumed that we'd merge the later consecutive stores when visiting them because we visited them first.
Also, it would be better to add a more obvious test case. I think this should work as an x86 test:

; Don't let a non-consecutive store thwart merging of the last two.
define void @almost_consecutive_stores(i8* %p) {
store i8 0, i8* %p
%p1 = getelementptr i8, i8* %p, i64 42
store i8 1, i8* %p1
%p2 = getelementptr i8, i8* %p, i64 2
store i8 2, i8* %p2
%p3 = getelementptr i8, i8* %p, i64 3
store i8 3, i8* %p3
ret void
}

If that's correct, can you add that test as an NFC change ahead of this patch, so we just see the diff from your patch? Thanks.

[niravd]

Update with Sanjay's suggestions and rebased with to show additional testcase

[spatel]

LGTM, but see inline for nits in code comments.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12380 ↗

(On Diff #95564)

comma after "However"

12382 ↗

(On Diff #95564)

mismatched (}
comma before "we"

12383 ↗

(On Diff #95564)

typo: mergeable

12392 ↗

(On Diff #95564)

we've -> we
Candidates -> candidates
to do a -> to

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

bool DAGCombiner::MergeConsecutiveStores(StoreSDNode *St) {
 
   // Sort the memory operands according to their distance from the
   // base pointer.
   std::sort(StoreNodes.begin(), StoreNodes.end(),
             [](MemOpLink LHS, MemOpLink RHS) {
               return LHS.OffsetFromBase < RHS.OffsetFromBase;
             });
 
+  // Store Merge attempts to merge the lowest stores. This generally
+  // works out as if successful, as the remaining stores are checked
+  // after the first collection of stores is merged. However, in the
+  // case that a non-mergeable store is found first, e.g., {p[-2],
+  // p[0], p[1], p[2], p[3]}, we would fail and miss the subsequent
+  // mergeable cases. To prevent this, we prune such stores from the
+  // front of StoreNodes here.
+
+  unsigned StartIdx = 0;
+  while ((StartIdx + 1 < StoreNodes.size()) &&
+         StoreNodes[StartIdx].OffsetFromBase + ElementSizeBytes !=
+             StoreNodes[StartIdx + 1].OffsetFromBase)
+    ++StartIdx;
+
+  // Bail if we don't have enough candidates to merge.
+  if (StartIdx + 1 >= StoreNodes.size())
+    return false;
+
+  if (StartIdx)
+    StoreNodes.erase(StoreNodes.begin(), StoreNodes.begin() + StartIdx);
+
   // Scan the memory operations on the chain and find the first non-consecutive
   // store memory address.
   unsigned NumConsecutiveStores = 0;
   int64_t StartAddress = StoreNodes[0].OffsetFromBase;
 
   // Check that the addresses are consecutive starting from the second
   // element in the list of stores.
   for (unsigned i = 1, e = StoreNodes.size(); i < e; ++i) {

llvm/trunk/test/CodeGen/AArch64/arm64-abi.ll

; FAST-DAG: ldrsb  {{w[0-9]+}}, [sp]
   ret i64 %conv17
 }
 
 define i32 @i8i16caller() nounwind readnone {
 entry:
 ; CHECK-LABEL: i8i16caller
 ; The 8th, 9th, 10th and 11th arguments are passed at sp, sp+2, sp+4, sp+5.
 ; They are i8, i16, i8 and i8.
-; CHECK-DAG: strb {{w[0-9]+}}, [sp, #5]
-; CHECK-DAG: strb {{w[0-9]+}}, [sp, #4]
-; CHECK-DAG: strh {{w[0-9]+}}, [sp, #2]
+; CHECK-DAG: stur {{w[0-9]+}}, [sp, #2]
 ; CHECK-DAG: strb {{w[0-9]+}}, [sp]
 ; CHECK: bl
 ; FAST-LABEL: i8i16caller
 ; FAST: strb {{w[0-9]+}}, [sp]
 ; FAST: strh {{w[0-9]+}}, [sp, #2]
 ; FAST: strb {{w[0-9]+}}, [sp, #4]
 ; FAST: strb {{w[0-9]+}}, [sp, #5]
 ; FAST: bl

llvm/trunk/test/CodeGen/X86/MergeConsecutiveStores.ll

define void @almost_consecutive_stores(i8* %p) {
   %p1 = getelementptr i8, i8* %p, i64 42
   store i8 1, i8* %p1
   %p2 = getelementptr i8, i8* %p, i64 2
   store i8 2, i8* %p2
   %p3 = getelementptr i8, i8* %p, i64 3
   store i8 3, i8* %p3
   ret void
 ; CHECK-LABEL: almost_consecutive_stores
-; CHECK-DAG: movb $0, (%rdi)
-; CHECK-DAG: movb $1, 42(%rdi)
-; CHECK-DAG: movb $2, 2(%rdi)
-; CHECK-DAG: movb $3, 3(%rdi)
+; CHECK-DAG: movb $0, (%rdi)
+; CHECK-DAG: movb $1, 42(%rdi)
+; CHECK-DAG: movw $770, 2(%rdi)
 ; CHECK: retq
-
-; We should able to merge the final two stores into a 16-bit store
-; FIXMECHECK-DAG: movw $770, 2(%rdi)
-
-
 }