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

Code before:
      if (R.mayWriteToMemory() && !match(&R, m_BranchOnCount()))
        return false;
    }
  }

  VPBasicBlock *LatchVPBB = LoopRegion->getExitingBasicBlock();
  VPBuilder LatchBuilder(LatchVPBB->getTerminator());
  VPValue *AllNaNLanes = nullptr;
  SmallPtrSet<VPValue *, 2> RdxResults;
  for (const auto &[_, MinOrMaxOp] : MinOrMaxNumReductionsToHandle) {
    VPValue *RedNaNLanes =
        LatchBuilder.createFCmp(CmpInst::FCMP_UNO, MinOrMaxOp, MinOrMaxOp);
    AllNaNLanes = AllNaNLanes ? LatchBuilder.createOr(AllNaNLanes, RedNaNLanes)
                              : RedNaNLanes;
  }

  VPValue *AnyNaNLane =
      LatchBuilder.createNaryOp(VPInstruction::AnyOf, {AllNaNLanes});
  VPBasicBlock *MiddleVPBB = Plan.getMiddleBlock();
  VPBuilder MiddleBuilder(MiddleVPBB, MiddleVPBB->begin());
  for (const auto &[RedPhiR, _] : MinOrMaxNumReductionsToHandle) {
    assert(RecurrenceDescriptor::isFPMinMaxNumRecurrenceKind(
               RedPhiR->getRecurrenceKind()) &&
           "unsupported reduction");

    // If we exit early due to NaNs, compute the final reduction result based on
    // the reduction phi at the beginning of the last vector iteration.
    auto *RdxResult = vputils::findComputeReductionResult(RedPhiR);
    assert(RdxResult && "must find a ComputeReductionResult");

    auto *NewSel = MiddleBuilder.createSelect(AnyNaNLane, RedPhiR,
                                              RdxResult->getOperand(0));
    RdxResult->setOperand(0, NewSel);
    assert(!RdxResults.contains(RdxResult) && "RdxResult already used");
    RdxResults.insert(RdxResult);
  }

  auto *LatchExitingBranch = LatchVPBB->getTerminator();
  assert(match(LatchExitingBranch, m_BranchOnCount(m_VPValue(), m_VPValue())) &&
         "Unexpected terminator");
  auto *IsLatchExitTaken = LatchBuilder.createICmp(
      CmpInst::ICMP_EQ, LatchExitingBranch->getOperand(0),
      LatchExitingBranch->getOperand(1));
  auto *AnyExitTaken = LatchBuilder.createOr(AnyNaNLane, IsLatchExitTaken);
  LatchBuilder.createNaryOp(VPInstruction::BranchOnCond, AnyExitTaken);
  LatchExitingBranch->eraseFromParent();

  // Update resume phis for inductions in the scalar preheader. If AnyNaNLane is
  // true, the resume from the start of the last vector iteration via the
  // canonical IV, otherwise from the original value.
  auto IsTC = [&Plan](VPValue *V) {

Review comment (bug):
I presume the only case where the IV operand of a VPDerivedIVRecipe would be `Plan.getTripCount()` is when using tail-folding? That then means the only way we could even enter the scalar preheader is if:

1. We skipped the tail-folded vector loop due to failing runtime checks, in which case we couldn't possibly be entering the scalar pre-header via the middle block and hence we're not resuming anything, or
2. The tail-folded vector loop deliberately exits the loop early with the intention of mopping up the remaining iterations in the scalar loop.

Is my understanding correct here?

Code after:
    // the reduction phi at the beginning of the last vector iteration.
    auto *RdxResult = vputils::findComputeReductionResult(RedPhiR);
    assert(RdxResult && "must find a ComputeReductionResult");

    auto *NewSel = MiddleBuilder.createSelect(AnyNaNLane, RedPhiR,
                                              RdxResult->getOperand(0));
    RdxResult->setOperand(0, NewSel);
    assert(!RdxResults.contains(RdxResult) && "RdxResult already used");
    RdxResults.insert(RdxResult);
  }

  auto *LatchExitingBranch = LatchVPBB->getTerminator();
  assert(match(LatchExitingBranch, m_BranchOnCount(m_VPValue(), m_VPValue())) &&
         "Unexpected terminator");
  auto *IsLatchExitTaken = LatchBuilder.createICmp(
      CmpInst::ICMP_EQ, LatchExitingBranch->getOperand(0),
      LatchExitingBranch->getOperand(1));
  auto *AnyExitTaken = LatchBuilder.createOr(AnyNaNLane, IsLatchExitTaken);
  LatchBuilder.createNaryOp(VPInstruction::BranchOnCond, AnyExitTaken);
  LatchExitingBranch->eraseFromParent();

  // Update resume phis for inductions in the scalar preheader. If AnyNaNLane is
  // true, the resume from the start of the last vector iteration via the
  // canonical IV, otherwise from the original value.
  auto IsTC = [&Plan](VPValue *V) {
    return V == &Plan.getVectorTripCount() || V == Plan.getTripCount();
  };
  for (auto &R : Plan.getScalarPreheader()->phis()) {
    auto *ResumeR = cast<VPPhi>(&R);
    VPValue *VecV = ResumeR->getOperand(0);
    if (RdxResults.contains(VecV))
      continue;
    if (auto *DerivedIV = dyn_cast<VPDerivedIVRecipe>(VecV)) {
      VPValue *DIVTC = DerivedIV->getOperand(1);
      if (DerivedIV->getNumUsers() == 1 && IsTC(DIVTC)) {
        auto *NewSel = MiddleBuilder.createSelect(
            AnyNaNLane, LoopRegion->getCanonicalIV(), DIVTC);
        DerivedIV->moveAfter(&*MiddleBuilder.getInsertPoint());
        DerivedIV->setOperand(1, NewSel);
        continue;
      }
    }
    // Bail out and abandon the current, partially modified, VPlan if we
    // encounter resume phi that cannot be updated yet.
    if (!IsTC(VecV)) {
      LLVM_DEBUG(dbgs() << "Found resume phi we cannot update for VPlan with "
                           "FMaxNum/FMinNum reduction.\n");
      return false;
    }
    auto *NewSel = MiddleBuilder.createSelect(
        AnyNaNLane, LoopRegion->getCanonicalIV(), VecV);