Teach BaseIndexOffset::match to identify base pointers in loops.

The small utility function that pattern matches Base + Index +
Offset patterns for loads and stores fails to recognize the base
pointer for loads/stores from/into an array at offset 0 inside a
loop. As a result DAGCombiner::MergeConsecutiveStores was not able
to merge all stores.

This commit fixes the issue by adding an additional pattern match
and also a test case.

Reviewer: Nadav

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@188936 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 44 insertions, 2 deletions

diff --git a/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index ba23aac..bbea113 100644
--- a/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -7859,17 +7859,29 @@ struct BaseIndexOffset {
   static BaseIndexOffset match(SDValue Ptr) {
     bool IsIndexSignExt = false;
 
-    // Just Base or possibly anything else.
+    // We only can pattern match BASE + INDEX + OFFSET. If Ptr is not an ADD
+    // instruction, then it could be just the BASE or everything else we don't
+    // know how to handle. Just use Ptr as BASE and give up.
     if (Ptr->getOpcode() != ISD::ADD)
       return BaseIndexOffset(Ptr, SDValue(), 0, IsIndexSignExt);
 
-    // Base + offset.
+    // We know that we have at least an ADD instruction. Try to pattern match
+    // the simple case of BASE + OFFSET.
     if (isa<ConstantSDNode>(Ptr->getOperand(1))) {
       int64_t Offset = cast<ConstantSDNode>(Ptr->getOperand(1))->getSExtValue();
       return  BaseIndexOffset(Ptr->getOperand(0), SDValue(), Offset,
                               IsIndexSignExt);
     }
 
+    // Inside a loop the current BASE pointer is calculated using an ADD and a
+    // MUL insruction. In this case Ptr is the actual BASE pointer.
+    // (i64 add (i64 %array_ptr)
+    //          (i64 mul (i64 %induction_var)
+    //                   (i64 %element_size)))
+    if (Ptr->getOperand(1)->getOpcode() == ISD::MUL) {
+      return BaseIndexOffset(Ptr, SDValue(), 0, IsIndexSignExt);
+    }
+
     // Look at Base + Index + Offset cases.
     SDValue Base = Ptr->getOperand(0);
     SDValue IndexOffset = Ptr->getOperand(1);
diff --git a/test/CodeGen/X86/merge_store.ll b/test/CodeGen/X86/merge_store.ll
new file mode 100644
index 0000000..940688c
--- /dev/null
+++ b/test/CodeGen/X86/merge_store.ll
@@ -0,0 +1,30 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mcpu=x86-64 | FileCheck %s
+
+define void @merge_store(i32* nocapture %a) {
+; CHECK-LABEL: merge_store:
+; CHECK: movq
+; CHECK: movq
+entry:
+  br label %for.body
+
+  for.body:
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32* %a, i64 %indvars.iv
+  store i32 1, i32* %arrayidx, align 4
+  %0 = or i64 %indvars.iv, 1
+  %arrayidx2 = getelementptr inbounds i32* %a, i64 %0
+  store i32 1, i32* %arrayidx2, align 4
+  %1 = or i64 %indvars.iv, 2
+  %arrayidx5 = getelementptr inbounds i32* %a, i64 %1
+  store i32 1, i32* %arrayidx5, align 4
+  %2 = or i64 %indvars.iv, 3
+  %arrayidx8 = getelementptr inbounds i32* %a, i64 %2
+  store i32 1, i32* %arrayidx8, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 4
+  %3 = trunc i64 %indvars.iv.next to i32
+  %cmp = icmp slt i32 %3, 1000
+  br i1 %cmp, label %for.body, label %for.end
+
+  for.end:
+  ret void
+}