[llvm] [LV][NFC] Refactor code for extracting first active element (PR #131118)
David Sherwood via llvm-commits
llvm-commits at lists.llvm.org
Thu Mar 13 06:58:21 PDT 2025
https://github.com/david-arm updated https://github.com/llvm/llvm-project/pull/131118
>From 5de057947662b4c9de71664dfc751b750e8eb596 Mon Sep 17 00:00:00 2001
From: David Sherwood <david.sherwood at arm.com>
Date: Thu, 13 Mar 2025 09:57:01 +0000
Subject: [PATCH] [LV][NFC] Refactor code for extracting first active element
Refactor the code to extract the first active element of a
vector in the early exit block, in preparation for PR #130766.
I've replaced the VPInstruction::ExtractFirstActive nodes with
a combination of a new VPInstruction::FirstActiveLane node and
a Instruction::ExtractElement node.
---
llvm/lib/Transforms/Vectorize/VPlan.h | 5 +--
.../Transforms/Vectorize/VPlanAnalysis.cpp | 5 ++-
.../lib/Transforms/Vectorize/VPlanRecipes.cpp | 42 +++++++++++--------
.../Transforms/Vectorize/VPlanTransforms.cpp | 12 ++++--
.../LoopVectorize/AArch64/early_exit_costs.ll | 12 ++++--
5 files changed, 46 insertions(+), 30 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 43dc30c40bb53..0b63dba76ff9b 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -879,9 +879,8 @@ class VPInstruction : public VPRecipeWithIRFlags,
// Returns a scalar boolean value, which is true if any lane of its (only
// boolean) vector operand is true.
AnyOf,
- // Extracts the first active lane of a vector, where the first operand is
- // the predicate, and the second operand is the vector to extract.
- ExtractFirstActive,
+ // Calculates the first active lane index of the vector predicate operand.
+ FirstActiveLane,
};
private:
diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index 6f6875f0e5e0e..031428b3e797f 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -50,6 +50,8 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
return SetResultTyFromOp();
switch (Opcode) {
+ case Instruction::ExtractElement:
+ return inferScalarType(R->getOperand(0));
case Instruction::Select: {
Type *ResTy = inferScalarType(R->getOperand(1));
VPValue *OtherV = R->getOperand(2);
@@ -78,7 +80,8 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
case VPInstruction::CanonicalIVIncrementForPart:
case VPInstruction::AnyOf:
return SetResultTyFromOp();
- case VPInstruction::ExtractFirstActive:
+ case VPInstruction::FirstActiveLane:
+ return Type::getIntNTy(Ctx, 64);
case VPInstruction::ExtractFromEnd: {
Type *BaseTy = inferScalarType(R->getOperand(0));
if (auto *VecTy = dyn_cast<VectorType>(BaseTy))
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 0d3c9aeec6be4..db8fb70e2ea81 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -461,6 +461,11 @@ Value *VPInstruction::generate(VPTransformState &State) {
Value *A = State.get(getOperand(0));
return Builder.CreateNot(A, Name);
}
+ case Instruction::ExtractElement: {
+ Value *Vec = State.get(getOperand(0));
+ Value *Idx = State.get(getOperand(1), true);
+ return Builder.CreateExtractElement(Vec, Idx, Name);
+ }
case Instruction::ICmp: {
bool OnlyFirstLaneUsed = vputils::onlyFirstLaneUsed(this);
Value *A = State.get(getOperand(0), OnlyFirstLaneUsed);
@@ -705,12 +710,10 @@ Value *VPInstruction::generate(VPTransformState &State) {
Value *A = State.get(getOperand(0));
return Builder.CreateOrReduce(A);
}
- case VPInstruction::ExtractFirstActive: {
- Value *Vec = State.get(getOperand(0));
- Value *Mask = State.get(getOperand(1));
- Value *Ctz = Builder.CreateCountTrailingZeroElems(
- Builder.getInt64Ty(), Mask, true, "first.active.lane");
- return Builder.CreateExtractElement(Vec, Ctz, "early.exit.value");
+ case VPInstruction::FirstActiveLane: {
+ Value *Mask = State.get(getOperand(0));
+ return Builder.CreateCountTrailingZeroElems(Builder.getInt64Ty(), Mask,
+ true, Name);
}
default:
llvm_unreachable("Unsupported opcode for instruction");
@@ -737,22 +740,24 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
}
switch (getOpcode()) {
+ case Instruction::ExtractElement: {
+ // Add on the cost of extracting the element.
+ auto *VecTy = toVectorTy(Ctx.Types.inferScalarType(getOperand(0)), VF);
+ return Ctx.TTI.getVectorInstrCost(Instruction::ExtractElement, VecTy,
+ Ctx.CostKind);
+ }
case VPInstruction::AnyOf: {
auto *VecTy = toVectorTy(Ctx.Types.inferScalarType(this), VF);
return Ctx.TTI.getArithmeticReductionCost(
Instruction::Or, cast<VectorType>(VecTy), std::nullopt, Ctx.CostKind);
}
- case VPInstruction::ExtractFirstActive: {
+ case VPInstruction::FirstActiveLane: {
// Calculate the cost of determining the lane index.
- auto *PredTy = toVectorTy(Ctx.Types.inferScalarType(getOperand(1)), VF);
+ auto *PredTy = toVectorTy(Ctx.Types.inferScalarType(getOperand(0)), VF);
IntrinsicCostAttributes Attrs(Intrinsic::experimental_cttz_elts,
Type::getInt64Ty(Ctx.LLVMCtx),
{PredTy, Type::getInt1Ty(Ctx.LLVMCtx)});
- InstructionCost Cost = Ctx.TTI.getIntrinsicInstrCost(Attrs, Ctx.CostKind);
- // Add on the cost of extracting the element.
- auto *VecTy = toVectorTy(Ctx.Types.inferScalarType(getOperand(0)), VF);
- return Cost + Ctx.TTI.getVectorInstrCost(Instruction::ExtractElement, VecTy,
- Ctx.CostKind);
+ return Ctx.TTI.getIntrinsicInstrCost(Attrs, Ctx.CostKind);
}
default:
// TODO: Compute cost other VPInstructions once the legacy cost model has
@@ -765,7 +770,8 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
bool VPInstruction::isVectorToScalar() const {
return getOpcode() == VPInstruction::ExtractFromEnd ||
- getOpcode() == VPInstruction::ExtractFirstActive ||
+ getOpcode() == Instruction::ExtractElement ||
+ getOpcode() == VPInstruction::FirstActiveLane ||
getOpcode() == VPInstruction::ComputeReductionResult ||
getOpcode() == VPInstruction::AnyOf;
}
@@ -824,13 +830,14 @@ bool VPInstruction::opcodeMayReadOrWriteFromMemory() const {
if (Instruction::isBinaryOp(getOpcode()))
return false;
switch (getOpcode()) {
+ case Instruction::ExtractElement:
case Instruction::ICmp:
case Instruction::Select:
case VPInstruction::AnyOf:
case VPInstruction::CalculateTripCountMinusVF:
case VPInstruction::CanonicalIVIncrementForPart:
case VPInstruction::ExtractFromEnd:
- case VPInstruction::ExtractFirstActive:
+ case VPInstruction::FirstActiveLane:
case VPInstruction::FirstOrderRecurrenceSplice:
case VPInstruction::LogicalAnd:
case VPInstruction::Not:
@@ -939,7 +946,6 @@ void VPInstruction::print(raw_ostream &O, const Twine &Indent,
case VPInstruction::Broadcast:
O << "broadcast";
break;
-
case VPInstruction::ExtractFromEnd:
O << "extract-from-end";
break;
@@ -955,8 +961,8 @@ void VPInstruction::print(raw_ostream &O, const Twine &Indent,
case VPInstruction::AnyOf:
O << "any-of";
break;
- case VPInstruction::ExtractFirstActive:
- O << "extract-first-active";
+ case VPInstruction::FirstActiveLane:
+ O << "first-active-lane";
break;
default:
O << Instruction::getOpcodeName(getOpcode());
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index f23795f87d467..c15bcfd1f36a5 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2157,10 +2157,14 @@ void VPlanTransforms::handleUncountableEarlyExit(
ExitIRI->extractLastLaneOfOperand(MiddleBuilder);
}
// Add the incoming value from the early exit.
- if (!IncomingFromEarlyExit->isLiveIn())
- IncomingFromEarlyExit =
- EarlyExitB.createNaryOp(VPInstruction::ExtractFirstActive,
- {IncomingFromEarlyExit, EarlyExitTakenCond});
+ if (!IncomingFromEarlyExit->isLiveIn()) {
+ VPValue *FirstActiveLane = EarlyExitB.createNaryOp(
+ VPInstruction::FirstActiveLane, {EarlyExitTakenCond}, nullptr,
+ "first.active.lane");
+ IncomingFromEarlyExit = EarlyExitB.createNaryOp(
+ Instruction::ExtractElement, {IncomingFromEarlyExit, FirstActiveLane},
+ nullptr, "early.exit.value");
+ }
ExitIRI->addOperand(IncomingFromEarlyExit);
}
MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {IsEarlyExitTaken});
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll b/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll
index 55c6c43b6306a..a9c9af2c99932 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll
@@ -11,8 +11,10 @@ define i64 @same_exit_block_pre_inc_use1_sve() #1 {
; CHECK-LABEL: LV: Checking a loop in 'same_exit_block_pre_inc_use1_sve'
; CHECK: LV: Selecting VF: vscale x 16
; CHECK: Calculating cost of work in exit block vector.early.exit
-; CHECK-NEXT: Cost of 6 for VF vscale x 16: EMIT vp<{{.*}}> = extract-first-active
-; CHECK-NEXT: Cost of 6 for VF vscale x 16: EMIT vp<{{.*}}> = extract-first-active
+; CHECK-NEXT: Cost of 4 for VF vscale x 16: EMIT vp<{{.*}}> = first-active-lane vp<{{.*}}>
+; CHECK-NEXT: Cost of 2 for VF vscale x 16: EMIT vp<{{.*}}> = extractelement ir<{{.*}}>, vp<{{.*}}>
+; CHECK-NEXT: Cost of 4 for VF vscale x 16: EMIT vp<{{.*}}>.1 = first-active-lane vp<{{.*}}>
+; CHECK-NEXT: Cost of 2 for VF vscale x 16: EMIT vp<{{.*}}>.1 = extractelement ir<{{.*}}>, vp<%first.active.lane>.1
; CHECK: LV: Minimum required TC for runtime checks to be profitable:32
entry:
%p1 = alloca [1024 x i8]
@@ -48,8 +50,10 @@ define i64 @same_exit_block_pre_inc_use1_nosve() {
; CHECK-LABEL: LV: Checking a loop in 'same_exit_block_pre_inc_use1_nosve'
; CHECK: LV: Selecting VF: 16
; CHECK: Calculating cost of work in exit block vector.early.exit
-; CHECK-NEXT: Cost of 50 for VF 16: EMIT vp<{{.*}}> = extract-first-active
-; CHECK-NEXT: Cost of 50 for VF 16: EMIT vp<{{.*}}> = extract-first-active
+; CHECK-NEXT: Cost of 48 for VF 16: EMIT vp<{{.*}}> = first-active-lane vp<{{.*}}>
+; CHECK-NEXT: Cost of 2 for VF 16: EMIT vp<{{.*}}> = extractelement ir<{{.*}}>, vp<{{.*}}>
+; CHECK-NEXT: Cost of 48 for VF 16: EMIT vp<{{.*}}>.1 = first-active-lane vp<{{.*}}>
+; CHECK-NEXT: Cost of 2 for VF 16: EMIT vp<{{.*}}>.1 = extractelement ir<{{.*}}>, vp<%first.active.lane>.1
; CHECK: LV: Minimum required TC for runtime checks to be profitable:176
; CHECK-NEXT: LV: Vectorization is not beneficial: expected trip count < minimum profitable VF (64 < 176)
; CHECK-NEXT: LV: Too many memory checks needed.
More information about the llvm-commits
mailing list