[SLP] Initial rework for min/max horizontal reduction vectorization, NFC.

[ABataev]

Summary

All getReductionCost() functions are renamed to getArithmeticReductionCost() + added basic infrastructure to handle non-binary reduction operations.

[mkuper]

lib/Analysis/CostModel.cpp

196 ↗

(On Diff #86662)

Can this ever be a Value that's not an Instruction?
(Same question applies to getReductionOpcode())

lib/Transforms/Vectorize/SLPVectorizer.cpp

4094 ↗

(On Diff #86662)

This basically complements the change in line 4803 - 4823, right, but taken together, those two changes are NFC, regardless of everything else here, right? Or am I confused and this depends on something else?

4197 ↗

(On Diff #86662)

This part of the change is completely orthogonal to the name change and using matchers, right?

Anyway, I think I see where this is going. This is possibly similar to what Ayal, Gil & company are doing for LV, in terms of recipes?
Adding them to take a look.

4198 ↗

(On Diff #86662)

What does being valid mean in this context? Just that this is a "non-null handle"

4245 ↗

(On Diff #86662)

only adds at this point. :-)
Also, if you're checking the op here, I'd expect a check of isAssociative() as well.

4249 ↗

(On Diff #86662)

So, by "of the same kind" you mean "both a reduction op or both a reduced value", which is something you determine by either both LHS/RHS being null, or both being non-null?
Can you please be more explicit about this? Also, maybe it's worth having an explicit flag for this, instead of relying on LHS and RHS being null?

4271 ↗

(On Diff #86662)

I've cleaned this up in the existing code, there's no need to special-case FAdd, you can just CreateBinOp for it.

[ABataev]

Update after review

[Ayal]

lib/Analysis/CostModel.cpp

211 ↗

(On Diff #86833)

Can PatternMatch be put to more use, by asking it to match a BinOp whose L and R sides are both shuffles, when Level > 0?

Passing L and R by reference here saves little compared to the original direct calls to getOperand(0) and getOperand(1), and seems mostly redundant elsewhere.

405 ↗

(On Diff #86833)

Check if RdxOp, or use cast instead of dyn_cast?

lib/Transforms/Vectorize/SLPVectorizer.cpp

4094 ↗

(On Diff #86833)

V is-surely-a<BinaryOperator>, being a non-nullptr to BinaryOperator, right?

[ABataev]

lib/Analysis/CostModel.cpp

211 ↗

(On Diff #86833)

This is just an initial patch. In future it will be extended with support of min/max patterns, that's why this kind of recognition of LHS/RHS operands is placed in function. Another thing, that it will be better just to return the record of opcode + LHS/RHS parts, rather than pass L/R by reference.

405 ↗

(On Diff #86833)

Yes, missed this, thanks. Will fix.

196 ↗

(On Diff #86662)

Of course no, only instructions can be used here. I'll convert these params to Instruction *.

lib/Transforms/Vectorize/SLPVectorizer.cpp

4094 ↗

(On Diff #86833)

Yes, forgot to change type of parameter, will do it

4094 ↗

(On Diff #86662)

Yes, you're right, I'll remove these changes from this patch and prepare another one for CmpInsts.

4197 ↗

(On Diff #86662)

No, actually it is basic infrastructure for supporting horizontal reduction vectorization of not-BinOp instructions. It does not change the functionality, but allows later easily to support vectorization of other complex instructions, like min/max.

4198 ↗

(On Diff #86662)

Will try to remove it

4245 ↗

(On Diff #86662)

Yes, will fix this.
isAssociative() requires actual instruction instead of Opcode only, that's why I did not check it here.

4249 ↗

(On Diff #86662)

Yes, it is exactly what I meant.
Ok, will think about it.

4271 ↗

(On Diff #86662)

Ok, will fix it.

[Ayal]

lib/Transforms/Vectorize/SLPVectorizer.cpp

4213 ↗

(On Diff #86833)

Should IsReducedValue be set to false here?

4267 ↗

(On Diff #86833)

re[d]uction

4249 ↗

(On Diff #86662)

The IsReducedValue explicit flag indicates if its a reduction op or reduced value, right?

To implement the behavior documented above: "Checks if two operation data are both a reduction op or both a reduced value", it's enough to have the "IsReducedValue == OD.IsReducedValue" part.

[ABataev]

Update after review

[ABataev]

Ping

[Ayal]

lib/Transforms/Vectorize/SLPVectorizer.cpp

4262 ↗

(On Diff #88877)

"+ validity" should be removed?

4285 ↗

(On Diff #88877)

Still a couple of "if (x.Opcode == 0)" cases, may as well make use of this bool().

4865 ↗

(On Diff #88877)

Not sure why you wanted this change; but how about simply "if (isa<BinaryOperator>(Inst))"?

4094 ↗

(On Diff #86833)

I must be missing something here; tryToVectorize() is still being fed by BinaryOperator*, right? Wondering why the parameter type change. Better include this change in follow-up patch which tries to vectorize non BinaryOperators, if that's the intention.

4267 ↗

(On Diff #86833)

typo

4269 ↗

(On Diff #86833)

Every other method fits with the passive, storage-type "OperationData" name of this struct, but this last method may look a bit odd doing "OperationData.createOp()". All that it's contributing is its Opcode which it can provide, and some asserts. Is there a better way of keeping OperationData's nature, or a better name for it?

4245 ↗

(On Diff #86662)

isVectorizable() is used only in conjuction with isAssociative(I), and as an assert when getOpcode() is called by getReductionCost(). Isn't it better to have isVectorizable(I) also call isAssociative?

4249 ↗

(On Diff #86662)

It sounds like you can

assert ((IsReducedValue != OD.IsReducedValue) || ((!LHS == !OD.LHS) && (!RHS == !OD.RHS) && Opcode == OD.Opcode) && "message");

no need to actually check them all to return true, right?

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

4262 ↗

(On Diff #88877)

Removed

4285 ↗

(On Diff #88877)

Tried to replace all

4865 ↗

(On Diff #88877)

Restored it

4094 ↗

(On Diff #86833)

Restored back.

4267 ↗

(On Diff #86833)

Yes, thanks, fixed

4269 ↗

(On Diff #86833)

Modified this method, now it accepts only Builder and optional name

4245 ↗

(On Diff #86662)

Added isVectorized(Instruction *I)

4249 ↗

(On Diff #86662)

Modified it a bit, still need to compare the opcode

[ABataev]

Update after review

[anemet]

lib/Analysis/CostModel.cpp

189 ↗

(On Diff #101209)

Contains

[RKSimon]

Some minor style comments

lib/Analysis/CostModel.cpp

202 ↗

(On Diff #101209)

No need to initialize these.

Value *L, *R;

204 ↗

(On Diff #101209)

Use the ReductionData constructor for clarity?

lib/Transforms/Vectorize/SLPVectorizer.cpp

4353 ↗

(On Diff #101209)

?

4358 ↗

(On Diff #101209)

Add assert message

4363 ↗

(On Diff #101209)

Add assert message

4368 ↗

(On Diff #101209)

Add assert message

4394 ↗

(On Diff #101209)

Add assert message

4445 ↗

(On Diff #101209)

Use OperationData constructor for clarity

[RKSimon]

@ABataev You marked a lot of the comments done but haven't updated the diff.

[ABataev]

Update after review.

[RKSimon]

A couple of NFC changes that will reduce this patch.

include/llvm/CodeGen/BasicTTIImpl.h

1121 ↗

(On Diff #101750)

You should probably move this comment (unchanged) as a NFC commit right away to reduce the patch. I can't tell if you've changed anything.

include/llvm/Transforms/Vectorize/SLPVectorizer.h

87 ↗

(On Diff #101750)

Do this as an NFC commit now - its a tidyup that isn't relevant to the patch.

[ABataev]

Update to trunk

[RKSimon]

What's happening with this? It seems almost ready to go.

lib/Analysis/CostModel.cpp

200 ↗

(On Diff #101934)

Worth making this an Optional<> instead of relying on Opcode == 0?

196 ↗

(On Diff #86662)

Still not Instruction*

[ABataev]

lib/Analysis/CostModel.cpp

200 ↗

(On Diff #101934)

Ok, will do

196 ↗

(On Diff #86662)

Investigated this problem once again, it maybe not of Instruction type, for example, it may be an Argument. So, we have to use Value * here.

[ABataev]

Update after review

[RKSimon]

There's still a lot going on in this patch but it looks alright to me (I noticed one minor).

Does anyone have any further comments?

lib/Transforms/Vectorize/SLPVectorizer.cpp

4337 ↗

(On Diff #108463)

If OperationData is a struct then this public can go?

[avt77]

Could we split it on 2 sub-parts:

1 - rename (maybe NFC?)
2 - new infrastructure

In this case it will be easier and smaller.

[RKSimon]

LGTM - if possible please commit the getArithmeticReductionCost rename as a separate pre-commit.

llvm/trunk/lib/Analysis/CostModel.cpp

 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Analysis/Passes.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace llvm;
+using namespace PatternMatch;
 
 #define CM_NAME "cost-model"
 #define DEBUG_TYPE CM_NAME
 
 static cl::opt<bool> EnableReduxCost("costmodel-reduxcost", cl::init(false),
                                      cl::Hidden,
                                      cl::desc("Recognize reduction patterns."));
 
static bool matchPairwiseShuffleMask(ShuffleVectorInst *SI, bool IsLeft,
   // we look at the left or right side.
   for (unsigned i = 0, e = (1 << Level), val = !IsLeft; i != e; ++i, val += 2)
     Mask[i] = val;
 
   SmallVector<int, 16> ActualMask = SI->getShuffleMask();
   return Mask == ActualMask;
 }
 
-static bool matchPairwiseReductionAtLevel(const BinaryOperator *BinOp,
-                                          unsigned Level, unsigned NumLevels) {
+namespace {
+/// Contains opcode + LHS/RHS parts of the reduction operations.
+struct ReductionData {
+  explicit ReductionData() = default;
+  ReductionData(unsigned Opcode, Value *LHS, Value *RHS)
+      : Opcode(Opcode), LHS(LHS), RHS(RHS) {}
+  unsigned Opcode = 0;
+  Value *LHS = nullptr;
+  Value *RHS = nullptr;
+};
+} // namespace
+
+static Optional<ReductionData> getReductionData(Instruction *I) {
+  Value *L, *R;
+  if (m_BinOp(m_Value(L), m_Value(R)).match(I))
+    return ReductionData(I->getOpcode(), L, R);
+  return llvm::None;
+}
+
+static bool matchPairwiseReductionAtLevel(Instruction *I, unsigned Level,
+                                          unsigned NumLevels) {
   // Match one level of pairwise operations.
   // %rdx.shuf.0.0 = shufflevector <4 x float> %rdx, <4 x float> undef,
   //       <4 x i32> <i32 0, i32 2 , i32 undef, i32 undef>
   // %rdx.shuf.0.1 = shufflevector <4 x float> %rdx, <4 x float> undef,
   //       <4 x i32> <i32 1, i32 3, i32 undef, i32 undef>
   // %bin.rdx.0 = fadd <4 x float> %rdx.shuf.0.0, %rdx.shuf.0.1
-  if (BinOp == nullptr)
+  if (!I)
     return false;
 
-  assert(BinOp->getType()->isVectorTy() && "Expecting a vector type");
+  assert(I->getType()->isVectorTy() && "Expecting a vector type");
 
-  unsigned Opcode = BinOp->getOpcode();
-  Value *L = BinOp->getOperand(0);
-  Value *R = BinOp->getOperand(1);
+  Optional<ReductionData> RD = getReductionData(I);
+  if (!RD)
+    return false;
 
-  ShuffleVectorInst *LS = dyn_cast<ShuffleVectorInst>(L);
+  ShuffleVectorInst *LS = dyn_cast<ShuffleVectorInst>(RD->LHS);
   if (!LS && Level)
     return false;
-  ShuffleVectorInst *RS = dyn_cast<ShuffleVectorInst>(R);
+  ShuffleVectorInst *RS = dyn_cast<ShuffleVectorInst>(RD->RHS);
   if (!RS && Level)
     return false;
 
   // On level 0 we can omit one shufflevector instruction.
   if (!Level && !RS && !LS)
     return false;
 
   // Shuffle inputs must match.
   Value *NextLevelOpL = LS ? LS->getOperand(0) : nullptr;
   Value *NextLevelOpR = RS ? RS->getOperand(0) : nullptr;
   Value *NextLevelOp = nullptr;
   if (NextLevelOpR && NextLevelOpL) {
     // If we have two shuffles their operands must match.
     if (NextLevelOpL != NextLevelOpR)
       return false;
 
     NextLevelOp = NextLevelOpL;
   } else if (Level == 0 && (NextLevelOpR || NextLevelOpL)) {
     // On the first level we can omit the shufflevector <0, undef,...>. So the
     // input to the other shufflevector <1, undef> must match with one of the
     // inputs to the current binary operation.
     // Example:
     //  %NextLevelOpL = shufflevector %R, <1, undef ...>
     //  %BinOp        = fadd          %NextLevelOpL, %R
-    if (NextLevelOpL && NextLevelOpL != R)
+    if (NextLevelOpL && NextLevelOpL != RD->RHS)
       return false;
-    else if (NextLevelOpR && NextLevelOpR != L)
+    else if (NextLevelOpR && NextLevelOpR != RD->LHS)
       return false;
 
-    NextLevelOp = NextLevelOpL ? R : L;
+    NextLevelOp = NextLevelOpL ? RD->RHS : RD->LHS;
   } else
     return false;
 
   // Check that the next levels binary operation exists and matches with the
   // current one.
-  BinaryOperator *NextLevelBinOp = nullptr;
   if (Level + 1 != NumLevels) {
-    if (!(NextLevelBinOp = dyn_cast<BinaryOperator>(NextLevelOp)))
-      return false;
-    else if (NextLevelBinOp->getOpcode() != Opcode)
+    Optional<ReductionData> NextLevelRD =
+        getReductionData(cast<Instruction>(NextLevelOp));
+    if (!NextLevelRD || RD->Opcode != NextLevelRD->Opcode)
       return false;
   }
 
   // Shuffle mask for pairwise operation must match.
-  if (matchPairwiseShuffleMask(LS, true, Level)) {
-    if (!matchPairwiseShuffleMask(RS, false, Level))
+  if (matchPairwiseShuffleMask(LS, /*IsLeft=*/true, Level)) {
+    if (!matchPairwiseShuffleMask(RS, /*IsLeft=*/false, Level))
       return false;
-  } else if (matchPairwiseShuffleMask(RS, true, Level)) {
-    if (!matchPairwiseShuffleMask(LS, false, Level))
+  } else if (matchPairwiseShuffleMask(RS, /*IsLeft=*/true, Level)) {
+    if (!matchPairwiseShuffleMask(LS, /*IsLeft=*/false, Level))
       return false;
   } else
     return false;
 
   if (++Level == NumLevels)
     return true;
 
   // Match next level.
-  return matchPairwiseReductionAtLevel(NextLevelBinOp, Level, NumLevels);
+  return matchPairwiseReductionAtLevel(cast<Instruction>(NextLevelOp), Level,
+                                       NumLevels);
 }
 
 static bool matchPairwiseReduction(const ExtractElementInst *ReduxRoot,
                                    unsigned &Opcode, Type *&Ty) {
   if (!EnableReduxCost)
     return false;
 
   // Need to extract the first element.
   ConstantInt *CI = dyn_cast<ConstantInt>(ReduxRoot->getOperand(1));
   unsigned Idx = ~0u;
   if (CI)
     Idx = CI->getZExtValue();
   if (Idx != 0)
     return false;
 
-  BinaryOperator *RdxStart = dyn_cast<BinaryOperator>(ReduxRoot->getOperand(0));
+  auto *RdxStart = dyn_cast<Instruction>(ReduxRoot->getOperand(0));
   if (!RdxStart)
     return false;
+  Optional<ReductionData> RD = getReductionData(RdxStart);
+  if (!RD)
+    return false;
 
-  Type *VecTy = ReduxRoot->getOperand(0)->getType();
+  Type *VecTy = RdxStart->getType();
   unsigned NumVecElems = VecTy->getVectorNumElements();
   if (!isPowerOf2_32(NumVecElems))
     return false;
 
   // We look for a sequence of shuffle,shuffle,add triples like the following
   // that builds a pairwise reduction tree.
   //
   //  (X0, X1, X2, X3)
static bool matchPairwiseReduction(const ExtractElementInst *ReduxRoot,
   //       <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
   // %rdx.shuf.1.1 = shufflevector <4 x float> %bin.rdx.0, <4 x float> undef,
   //       <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
   // %bin.rdx8 = fadd <4 x float> %rdx.shuf.1.0, %rdx.shuf.1.1
   // %r = extractelement <4 x float> %bin.rdx8, i32 0
   if (!matchPairwiseReductionAtLevel(RdxStart, 0,  Log2_32(NumVecElems)))
     return false;
 
-  Opcode = RdxStart->getOpcode();
+  Opcode = RD->Opcode;
   Ty = VecTy;
 
   return true;
 }
 
 static std::pair<Value *, ShuffleVectorInst *>
-getShuffleAndOtherOprd(BinaryOperator *B) {
-
-  Value *L = B->getOperand(0);
-  Value *R = B->getOperand(1);
+getShuffleAndOtherOprd(Value *L, Value *R) {
   ShuffleVectorInst *S = nullptr;
 
   if ((S = dyn_cast<ShuffleVectorInst>(L)))
     return std::make_pair(R, S);
 
   S = dyn_cast<ShuffleVectorInst>(R);
   return std::make_pair(L, S);
 }
 
 static bool matchVectorSplittingReduction(const ExtractElementInst *ReduxRoot,
                                           unsigned &Opcode, Type *&Ty) {
   if (!EnableReduxCost)
     return false;
 
   // Need to extract the first element.
   ConstantInt *CI = dyn_cast<ConstantInt>(ReduxRoot->getOperand(1));
   unsigned Idx = ~0u;
   if (CI)
     Idx = CI->getZExtValue();
   if (Idx != 0)
     return false;
 
-  BinaryOperator *RdxStart = dyn_cast<BinaryOperator>(ReduxRoot->getOperand(0));
+  auto *RdxStart = dyn_cast<Instruction>(ReduxRoot->getOperand(0));
   if (!RdxStart)
     return false;
-  unsigned RdxOpcode = RdxStart->getOpcode();
+  Optional<ReductionData> RD = getReductionData(RdxStart);
+  if (!RD)
+    return false;
 
   Type *VecTy = ReduxRoot->getOperand(0)->getType();
   unsigned NumVecElems = VecTy->getVectorNumElements();
   if (!isPowerOf2_32(NumVecElems))
     return false;
 
   // We look for a sequence of shuffles and adds like the following matching one
   // fadd, shuffle vector pair at a time.
   //
   // %rdx.shuf = shufflevector <4 x float> %rdx, <4 x float> undef,
   //                           <4 x i32> <i32 2, i32 3, i32 undef, i32 undef>
   // %bin.rdx = fadd <4 x float> %rdx, %rdx.shuf
   // %rdx.shuf7 = shufflevector <4 x float> %bin.rdx, <4 x float> undef,
   //                          <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
   // %bin.rdx8 = fadd <4 x float> %bin.rdx, %rdx.shuf7
   // %r = extractelement <4 x float> %bin.rdx8, i32 0
 
   unsigned MaskStart = 1;
-  Value *RdxOp = RdxStart;
+  Instruction *RdxOp = RdxStart;
   SmallVector<int, 32> ShuffleMask(NumVecElems, 0);
   unsigned NumVecElemsRemain = NumVecElems;
   while (NumVecElemsRemain - 1) {
     // Check for the right reduction operation.
-    BinaryOperator *BinOp;
-    if (!(BinOp = dyn_cast<BinaryOperator>(RdxOp)))
+    if (!RdxOp)
       return false;
-    if (BinOp->getOpcode() != RdxOpcode)
+    Optional<ReductionData> RDLevel = getReductionData(RdxOp);
+    if (!RDLevel || RDLevel->Opcode != RD->Opcode)
       return false;
 
     Value *NextRdxOp;
     ShuffleVectorInst *Shuffle;
-    std::tie(NextRdxOp, Shuffle) = getShuffleAndOtherOprd(BinOp);
+    std::tie(NextRdxOp, Shuffle) =
+        getShuffleAndOtherOprd(RDLevel->LHS, RDLevel->RHS);
 
     // Check the current reduction operation and the shuffle use the same value.
     if (Shuffle == nullptr)
       return false;
     if (Shuffle->getOperand(0) != NextRdxOp)
       return false;
 
     // Check that shuffle masks matches.
     for (unsigned j = 0; j != MaskStart; ++j)
       ShuffleMask[j] = MaskStart + j;
     // Fill the rest of the mask with -1 for undef.
     std::fill(&ShuffleMask[MaskStart], ShuffleMask.end(), -1);
 
     SmallVector<int, 16> Mask = Shuffle->getShuffleMask();
     if (ShuffleMask != Mask)
       return false;
 
-    RdxOp = NextRdxOp;
+    RdxOp = dyn_cast<Instruction>(NextRdxOp);
     NumVecElemsRemain /= 2;
     MaskStart *= 2;
   }
 
-  Opcode = RdxOpcode;
+  Opcode = RD->Opcode;
   Ty = VecTy;
   return true;
 }
 
 unsigned CostModelAnalysis::getInstructionCost(const Instruction *I) const {
   if (!TTI)
     return -1;
 
case Instruction::ExtractElement: {
     if (CI)
       Idx = CI->getZExtValue();
 
     // Try to match a reduction sequence (series of shufflevector and vector
     // adds followed by a extractelement).
     unsigned ReduxOpCode;
     Type *ReduxType;
 
-    if (matchVectorSplittingReduction(EEI, ReduxOpCode, ReduxType))
-      return TTI->getArithmeticReductionCost(ReduxOpCode, ReduxType, false);
-    else if (matchPairwiseReduction(EEI, ReduxOpCode, ReduxType))
-      return TTI->getArithmeticReductionCost(ReduxOpCode, ReduxType, true);
+    if (matchVectorSplittingReduction(EEI, ReduxOpCode, ReduxType)) {
+      return TTI->getArithmeticReductionCost(ReduxOpCode, ReduxType,
+                                             /*IsPairwiseForm=*/false);
+    }
+    if (matchPairwiseReduction(EEI, ReduxOpCode, ReduxType)) {
+      return TTI->getArithmeticReductionCost(ReduxOpCode, ReduxType,
+                                             /*IsPairwiseForm=*/true);
+    }
 
     return TTI->getVectorInstrCost(I->getOpcode(),
                                    EEI->getOperand(0)->getType(), Idx);
   }
   case Instruction::InsertElement: {
     const InsertElementInst * IE = cast<InsertElementInst>(I);
     ConstantInt *CI = dyn_cast<ConstantInt>(IE->getOperand(2));
     unsigned Idx = -1;

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

 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/NoFolder.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/GraphWriter.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Vectorize.h"
 #include <algorithm>
 #include <memory>
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 using namespace slpvectorizer;
 
 #define SV_NAME "slp-vectorizer"
 #define DEBUG_TYPE "SLP"
 
 STATISTIC(NumVectorInstructions, "Number of vector instructions generated");
 
 static cl::opt<int>

 ///   *p =
 ///
 class HorizontalReduction {
   SmallVector<Value *, 16> ReductionOps;
   SmallVector<Value *, 32> ReducedVals;
   // Use map vector to make stable output.
   MapVector<Instruction *, Value *> ExtraArgs;
 
-  BinaryOperator *ReductionRoot = nullptr;
+  /// Contains info about operation, like its opcode, left and right operands.
+  struct OperationData {
+    /// true if the operation is a reduced value, false if reduction operation.
+    bool IsReducedValue = false;
+    /// Opcode of the instruction.
+    unsigned Opcode = 0;
+    /// Left operand of the reduction operation.
+    Value *LHS = nullptr;
+    /// Right operand of the reduction operation.
+    Value *RHS = nullptr;
+
+    /// Checks if the reduction operation can be vectorized.
+    bool isVectorizable() const {
+      return LHS && RHS &&
+             // We currently only support adds.
+             (Opcode == Instruction::Add || Opcode == Instruction::FAdd);
+    }
 
-  /// The opcode of the reduction.
-  Instruction::BinaryOps ReductionOpcode = Instruction::BinaryOpsEnd;
-  /// The opcode of the values we perform a reduction on.
-  unsigned ReducedValueOpcode = 0;
+  public:
+    explicit OperationData() = default;
+    /// Construction for reduced values. They are identified by opcode only and
+    /// don't have associated LHS/RHS values.
+    explicit OperationData(Value *V) : IsReducedValue(true) {
+      if (auto *I = dyn_cast<Instruction>(V))
+        Opcode = I->getOpcode();
+    }
+    /// Constructor for binary reduction operations with opcode and its left and
+    /// right operands.
+    OperationData(unsigned Opcode, Value *LHS, Value *RHS)
+        : IsReducedValue(false), Opcode(Opcode), LHS(LHS), RHS(RHS) {}
+    explicit operator bool() const { return Opcode; }
+    /// Get the index of the first operand.
+    unsigned getFirstOperandIndex() const {
+      assert(!!*this && "The opcode is not set.");
+      return 0;
+    }
+    /// Total number of operands in the reduction operation.
+    unsigned getNumberOfOperands() const {
+      assert(!IsReducedValue && !!*this && LHS && RHS &&
+             "Expected reduction operation.");
+      return 2;
+    }
+    /// Expected number of uses for reduction operations/reduced values.
+    unsigned getRequiredNumberOfUses() const {
+      assert(!IsReducedValue && !!*this && LHS && RHS &&
+             "Expected reduction operation.");
+      return 1;
+    }
+    /// Checks if instruction is associative and can be vectorized.
+    bool isAssociative(Instruction *I) const {
+      assert(!IsReducedValue && *this && LHS && RHS &&
+             "Expected reduction operation.");
+      return I->isAssociative();
+    }
+    /// Checks if the reduction operation can be vectorized.
+    bool isVectorizable(Instruction *I) const {
+      return isVectorizable() && isAssociative(I);
+    }
+
+    /// Checks if two operation data are both a reduction op or both a reduced
+    /// value.
+    bool operator==(const OperationData &OD) {
+      assert((IsReducedValue != OD.IsReducedValue) ||
+             ((!LHS == !OD.LHS) && (!RHS == !OD.RHS)) &&
+                 "One of the comparing operations is incorrect.");
+      return this == &OD ||
+             (IsReducedValue == OD.IsReducedValue && Opcode == OD.Opcode);
+    }
+    bool operator!=(const OperationData &OD) { return !(*this == OD); }
+    void clear() {
+      IsReducedValue = false;
+      Opcode = 0;
+      LHS = nullptr;
+      RHS = nullptr;
+    }
+    /// Get the opcode of the reduction operation.
+    unsigned getOpcode() const {
+      assert(isVectorizable() && "Expected vectorizable operation.");
+      return Opcode;
+    }
+    Value *getLHS() const { return LHS; }
+    Value *getRHS() const { return RHS; }
+    /// Creates reduction operation with the current opcode.
+    Value *createOp(IRBuilder<> &Builder, const Twine &Name = "") const {
+      assert(!IsReducedValue &&
+             (Opcode == Instruction::FAdd || Opcode == Instruction::Add) &&
+             "Expected add|fadd reduction operation.");
+      return Builder.CreateBinOp((Instruction::BinaryOps)Opcode, LHS, RHS,
+                                 Name);
+    }
+  };
+
+  Instruction *ReductionRoot = nullptr;
+
+  /// The operation data of the reduction operation.
+  OperationData ReductionData;
+  /// The operation data of the values we perform a reduction on.
+  OperationData ReducedValueData;
   /// Should we model this reduction as a pairwise reduction tree or a tree that
   /// splits the vector in halves and adds those halves.
   bool IsPairwiseReduction = false;
 
   /// Checks if the ParentStackElem.first should be marked as a reduction
   /// operation with an extra argument or as extra argument itself.
   void markExtraArg(std::pair<Instruction *, unsigned> &ParentStackElem,
                     Value *ExtraArg) {
     if (ExtraArgs.count(ParentStackElem.first)) {
       ExtraArgs[ParentStackElem.first] = nullptr;
       // We ran into something like:
       // ParentStackElem.first = ExtraArgs[ParentStackElem.first] + ExtraArg.
       // The whole ParentStackElem.first should be considered as an extra value
       // in this case.
       // Do not perform analysis of remaining operands of ParentStackElem.first
       // instruction, this whole instruction is an extra argument.
       ParentStackElem.second = ParentStackElem.first->getNumOperands();
     } else {
       // We ran into something like:
       // ParentStackElem.first += ... + ExtraArg + ...
       ExtraArgs[ParentStackElem.first] = ExtraArg;
     }
   }
 
+  static OperationData getOperationData(Value *V) {
+    if (!V)
+      return OperationData();
+
+    Value *LHS;
+    Value *RHS;
+    if (m_BinOp(m_Value(LHS), m_Value(RHS)).match(V))
+      return OperationData(cast<BinaryOperator>(V)->getOpcode(), LHS, RHS);
+    return OperationData(V);
+  }
+
 public:
   HorizontalReduction() = default;
 
   /// \brief Try to find a reduction tree.
-  bool matchAssociativeReduction(PHINode *Phi, BinaryOperator *B) {
+  bool matchAssociativeReduction(PHINode *Phi, Instruction *B) {
     assert((!Phi || is_contained(Phi->operands(), B)) &&
            "Thi phi needs to use the binary operator");
 
+    ReductionData = getOperationData(B);
+
     // We could have a initial reductions that is not an add.
     //  r *= v1 + v2 + v3 + v4
     // In such a case start looking for a tree rooted in the first '+'.
     if (Phi) {
-      if (B->getOperand(0) == Phi) {
+      if (ReductionData.getLHS() == Phi) {
         Phi = nullptr;
-        B = dyn_cast<BinaryOperator>(B->getOperand(1));
-      } else if (B->getOperand(1) == Phi) {
+        B = dyn_cast<Instruction>(ReductionData.getRHS());
+        ReductionData = getOperationData(B);
+      } else if (ReductionData.getRHS() == Phi) {
         Phi = nullptr;
-        B = dyn_cast<BinaryOperator>(B->getOperand(0));
+        B = dyn_cast<Instruction>(ReductionData.getLHS());
+        ReductionData = getOperationData(B);
       }
     }
 
-    if (!B)
+    if (!ReductionData.isVectorizable(B))
       return false;
 
     Type *Ty = B->getType();
     if (!isValidElementType(Ty))
       return false;
 
-    ReductionOpcode = B->getOpcode();
-    ReducedValueOpcode = 0;
+    ReducedValueData.clear();
     ReductionRoot = B;
 
-    // We currently only support adds.
-    if ((ReductionOpcode != Instruction::Add &&
-         ReductionOpcode != Instruction::FAdd) ||
-        !B->isAssociative())
-      return false;
-
     // Post order traverse the reduction tree starting at B. We only handle true
-    // trees containing only binary operators or selects.
+    // trees containing only binary operators.
     SmallVector<std::pair<Instruction *, unsigned>, 32> Stack;
-    Stack.push_back(std::make_pair(B, 0));
+    Stack.push_back(std::make_pair(B, ReductionData.getFirstOperandIndex()));
+    const unsigned NUses = ReductionData.getRequiredNumberOfUses();
     while (!Stack.empty()) {
       Instruction *TreeN = Stack.back().first;
       unsigned EdgeToVist = Stack.back().second++;
-      bool IsReducedValue = TreeN->getOpcode() != ReductionOpcode;
+      OperationData OpData = getOperationData(TreeN);
+      bool IsReducedValue = OpData != ReductionData;
 
       // Postorder vist.
-      if (EdgeToVist == 2 || IsReducedValue) {
+      if (IsReducedValue || EdgeToVist == OpData.getNumberOfOperands()) {
         if (IsReducedValue)
           ReducedVals.push_back(TreeN);
         else {
           auto I = ExtraArgs.find(TreeN);
           if (I != ExtraArgs.end() && !I->second) {
             // Check if TreeN is an extra argument of its parent operation.
             if (Stack.size() <= 1) {
               // TreeN can't be an extra argument as it is a root reduction
while (!Stack.empty()) {
         Stack.pop_back();
         continue;
       }
 
       // Visit left or right.
       Value *NextV = TreeN->getOperand(EdgeToVist);
       if (NextV != Phi) {
         auto *I = dyn_cast<Instruction>(NextV);
+        OpData = getOperationData(I);
         // 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 (I && (!ReducedValueData || OpData == ReducedValueData ||
+                  OpData == ReductionData)) {
           // Only handle trees in the current basic block.
           if (I->getParent() != B->getParent()) {
             // I is an extra argument for TreeN (its parent operation).
             markExtraArg(Stack.back(), I);
             continue;
           }
 
-          // Each tree node needs to have one user except for the ultimate
-          // reduction.
-          if (!I->hasOneUse() && I != B) {
+          // Each tree node needs to have minimal number of users except for the
+          // ultimate reduction.
+          if (!I->hasNUses(NUses) && I != B) {
             // I is an extra argument for TreeN (its parent operation).
             markExtraArg(Stack.back(), I);
             continue;
           }
 
-          if (I->getOpcode() == ReductionOpcode) {
+          if (OpData == ReductionData) {
             // We need to be able to reassociate the reduction operations.
-            if (!I->isAssociative()) {
+            if (!OpData.isAssociative(I)) {
               // I is an extra argument for TreeN (its parent operation).
               markExtraArg(Stack.back(), I);
               continue;
             }
-          } else if (ReducedValueOpcode &&
-                     ReducedValueOpcode != I->getOpcode()) {
+          } else if (ReducedValueData &&
+                     ReducedValueData != OpData) {
             // Make sure that the opcodes of the operations that we are going to
             // reduce match.
             // I is an extra argument for TreeN (its parent operation).
             markExtraArg(Stack.back(), I);
             continue;
-          } else if (!ReducedValueOpcode)
-            ReducedValueOpcode = I->getOpcode();
+          } else if (!ReducedValueData)
+            ReducedValueData = OpData;
 
-          Stack.push_back(std::make_pair(I, 0));
+          Stack.push_back(std::make_pair(I, OpData.getFirstOperandIndex()));
           continue;
         }
       }
       // NextV is an extra argument for TreeN (its parent operation).
       markExtraArg(Stack.back(), NextV);
     }
     return true;
   }
while (i < NumReducedVals - ReduxWidth + 1 && ReduxWidth > 2) {
       DebugLoc Loc = cast<Instruction>(ReducedVals[i])->getDebugLoc();
       Value *VectorizedRoot = V.vectorizeTree(ExternallyUsedValues);
 
       // Emit a reduction.
       Value *ReducedSubTree =
           emitReduction(VectorizedRoot, Builder, ReduxWidth, ReductionOps, TTI);
       if (VectorizedTree) {
         Builder.SetCurrentDebugLocation(Loc);
-        VectorizedTree = Builder.CreateBinOp(ReductionOpcode, VectorizedTree,
-                                             ReducedSubTree, "bin.rdx");
+        OperationData VectReductionData(ReductionData.getOpcode(),
+                                        VectorizedTree, ReducedSubTree);
+        VectorizedTree = VectReductionData.createOp(Builder, "bin.rdx");
         propagateIRFlags(VectorizedTree, ReductionOps);
       } else
         VectorizedTree = ReducedSubTree;
       i += ReduxWidth;
       ReduxWidth = PowerOf2Floor(NumReducedVals - i);
     }
 
     if (VectorizedTree) {
       // Finish the reduction.
       for (; i < NumReducedVals; ++i) {
         auto *I = cast<Instruction>(ReducedVals[i]);
         Builder.SetCurrentDebugLocation(I->getDebugLoc());
-        VectorizedTree =
-            Builder.CreateBinOp(ReductionOpcode, VectorizedTree, I);
+        OperationData VectReductionData(ReductionData.getOpcode(),
+                                        VectorizedTree, I);
+        VectorizedTree = VectReductionData.createOp(Builder);
         propagateIRFlags(VectorizedTree, ReductionOps);
       }
       for (auto &Pair : ExternallyUsedValues) {
         assert(!Pair.second.empty() &&
                "At least one DebugLoc must be inserted");
         // Add each externally used value to the final reduction.
         for (auto *I : Pair.second) {
           Builder.SetCurrentDebugLocation(I->getDebugLoc());
-          VectorizedTree = Builder.CreateBinOp(ReductionOpcode, VectorizedTree,
-                                               Pair.first, "bin.extra");
+          OperationData VectReductionData(ReductionData.getOpcode(),
+                                          VectorizedTree, Pair.first);
+          VectorizedTree = VectReductionData.createOp(Builder, "bin.extra");
           propagateIRFlags(VectorizedTree, I);
         }
       }
       // Update users.
       ReductionRoot->replaceAllUsesWith(VectorizedTree);
     }
     return VectorizedTree != nullptr;
   }
 
   unsigned numReductionValues() const {
     return ReducedVals.size();
   }
 
 private:
   /// \brief Calculate the cost of a reduction.
   int getReductionCost(TargetTransformInfo *TTI, Value *FirstReducedVal,
                        unsigned ReduxWidth) {
     Type *ScalarTy = FirstReducedVal->getType();
     Type *VecTy = VectorType::get(ScalarTy, ReduxWidth);
 
     int PairwiseRdxCost =
-        TTI->getArithmeticReductionCost(ReductionOpcode, VecTy, true);
+        TTI->getArithmeticReductionCost(ReductionData.getOpcode(), VecTy,
+                                        /*IsPairwiseForm=*/true);
     int SplittingRdxCost =
-        TTI->getArithmeticReductionCost(ReductionOpcode, VecTy, false);
+        TTI->getArithmeticReductionCost(ReductionData.getOpcode(), VecTy,
+                                        /*IsPairwiseForm=*/false);
 
     IsPairwiseReduction = PairwiseRdxCost < SplittingRdxCost;
     int VecReduxCost = IsPairwiseReduction ? PairwiseRdxCost : SplittingRdxCost;
 
     int ScalarReduxCost =
         (ReduxWidth - 1) *
-        TTI->getArithmeticInstrCost(ReductionOpcode, ScalarTy);
+        TTI->getArithmeticInstrCost(ReductionData.getOpcode(), ScalarTy);
 
     DEBUG(dbgs() << "SLP: Adding cost " << VecReduxCost - ScalarReduxCost
                  << " for reduction that starts with " << *FirstReducedVal
                  << " (It is a "
                  << (IsPairwiseReduction ? "pairwise" : "splitting")
                  << " reduction)\n");
 
     return VecReduxCost - ScalarReduxCost;
   }
 
   /// \brief Emit a horizontal reduction of the vectorized value.
   Value *emitReduction(Value *VectorizedValue, IRBuilder<> &Builder,
                        unsigned ReduxWidth, ArrayRef<Value *> RedOps,
                        const TargetTransformInfo *TTI) {
     assert(VectorizedValue && "Need to have a vectorized tree node");
     assert(isPowerOf2_32(ReduxWidth) &&
            "We only handle power-of-two reductions for now");
 
     if (!IsPairwiseReduction)
       return createSimpleTargetReduction(
-          Builder, TTI, ReductionOpcode, VectorizedValue,
+          Builder, TTI, ReductionData.getOpcode(), VectorizedValue,
           TargetTransformInfo::ReductionFlags(), RedOps);
 
     Value *TmpVec = VectorizedValue;
     for (unsigned i = ReduxWidth / 2; i != 0; i >>= 1) {
       Value *LeftMask =
           createRdxShuffleMask(ReduxWidth, i, true, true, Builder);
       Value *RightMask =
           createRdxShuffleMask(ReduxWidth, i, true, false, Builder);
 
       Value *LeftShuf = Builder.CreateShuffleVector(
           TmpVec, UndefValue::get(TmpVec->getType()), LeftMask, "rdx.shuf.l");
       Value *RightShuf = Builder.CreateShuffleVector(
           TmpVec, UndefValue::get(TmpVec->getType()), (RightMask),
           "rdx.shuf.r");
-      TmpVec =
-          Builder.CreateBinOp(ReductionOpcode, LeftShuf, RightShuf, "bin.rdx");
+      OperationData VectReductionData(ReductionData.getOpcode(), LeftShuf,
+                                      RightShuf);
+      TmpVec = VectReductionData.createOp(Builder, "bin.rdx");
       propagateIRFlags(TmpVec, RedOps);
     }
 
     // The result is in the first element of the vector.
     return Builder.CreateExtractElement(TmpVec, Builder.getInt32(0));
   }
 };
 } // end anonymous namespace