[NFC][SLP] Cleanup: Simplify traversal loop in SLPVectorizerPass::vectorizeHorReduction().

[vporpo]

Summary

This includes a couple of changes:

1. Moves the code that changes the root node out of the TryToReduce lambda and out of the traversal loop.
2. Since that code moved, there isn't much left in TryToReduce so the code was inlined.
3. The phi node variable P was also being used as a flag that turns on/off the exploration of operands as new seeds. This patch uses a new variable TryOperandsAsNewSeeds for this.
4. Simplifies the code executed when vectorization fails.

The logic of the code should be identical to the original, but I may be missing something not caught by tests.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

13985

1. Make a separate NFC patch.
2. Format, the second parameter is not aligned

14027–14029

bool TryOperandsAsNewSeeds = P && isa<BinaryOperator>(Root);

14035–14042

if (P && isa<BinaryOperator>(Root) && HorizontalReduction::getRdxKind(Root) != RecurKind::None)
  Root = tryGetSecondaryReductionRoot(P, Root);

14071

What if Inst is one of the P operands here, but not Root?

14086

I think this is not quite same as before, it should not be a Root, it may be another P operand.

[vporpo]

Addressed comments.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

13985

OK I will do this in a separate patch.

14035–14042

I kept NewRoot because we need the original Root for Inst == Root later, otherwise the logic changes a bit. I also prefer to use TryOperandsAsNewSeeds instead of P as it shows intent. Same for isReductionCandidate() instead of isa<BinaryOperator>().

14071

This should work as the original code. This condition is the same as line 14028 and 14032.

14086

In the original code P was set to nullptr after visiting the root node, so this would only enter if Inst == Root.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14008

auto *Op

14008–14010

return dyn_cast<Instruction>(Op0 == Phi ? Op1 : Op0);

14070–14077

Better to keep this as lambda.

14086

The original code does not compare with Root, it compares with P. What if Inst is an instruction with P operand, but not Root?

14089

Why break here?

[vporpo]

Brought back the lambda and addressed comments.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14008–14010

Phi and Op are already Instruction * so I will use this instead: return Op == Phi ? dyn_cast<Instruction>(Val: Op1) : Op;

14086

P is set to null after visiting the root node (lines 14060, 14075 and 14081). So the code will enter this if only if Inst is the root node. If it is any other instruction it won't enter.

14089

This can only happen if Inst is the root node and at this point Stack is empty so it effectively breaks out of the loop, so it is better to be explicit about it.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14006

Better:

if (!match(...))
  return nullptr;

14008–14010

You just don't need Op at all, just
return dyn_cast<Instruction>(Op0 == Phi ? Op1 : Op0);
after assert is enough.

14085–14088

This logic also better to represent as lambda

14089

Add an assert(Stack.empty()) here and check with at least Spec that it does not crash

[vporpo]

Created lambda TryAppendToPostponedInsts and addressed comments.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14089

I checked with Spec 2017 and our presubmit tests and Stack is always empty.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14089

Add the assert here

[vporpo]

Added the assertion.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14034–14038

Can you turn into lambda too:

auto SelectRoot = [&]() {
  if (TryOperandsAsNewSeeds && isReductionCandidate(Root) &&
      HorizontalReduction::getRdxKind(Root) != RecurKind::None)
    return tryGetSecondaryReductionRoot(P, Root);
  return Root;
}
...
Stack.emplace(SelectRoot(), 0);

14036

TryOperandsAsNewSeeds && isReductionCandidate(Root) && to avoid extra call when TryOperandsAsNewSeeds is false.

14065–14066

auto TryAppendToPostponedInsts = [&]

[vporpo]

Addressed comments.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

Replace the message with something meaningful, like 'Empty Stack is expected.'

[vporpo]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

I think we should also mention why we are checking that Stack is empty, so that we know how to fix this if it triggers the assertion. How about: "Empty Stack expected. Change break to continue?"

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

The change continue to break is actual only for this patch, not for the codebase. So just Empty stack expected, the original root is reached or something like this should be enough

[vporpo]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

I don't understand what you mean. We are adding the assertion specifically for checking that changing continue to break was correct, and you are arguing that we should not mention this in the assertion message, and instead write some generic message that provides no information about why this assertion exists.
How will a message like "Empty stack expected, the original root is reached" is supposed to help a third person fix the code if the assertion triggers?

[vporpo]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

Anyway I don't feel to strongly about this I will change it to what you are proposing.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

The assert is not about the change itself. It should not inform about changes you made but about the precondition expected here. asserts are for precondition/postcondition messages, not change history.

[vporpo]

So a message like "expecting loop to exit" would

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

Yes the precondition is that the loop will exit right after this point, which explains why we are using break. Shouldn't we mention this in the message? I just don't like assertion messages that repeat the obvious like mentioning "Expected empty stack" when the assertion is assert(Stack.empty()).

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

You don't need to describe git history in the message, better describe better why Stack is expected to be empty (the original root is reached should help with this). Also, check https://llvm.org/docs/CodingStandards.html#assert-liberally and message examples.

[vporpo]

Updated assertion message. How about this?

[ABataev]

In D149627#4316334, @vporpo wrote:

Updated assertion message. How about this?

The break below already means it, better to describe why Stack is expected to be empty here.

[vporpo]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

The break below already means it, better to describe why Stack is expected to be empty here.

I don't understand why you don't like the message "Expected empty stack, loop should exit".
The precondition is that the loop should exit at this point. The fact that there is a break is exactly what is being checked by the assertion.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

It is obvious already, that the loop should exit.

[vporpo]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

14087

How is it obvious? This is exactly what the assertion is checking: if the assertion triggers it means that the loop should not exit and that break is wrong.

Anyway, I think we have already spent too much time on this and it is not that important. I will change it to "Expected stack is epxected" as you originally proposed, even though I think such messages are not very useful.

[vporpo]

Changed assertion message to "Expected empty stack".

[ABataev]

Still LG

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

if (match(I, m_Intrinsic<Intrinsic::umin>(m_Value(V0), m_Value(V1))))
     return true;
   return false;
 }
 
 /// We could have an initial reduction that is not an add.
 ///  r *= v1 + v2 + v3 + v4
 /// In such a case start looking for a tree rooted in the first '+'.
 /// \Returns the new root if found, which may be nullptr if not an instruction.
-static Instruction *tryGetScondaryReductionRoot(PHINode *Phi,
+static Instruction *tryGetSecondaryReductionRoot(PHINode *Phi,
                                                 Instruction *Root) {
   assert((isa<BinaryOperator>(Root) || isa<SelectInst>(Root) ||
           isa<IntrinsicInst>(Root)) &&
          "Expected binop, select, or intrinsic for reduction matching");
   Value *LHS =
       Root->getOperand(HorizontalReduction::getFirstOperandIndex(Root));
   Value *RHS =
       Root->getOperand(HorizontalReduction::getFirstOperandIndex(Root) + 1);
   if (LHS == Phi)
     return dyn_cast<Instruction>(RHS);
   if (RHS == Phi)
     return dyn_cast<Instruction>(LHS);
   return nullptr;
 }
 
+/// \p Returns the first operand of \p I that does not match \p Phi. If
+/// operand is not an instruction it returns nullptr.
+static Instruction *getNonPhiOperand(Instruction *I, PHINode *Phi) {
+  Value *Op0 = nullptr;
+  Value *Op1 = nullptr;
+  matchRdxBop(I, Op0, Op1);
+  assert(Op0 && Op1 && "Expected Binop");
+  Instruction *Op = dyn_cast<Instruction>(Op0);
+  if (Op == Phi)
+    Op = dyn_cast<Instruction>(Op1);
+  return Op;
+}
+
+/// \Returns true if \p I is a candidate instruction for reduction vectorization.
+static bool isReductionCandidate(Instruction *I) {
+  bool IsSelect = match(I, m_Select(m_Value(), m_Value(), m_Value()));
+  Value *B0 = nullptr, *B1 = nullptr;
+  bool IsBinop = matchRdxBop(I, B0, B1);
+  return IsBinop || IsSelect;
+}
+
 bool SLPVectorizerPass::vectorizeHorReduction(
     PHINode *P, Instruction *Root, BasicBlock *BB, BoUpSLP &R, TargetTransformInfo *TTI,
     SmallVectorImpl<WeakTrackingVH> &PostponedInsts) {
   if (!ShouldVectorizeHor)
     return false;
+  bool TryOperandsAsNewSeeds = P != nullptr;
   if (!isa<BinaryOperator>(Root))
-    P = nullptr;
+    TryOperandsAsNewSeeds = false;
 
   if (Root->getParent() != BB || isa<PHINode>(Root))
     return false;
+
+  // If we can find a secondary reduction root, use that instead.
+  Instruction *NewRoot = nullptr;
+  if (isReductionCandidate(Root)) {
+    assert((!TryOperandsAsNewSeeds || is_contained(P->operands(), Root)) &&
+           "Phi needs to use the binary operator");
+    if (TryOperandsAsNewSeeds &&
+        HorizontalReduction::getRdxKind(Root) != RecurKind::None)
+      NewRoot = tryGetSecondaryReductionRoot(P, Root);
+  }
+
   // Start analysis starting from Root instruction. If horizontal reduction is
   // found, try to vectorize it. If it is not a horizontal reduction or
   // vectorization is not possible or not effective, and currently analyzed
   // instruction is a binary operation, try to vectorize the operands, using
   // pre-order DFS traversal order. If the operands were not vectorized, repeat
   // the same procedure considering each operand as a possible root of the
   // horizontal reduction.
   // Interrupt the process if the Root instruction itself was vectorized or all
   // sub-trees not higher that RecursionMaxDepth were analyzed/vectorized.
   // If a horizintal reduction was not matched or vectorized we collect
   // instructions for possible later attempts for vectorization.
   std::queue<std::pair<Instruction *, unsigned>> Stack;
-  Stack.emplace(Root, 0);
+  Stack.emplace(NewRoot ? NewRoot : Root, 0);
   SmallPtrSet<Value *, 8> VisitedInstrs;
   bool Res = false;
-  auto &&TryToReduce = [this, TTI, &P, &R](Instruction *Inst, Value *&B0,
-                                           Value *&B1) -> Value * {
-    if (R.isAnalyzedReductionRoot(Inst))
-      return nullptr;
-    bool IsBinop = matchRdxBop(Inst, B0, B1);
-    bool IsSelect = match(Inst, m_Select(m_Value(), m_Value(), m_Value()));
-    if (IsBinop || IsSelect) {
-      assert((!P || is_contained(P->operands(), Inst)) &&
-             "Phi needs to use the binary operator");
-      if (P && HorizontalReduction::getRdxKind(Inst) != RecurKind::None)
-        if (Instruction *NewRoot = tryGetScondaryReductionRoot(P, Inst))
-          Inst = NewRoot;
-
-      HorizontalReduction HorRdx;
-      if (HorRdx.matchAssociativeReduction(Inst, *SE, *DL, *TLI))
-        return HorRdx.tryToReduce(R, TTI, *TLI);
-    }
-    return nullptr;
-  };
   while (!Stack.empty()) {
     Instruction *Inst;
     unsigned Level;
     std::tie(Inst, Level) = Stack.front();
     Stack.pop();
+
     // Do not try to analyze instruction that has already been vectorized.
     // This may happen when we vectorize instruction operands on a previous
     // iteration while stack was populated before that happened.
     if (R.isDeleted(Inst))
       continue;
-    Value *B0 = nullptr, *B1 = nullptr;
-    if (Value *V = TryToReduce(Inst, B0, B1)) {
+    Value *VectorizedV = nullptr;
+    if (!R.isAnalyzedReductionRoot(Inst) && isReductionCandidate(Inst)) {
+      HorizontalReduction HorRdx;
+      if (HorRdx.matchAssociativeReduction(Inst, *SE, *DL, *TLI))
+        VectorizedV = HorRdx.tryToReduce(R, TTI, *TLI);
+    }
+    if (VectorizedV != nullptr) {
       Res = true;
-      // Set P to nullptr to avoid re-analysis of phi node in
-      // matchAssociativeReduction function unless this is the root node.
-      P = nullptr;
-      if (auto *I = dyn_cast<Instruction>(V)) {
+      if (auto *I = dyn_cast<Instruction>(VectorizedV)) {
         // Try to find another reduction.
         Stack.emplace(I, Level);
         continue;
       }
     } else {
-      bool IsBinop = B0 && B1;
-      if (P && IsBinop) {
-        Inst = dyn_cast<Instruction>(B0);
-        if (Inst == P)
-          Inst = dyn_cast<Instruction>(B1);
-        if (!Inst) {
-          // Set P to nullptr to avoid re-analysis of phi node in
-          // matchAssociativeReduction function unless this is the root node.
-          P = nullptr;
-          continue;
-        }
+      // We could not vectorize `Inst` so try to use it as a future seed.
+      Instruction *FutureSeed = Inst;
+      if (TryOperandsAsNewSeeds && Inst == Root) {
+        FutureSeed = getNonPhiOperand(Root, P);
+        if (!FutureSeed)
+          break;
       }
-      // Set P to nullptr to avoid re-analysis of phi node in
-      // matchAssociativeReduction function unless this is the root node.
-      P = nullptr;
       // Do not collect CmpInst or InsertElementInst/InsertValueInst as their
       // analysis is done separately.
-      if (!isa<CmpInst, InsertElementInst, InsertValueInst>(Inst))
-        PostponedInsts.push_back(Inst);
+      if (!isa<CmpInst, InsertElementInst, InsertValueInst>(FutureSeed))
+        PostponedInsts.push_back(FutureSeed);
     }
 
     // Try to vectorize operands.
     // Continue analysis for the instruction from the same basic block only to
     // save compile time.
     if (++Level < RecursionMaxDepth)
       for (auto *Op : Inst->operand_values())
         if (VisitedInstrs.insert(Op).second)