Repository: llvm/llvm-project
File: llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (C++)

Code before:
    // and the users outside the vector region of each reduction, at the
    // beginning of the dedicated latch block.
    auto *OrigExitingVPV = PhiR->getBackedgeValue();
    auto *NewExitingVPV = PhiR->getBackedgeValue();
    // Don't output selects for partial reductions because they have an output
    // with fewer lanes than the VF. So the operands of the select would have
    // different numbers of lanes. Partial reductions mask the input instead.
    auto *RR = dyn_cast<VPReductionRecipe>(OrigExitingVPV->getDefiningRecipe());
    if (!PhiR->isInLoop() && CM.foldTailByMasking() &&
        (!RR || !RR->isPartialReduction())) {
      VPValue *Cond = vputils::findHeaderMask(*Plan);
      NewExitingVPV =
          Builder.createSelect(Cond, OrigExitingVPV, PhiR, {}, "", *PhiR);
      OrigExitingVPV->replaceUsesWithIf(NewExitingVPV, [](VPUser &U, unsigned) {
        using namespace VPlanPatternMatch;
        return match(
            &U, m_CombineOr(
                    m_VPInstruction<VPInstruction::ComputeAnyOfResult>(),
                    m_VPInstruction<VPInstruction::ComputeReductionResult>()));
      });

      // Note: For FindLast recurrences we prefer a predicated select to
      // simplify matching in handleFindLastReductions(), rather than handle
      // multiple cases.
      if (CM.usePredicatedReductionSelect() ||
          RecurrenceDescriptor::isFindLastRecurrenceKind(RecurrenceKind))
        PhiR->setOperand(1, NewExitingVPV);
    }

    // We want code in the middle block to appear to execute on the location of
    // the scalar loop's latch terminator because: (a) it is all compiler
    // generated, (b) these instructions are always executed after evaluating
    // the latch conditional branch, and (c) other passes may add new
    // predecessors which terminate on this line. This is the easiest way to
    // ensure we don't accidentally cause an extra step back into the loop while
    // debugging.
    DebugLoc ExitDL = OrigLoop->getLoopLatch()->getTerminator()->getDebugLoc();

    // TODO: At the moment ComputeReductionResult also drives creation of the
    // bc.merge.rdx phi nodes, hence it needs to be created unconditionally here
    // even for in-loop reductions, until the reduction resume value handling is
    // also modeled in VPlan.
    VPInstruction *FinalReductionResult;
    VPBuilder::InsertPointGuard Guard(Builder);
    Builder.setInsertPoint(MiddleVPBB, IP);
    // For AnyOf reductions, find the select among PhiR's users. This is used
    // both to find NewVal for ComputeAnyOfResult and to adjust the reduction.
    VPRecipeBase *AnyOfSelect = nullptr;
    if (RecurrenceDescriptor::isAnyOfRecurrenceKind(RecurrenceKind)) {
      AnyOfSelect = cast<VPRecipeBase>(*find_if(PhiR->users(), [](VPUser *U) {
        return match(U, m_Select(m_VPValue(), m_VPValue(), m_VPValue()));

Review comment (suggestion):
Can we pass the recurrence kind to the `usePredicatedReductionSelect`, and have that check if this is a FindLastRecurrenceKind, then here we can just write:

```
if (CM.usePredicatedReductionSelect(RecurrenceKind))
  PhiR->setOperand(1, NewExitingVPV);
```
and for any other uses we don't have to distinguish this case (which might be error prone).

Code after:
      continue;

    RecurKind RecurrenceKind = PhiR->getRecurrenceKind();
    const RecurrenceDescriptor &RdxDesc = Legal->getRecurrenceDescriptor(
        cast<PHINode>(PhiR->getUnderlyingInstr()));
    Type *PhiTy = TypeInfo.inferScalarType(PhiR);
    // If tail is folded by masking, introduce selects between the phi
    // and the users outside the vector region of each reduction, at the
    // beginning of the dedicated latch block.
    auto *OrigExitingVPV = PhiR->getBackedgeValue();
    auto *NewExitingVPV = PhiR->getBackedgeValue();
    // Don't output selects for partial reductions because they have an output
    // with fewer lanes than the VF. So the operands of the select would have
    // different numbers of lanes. Partial reductions mask the input instead.
    auto *RR = dyn_cast<VPReductionRecipe>(OrigExitingVPV->getDefiningRecipe());
    if (!PhiR->isInLoop() && CM.foldTailByMasking() &&
        (!RR || !RR->isPartialReduction())) {
      VPValue *Cond = vputils::findHeaderMask(*Plan);
      NewExitingVPV =
          Builder.createSelect(Cond, OrigExitingVPV, PhiR, {}, "", *PhiR);
      OrigExitingVPV->replaceUsesWithIf(NewExitingVPV, [](VPUser &U, unsigned) {
        using namespace VPlanPatternMatch;
        return match(
            &U, m_CombineOr(
                    m_VPInstruction<VPInstruction::ComputeAnyOfResult>(),
                    m_VPInstruction<VPInstruction::ComputeReductionResult>()));
      });

      if (CM.usePredicatedReductionSelect(RecurrenceKind))
        PhiR->setOperand(1, NewExitingVPV);
    }

    // We want code in the middle block to appear to execute on the location of
    // the scalar loop's latch terminator because: (a) it is all compiler
    // generated, (b) these instructions are always executed after evaluating
    // the latch conditional branch, and (c) other passes may add new
    // predecessors which terminate on this line. This is the easiest way to
    // ensure we don't accidentally cause an extra step back into the loop while
    // debugging.
    DebugLoc ExitDL = OrigLoop->getLoopLatch()->getTerminator()->getDebugLoc();

    // TODO: At the moment ComputeReductionResult also drives creation of the
    // bc.merge.rdx phi nodes, hence it needs to be created unconditionally here
    // even for in-loop reductions, until the reduction resume value handling is
    // also modeled in VPlan.
    VPInstruction *FinalReductionResult;
    VPBuilder::InsertPointGuard Guard(Builder);
    Builder.setInsertPoint(MiddleVPBB, IP);
    // For AnyOf reductions, find the select among PhiR's users. This is used
    // both to find NewVal for ComputeAnyOfResult and to adjust the reduction.
    VPRecipeBase *AnyOfSelect = nullptr;