Skip to content

[LV] Optimise latch exit induction users for some early exit loops #128880

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.

Sign up for GitHub

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 2 commits into from
Mar 11, 2025
Merged

[LV] Optimise latch exit induction users for some early exit loops #128880

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
69b05fe
[LV] Optimise latch exit induction users for some early exit loops
david-arm Feb 26, 2025
02e5552
Address review comments
david-arm Mar 3, 2025
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
117 changes: 63 additions & 54 deletions llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -730,67 +730,76 @@ static VPWidenInductionRecipe *getOptimizableIVOf(VPValue *VPV) {
return IsWideIVInc() ? WideIV : nullptr;
}

void VPlanTransforms::optimizeInductionExitUsers(
VPlan &Plan, DenseMap<VPValue *, VPValue *> &EndValues) {
/// Attempts to optimize the induction variable exit values for users in the
/// exit block coming from the latch in the original scalar loop.
static VPValue *
optimizeLatchExitInductionUser(VPlan &Plan, VPTypeAnalysis &TypeInfo,
VPBlockBase *PredVPBB, VPValue *Op,
DenseMap<VPValue *, VPValue *> &EndValues) {
using namespace VPlanPatternMatch;
SmallVector<VPIRBasicBlock *> ExitVPBBs(Plan.getExitBlocks());
if (ExitVPBBs.size() != 1)
return;

VPIRBasicBlock *ExitVPBB = ExitVPBBs[0];
VPBlockBase *PredVPBB = ExitVPBB->getSinglePredecessor();
if (!PredVPBB)
return;
assert(PredVPBB == Plan.getMiddleBlock() &&
"predecessor must be the middle block");

VPTypeAnalysis TypeInfo(Plan.getCanonicalIV()->getScalarType());
VPBuilder B(Plan.getMiddleBlock()->getTerminator());
for (VPRecipeBase &R : *ExitVPBB) {
auto *ExitIRI = cast<VPIRInstruction>(&R);
if (!isa<PHINode>(ExitIRI->getInstruction()))
break;
VPValue *Incoming;
if (!match(Op, m_VPInstruction<VPInstruction::ExtractFromEnd>(
m_VPValue(Incoming), m_SpecificInt(1))))
return nullptr;

VPValue *Incoming;
if (!match(ExitIRI->getOperand(0),
m_VPInstruction<VPInstruction::ExtractFromEnd>(
m_VPValue(Incoming), m_SpecificInt(1))))
continue;
auto *WideIV = getOptimizableIVOf(Incoming);
if (!WideIV)
return nullptr;

auto *WideIV = getOptimizableIVOf(Incoming);
if (!WideIV)
continue;
VPValue *EndValue = EndValues.lookup(WideIV);
assert(EndValue && "end value must have been pre-computed");
VPValue *EndValue = EndValues.lookup(WideIV);
Copy link
Member

Choose a reason for hiding this comment

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

I see the end values are populated in addScalarResumePhis(), am I right in thinking this map goes from IV -> exit IV value from scalar loop? Could this be non-constant? (in the changed tests it is always a constant).

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes in theory this could be non-constant, although since the tests are all testing users of the induction variable used to control the latch exit they end up being constant. This is because the only early exit loops we can currently vectorise are effectively ones with known trip counts, since the trip counts are used to prove that the loads always fall within the bounds of the object, such as an alloca. I'll see if I can add a test for a user of something that isn't the trip count.

assert(EndValue && "end value must have been pre-computed");

// `getOptimizableIVOf()` always returns the pre-incremented IV, so if it
// changed it means the exit is using the incremented value, so we don't
// need to subtract the step.
if (Incoming != WideIV)
return EndValue;

// Otherwise, subtract the step from the EndValue.
VPBuilder B(cast<VPBasicBlock>(PredVPBB)->getTerminator());
VPValue *Step = WideIV->getStepValue();
Type *ScalarTy = TypeInfo.inferScalarType(WideIV);
if (ScalarTy->isIntegerTy())
return B.createNaryOp(Instruction::Sub, {EndValue, Step}, {}, "ind.escape");
if (ScalarTy->isPointerTy()) {
auto *Zero = Plan.getOrAddLiveIn(
ConstantInt::get(Step->getLiveInIRValue()->getType(), 0));
return B.createPtrAdd(EndValue,
B.createNaryOp(Instruction::Sub, {Zero, Step}), {},
"ind.escape");
}
if (ScalarTy->isFloatingPointTy()) {
const auto &ID = WideIV->getInductionDescriptor();
return B.createNaryOp(
ID.getInductionBinOp()->getOpcode() == Instruction::FAdd
? Instruction::FSub
: Instruction::FAdd,
{EndValue, Step}, {ID.getInductionBinOp()->getFastMathFlags()});
}
llvm_unreachable("all possible induction types must be handled");
return nullptr;
}

if (Incoming != WideIV) {
ExitIRI->setOperand(0, EndValue);
continue;
}
void VPlanTransforms::optimizeInductionExitUsers(
VPlan &Plan, DenseMap<VPValue *, VPValue *> &EndValues) {
VPBlockBase *MiddleVPBB = Plan.getMiddleBlock();
VPTypeAnalysis TypeInfo(Plan.getCanonicalIV()->getScalarType());
for (VPIRBasicBlock *ExitVPBB : Plan.getExitBlocks()) {
for (VPRecipeBase &R : *ExitVPBB) {
auto *ExitIRI = cast<VPIRInstruction>(&R);
if (!isa<PHINode>(ExitIRI->getInstruction()))
break;

VPValue *Escape = nullptr;
VPValue *Step = WideIV->getStepValue();
Type *ScalarTy = TypeInfo.inferScalarType(WideIV);
if (ScalarTy->isIntegerTy()) {
Escape =
B.createNaryOp(Instruction::Sub, {EndValue, Step}, {}, "ind.escape");
} else if (ScalarTy->isPointerTy()) {
auto *Zero = Plan.getOrAddLiveIn(
ConstantInt::get(Step->getLiveInIRValue()->getType(), 0));
Escape = B.createPtrAdd(EndValue,
B.createNaryOp(Instruction::Sub, {Zero, Step}),
{}, "ind.escape");
} else if (ScalarTy->isFloatingPointTy()) {
const auto &ID = WideIV->getInductionDescriptor();
Escape = B.createNaryOp(
ID.getInductionBinOp()->getOpcode() == Instruction::FAdd
? Instruction::FSub
: Instruction::FAdd,
{EndValue, Step}, {ID.getInductionBinOp()->getFastMathFlags()});
} else {
llvm_unreachable("all possible induction types must be handled");
for (auto [Idx, PredVPBB] : enumerate(ExitVPBB->getPredecessors())) {
if (PredVPBB == MiddleVPBB)
if (VPValue *Escape = optimizeLatchExitInductionUser(
Plan, TypeInfo, PredVPBB, ExitIRI->getOperand(Idx),
EndValues))
ExitIRI->setOperand(Idx, Escape);
// TODO: Optimize early exit induction users in follow-on patch.
}
}
ExitIRI->setOperand(0, Escape);
}
}

Expand Down
Loading