[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

Why do you need this change?

4171

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

4578–4581

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

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

Ok, will do it

4578–4581

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

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

4171

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?

4578–4581

Ok, that sounds reasonable.

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797

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

Yes, that's right

4171

Ok

[ABataev]

Added a comment + updated

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797

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

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

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797

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

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1797

Sounds good to me.

[ABataev]

Other comments?

[mkuper]

LGTM

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))
+      if (isa<BinaryOperator>(NextV) || isa<SelectInst>(NextV) ||
+          isa<LoadInst>(NextV))
         Stack.push_back(std::make_pair(cast<Instruction>(NextV), 0));
       else if (NextV != Phi)
         return false;
     }
     return true;
   }
 
   /// \brief Attempt to vectorize the tree found by
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.

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
+}