Skip to content

[LV][NFC] Clean up tail-folding check for early-exit loops #133931

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 18, 2025
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1924,6 +1924,16 @@ bool LoopVectorizationLegality::canFoldTailByMasking() const {
}
}

// The only loops we can vectorize without a scalar epilogue, are loops with
// a bottom-test and a single exiting block. We'd have to handle the fact
// that not every instruction executes on the last iteration. This will
// require a lane mask which varies through the vector loop body. (TODO)
if (TheLoop->getExitingBlock() != TheLoop->getLoopLatch()) {
LLVM_DEBUG(
dbgs()
<< "LV: Cannot fold tail by masking. Requires a singe latch exit\n");
return false;
}
LLVM_DEBUG(dbgs() << "LV: can fold tail by masking.\n");

return true;
Expand Down
23 changes: 5 additions & 18 deletions llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3987,22 +3987,6 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
break;
}

// The only loops we can vectorize without a scalar epilogue, are loops with
// a bottom-test and a single exiting block. We'd have to handle the fact
// that not every instruction executes on the last iteration. This will
// require a lane mask which varies through the vector loop body. (TODO)
if (TheLoop->getExitingBlock() != TheLoop->getLoopLatch()) {
// If there was a tail-folding hint/switch, but we can't fold the tail by
// masking, fallback to a vectorization with a scalar epilogue.
if (ScalarEpilogueStatus == CM_ScalarEpilogueNotNeededUsePredicate) {
LLVM_DEBUG(dbgs() << "LV: Cannot fold tail by masking: vectorize with a "
"scalar epilogue instead.\n");
ScalarEpilogueStatus = CM_ScalarEpilogueAllowed;
return computeFeasibleMaxVF(MaxTC, UserVF, false);
}
return FixedScalableVFPair::getNone();
}

// Now try the tail folding

// Invalidate interleave groups that require an epilogue if we can't mask
Expand Down Expand Up @@ -4049,7 +4033,9 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
return Rem->isZero();
};

if (MaxPowerOf2RuntimeVF > 0u) {
bool HasSingleLatchExit =
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I understand what you're trying to do here, but I think we should remove the extra HasSingleLatchExit check and instead update NoScalarEpilogueNeeded for the case when a scalar epilogue is genuinely required. For loops with uncountable early exits we don't actually require a scalar epilogue at the moment, so we can still benefit from

  1. If we have a power-of-2 runtime VF then there is no point tail-folding.
  2. If we have a very low trip count we should fall back on a runtime power-of-2 fixed-width VF if possible.

In NoScalarEpilogueNeeded I think you can then add an extra check like this:

  auto NoScalarEpilogueNeeded = [this, &UserIC](unsigned MaxVF) {
    if (TheLoop->getExitingBlock() != TheLoop->getLoopLatch() &&
        !Legal->hasUncountableEarlyExit())
      return false;

What do you think?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! It does streamline the code.
I had been thinking that enabling the EVL transform would require going through setTailFoldingStyles(). But you're right — for loops where the trip count is known to be a multiple of the VF, using a non-predicated form should work well.

TheLoop->getExitingBlock() == TheLoop->getLoopLatch();
if (HasSingleLatchExit && MaxPowerOf2RuntimeVF > 0u) {
assert((UserVF.isNonZero() || isPowerOf2_32(*MaxPowerOf2RuntimeVF)) &&
"MaxFixedVF must be a power of 2");
if (NoScalarEpilogueNeeded(*MaxPowerOf2RuntimeVF)) {
Expand All @@ -4060,7 +4046,8 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
}

auto ExpectedTC = getSmallBestKnownTC(PSE, TheLoop);
if (ExpectedTC && ExpectedTC <= TTI.getMinTripCountTailFoldingThreshold()) {
if (HasSingleLatchExit && ExpectedTC &&
ExpectedTC <= TTI.getMinTripCountTailFoldingThreshold()) {
if (MaxPowerOf2RuntimeVF > 0u) {
// If we have a low-trip-count, and the fixed-width VF is known to divide
// the trip count but the scalable factor does not, use the fixed-width
Expand Down