[SLPVectorizer] Fix for PR25748: reduction vectorization after loop unrolling.

[ABataev]

Summary

The next code is not vectorized by the SLPVectorizer:

int test(unsigned int *p) {
  int sum = 0;
  for (int i = 0; i < 8; i++)
    sum += p[i];
  return sum;
 }

During optimization this loop is fully unrolled and SLPVectorizer is unable to vectorize it. Patch tries to fix this problem.

[ABataev]

Reworked handling of not BinaryOps/SelectInst instructions

[mkuper]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

Why do you need this change?

4171 ↗

(On Diff #72055)

This comment looks wrong now. Could you replace it with a comment that explains the current situation?

4586 ↗

(On Diff #72055)

I actually started coding a similar patch at some point, but decided against it, because the "return" here seems really fairly accidental.
You could just as easily have something like:

int test(unsigned int *p) {
  int sum = 0;
  for (int i = 0; i < 8; i++)
    sum += p[i];
  return sum + 5;
}

It may be the case that having it feed a return is an important special case, but maybe we can work out a full solution for this.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

It is required to handle the next kind of hor reduction:

sum = a + b + c + ...

It comes from:

int test(unsigned int *p) {
  int sum = 0;
  for (int i = 0; i < 8; i++)
    sum += p[i];
  return sum;
}

4171 ↗

(On Diff #72055)

Ok, will do it

4586 ↗

(On Diff #72055)

Michael,
I'm working on a general patch, that is able to vectorize your example (I have it already but need some time to improve it). But this patch is still required as a first step for the second patch. Later we can rework this code for better compatibility

[ABataev]

Address Michael's comments

[mkuper]

Hi Alexey,
Sorry for the slow response time, I'm intermittently on vacation.

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

What I'm asking is, how do you even get trees of height 1?

4171 ↗

(On Diff #72055)

The code below is really not entirely obvious (at least to me - admittedly, I'm not an expert on the SLP vectorizer...).
Could you actually update the comment to explain what it does - that is, why you have that specific set of conditions - instead of just removing it?

4586 ↗

(On Diff #72055)

Ok, that sounds reasonable.

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

I could be wrong about this, but I think the horizontal reduction matching logic looks through the actual reduction step, since it will generate customized vector operations for the reduce. So it will build trees starting from the reduction operands, not the reduced value. Looking at the test case, this means the adds won't be in the tree. Only the loads will be, so the height will be one.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

Yes, that's right

4171 ↗

(On Diff #72055)

Ok

[ABataev]

Added a comment + updated

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

I think it would be a good idea to rework this function so that it checks the fully vectorizable property up to max(MinTreeSize, VectorizeableTree.size()). The current state of things is a bit confusing. MinTreeSize defines what is "tiny" yet here we assume tiny means a height of two. What do you think?

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

That should be min(MinTreeSize, VectorizeableTree.size()), not max.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

Yes, I agree. But I believe it must done in a different patch

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797 ↗

(On Diff #72055)

Sounds good to me.

[ABataev]

Other comments?

[mkuper]

LGTM

llvm/trunk/lib/Transforms/Vectorize/SLPVectorizer.cpp

default:
       llvm_unreachable("Unknown instruction");
   }
 }
 
 bool BoUpSLP::isFullyVectorizableTinyTree() {
   DEBUG(dbgs() << "SLP: Check whether the tree with height " <<
         VectorizableTree.size() << " is fully vectorizable .\n");
 
-  // We only handle trees of height 2.
+  // We only handle trees of heights 1 and 2.
+  if (VectorizableTree.size() == 1 && !VectorizableTree[0].NeedToGather)
+    return true;
+
   if (VectorizableTree.size() != 2)
     return false;
 
   // Handle splat and all-constants stores.
   if (!VectorizableTree[0].NeedToGather &&
       (allConstant(VectorizableTree[1].Scalars) ||
        isSplat(VectorizableTree[1].Scalars)))
     return true;
while (!Stack.empty()) {
         }
         // Retract.
         Stack.pop_back();
         continue;
       }
 
       // Visit left or right.
       Value *NextV = TreeN->getOperand(EdgeToVist);
-      // We currently only allow BinaryOperator's and SelectInst's as reduction
-      // values in our tree.
-      if (isa<BinaryOperator>(NextV) || isa<SelectInst>(NextV))
-        Stack.push_back(std::make_pair(cast<Instruction>(NextV), 0));
-      else if (NextV != Phi)
+      if (NextV != Phi) {
+        auto *I = dyn_cast<Instruction>(NextV);
+        // Continue analysis if the next operand is a reduction operation or
+        // (possibly) a reduced value. If the reduced value opcode is not set,
+        // the first met operation != reduction operation is considered as the
+        // reduced value class.
+        if (I && (!ReducedValueOpcode || I->getOpcode() == ReducedValueOpcode ||
+                  I->getOpcode() == ReductionOpcode)) {
+          if (!ReducedValueOpcode && I->getOpcode() != ReductionOpcode)
+            ReducedValueOpcode = I->getOpcode();
+          Stack.push_back(std::make_pair(I, 0));
+          continue;
+        }
         return false;
       }
+    }
     return true;
   }
 
   /// \brief Attempt to vectorize the tree found by
   /// matchAssociativeReduction.
   bool tryToReduce(BoUpSLP &V, TargetTransformInfo *TTI) {
     if (ReducedVals.empty())
       return false;
if (ShouldStartVectorizeHorAtStore)
         }
 
     // Try to vectorize horizontal reductions feeding into a return.
     if (ReturnInst *RI = dyn_cast<ReturnInst>(it))
       if (RI->getNumOperands() != 0)
         if (BinaryOperator *BinOp =
                 dyn_cast<BinaryOperator>(RI->getOperand(0))) {
           DEBUG(dbgs() << "SLP: Found a return to vectorize.\n");
-          if (tryToVectorizePair(BinOp->getOperand(0),
-                                 BinOp->getOperand(1), R)) {
+          if (canMatchHorizontalReduction(nullptr, BinOp, R, TTI,
+                                          R.getMinVecRegSize()) ||
+              tryToVectorizePair(BinOp->getOperand(0), BinOp->getOperand(1),
+                                 R)) {
             Changed = true;
             it = BB->begin();
             e = BB->end();
             continue;
           }
         }
 
     // Try to vectorize trees that start at compare instructions.

llvm/trunk/test/Transforms/SLPVectorizer/X86/reduction_unrolled.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -slp-vectorizer -slp-vectorize-hor -S -mtriple=x86_64-unknown-linux-gnu -mcpu=bdver2 < %s | FileCheck %s
+
+; int test(unsigned int *p) {
+;   int sum = 0;
+;   for (int i = 0; i < 8; i++)
+;     sum += p[i];
+;   return sum;
+; }
+
+define i32 @test(i32* nocapture readonly %p) {
+; CHECK-LABEL: @test(
+; CHECK:         [[BC:%.*]] = bitcast i32* %p to <8 x i32>*
+; CHECK-NEXT:    [[LD:%.*]] = load <8 x i32>, <8 x i32>* [[BC]], align 4
+; CHECK:         [[RDX_SHUF:%.*]] = shufflevector <8 x i32> [[LD]], <8 x i32> undef, <8 x i32> <i32 4, i32 5, i32 6, i32 7, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX:%.*]] = add <8 x i32> [[LD]], [[RDX_SHUF]]
+; CHECK-NEXT:    [[RDX_SHUF1:%.*]] = shufflevector <8 x i32> [[BIN_RDX]], <8 x i32> undef, <8 x i32> <i32 2, i32 3, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX2:%.*]] = add <8 x i32> [[BIN_RDX]], [[RDX_SHUF1]]
+; CHECK-NEXT:    [[RDX_SHUF3:%.*]] = shufflevector <8 x i32> [[BIN_RDX2]], <8 x i32> undef, <8 x i32> <i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX4:%.*]] = add <8 x i32> [[BIN_RDX2]], [[RDX_SHUF3]]
+; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <8 x i32> [[BIN_RDX4]], i32 0
+; CHECK:         ret i32 [[TMP2]]
+;
+entry:
+  %0 = load i32, i32* %p, align 4
+  %arrayidx.1 = getelementptr inbounds i32, i32* %p, i64 1
+  %1 = load i32, i32* %arrayidx.1, align 4
+  %mul.18 = add i32 %1, %0
+  %arrayidx.2 = getelementptr inbounds i32, i32* %p, i64 2
+  %2 = load i32, i32* %arrayidx.2, align 4
+  %mul.29 = add i32 %2, %mul.18
+  %arrayidx.3 = getelementptr inbounds i32, i32* %p, i64 3
+  %3 = load i32, i32* %arrayidx.3, align 4
+  %mul.310 = add i32 %3, %mul.29
+  %arrayidx.4 = getelementptr inbounds i32, i32* %p, i64 4
+  %4 = load i32, i32* %arrayidx.4, align 4
+  %mul.411 = add i32 %4, %mul.310
+  %arrayidx.5 = getelementptr inbounds i32, i32* %p, i64 5
+  %5 = load i32, i32* %arrayidx.5, align 4
+  %mul.512 = add i32 %5, %mul.411
+  %arrayidx.6 = getelementptr inbounds i32, i32* %p, i64 6
+  %6 = load i32, i32* %arrayidx.6, align 4
+  %mul.613 = add i32 %6, %mul.512
+  %arrayidx.7 = getelementptr inbounds i32, i32* %p, i64 7
+  %7 = load i32, i32* %arrayidx.7, align 4
+  %mul.714 = add i32 %7, %mul.613
+  ret i32 %mul.714
+}

llvm/trunk/test/Transforms/SLPVectorizer/X86/scheduling.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -basicaa -slp-vectorizer -S -mtriple=i386-apple-macosx10.8.0 -mcpu=corei7-avx | FileCheck %s
 
-target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
-target triple = "x86_64-apple-macosx10.9.0"
-
-;CHECK-LABEL: @foo
-;CHECK: load <4 x i32>
-;CHECK: load <4 x i32>
-;CHECK: %[[S1:.+]] = add nsw <4 x i32>
-;CHECK-DAG: store <4 x i32> %[[S1]]
-;CHECK-DAG: %[[A1:.+]] = add nsw i32
-;CHECK-DAG: %[[A2:.+]] = add nsw i32 %[[A1]]
-;CHECK-DAG: %[[A3:.+]] = add nsw i32 %[[A2]]
-;CHECK-DAG: %[[A4:.+]] = add nsw i32 %[[A3]]
-;CHECK: ret i32 %[[A4]] 
-
-define i32 @foo(i32* nocapture readonly %diff) #0 {
+define i32 @foo(i32* nocapture readonly %diff) #0 {
+; CHECK-LABEL: @foo(
+; CHECK: load <4 x i32>
+; CHECK: load <4 x i32>
+; CHECK: [[S1:%.+]] = add nsw <4 x i32>
+; CHECK: store <4 x i32> [[S1]],
+; CHECK:         [[RDX_SHUF:%.*]] = shufflevector <4 x i32> [[S1]], <4 x i32> undef, <4 x i32> <i32 2, i32 3, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX:%.*]] = add <4 x i32> [[S1]], [[RDX_SHUF]]
+; CHECK-NEXT:    [[RDX_SHUF1:%.*]] = shufflevector <4 x i32> [[BIN_RDX]], <4 x i32> undef, <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX2:%.*]] = add <4 x i32> [[BIN_RDX]], [[RDX_SHUF1]]
+; CHECK-NEXT:    [[TMP15:%.*]] = extractelement <4 x i32> [[BIN_RDX2]], i32 0
+; CHECK-NEXT:    [[ADD52:%.*]] = add nsw i32 [[TMP15]],
+; CHECK:          ret i32 [[ADD52]]
+;
 entry:
   %m2 = alloca [8 x [8 x i32]], align 16
   %0 = bitcast [8 x [8 x i32]]* %m2 to i8*
   br label %for.body
 
 for.body:                                         ; preds = %for.body, %entry
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %a.088 = phi i32 [ 0, %entry ], [ %add52, %for.body ]