[llvm] [DAG] Use getLegalMaskAndStepVector in expandGetActiveLaneMask (PR #226506)
Luke Lau via llvm-commits
llvm-commits at lists.llvm.org
Fri Sep 25 07:27:15 PDT 2026
https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/226506
Currently we will either use the mask vector type if it's big enough to fit the elements, or promote the element width to the scalar operand type.
In the latter case we will end up with a vector element type larger than strictly necessary, which means a v32i1 mask will get promoted to v32i64 on RISC-V, which isn't a legal type, and expansion fails.
We can use a legal vector type that's just small enough via getLegalMaskAndStepVector, which is already used by expandVectorFindLastActive, expandCttzElts etc. This fixes crashes on RISC-V and also uses a more compact step vector. It also means we can handle scalable vectors.
We need to move the function into TargetLowering.cpp to use the getLegalMaskAndStepVector, but this brings it inline with the other expansions.
>From 3bf37375e4fab4267fbdaf3d613f23d36cfa031b Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 25 Sep 2026 12:54:23 +0100
Subject: [PATCH] [DAG] Use getLegalMaskAndStepVector in
expandGetActiveLaneMask
Currently we will either use the mask vector type if it's big enough to fit the elements, or promote the element width to the scalar operand type.
In the latter case we will end up with a vector element type larger than strictly necessary, which means a v32i1 mask will get promoted to v32i64 on RISC-V, which isn't a legal type, and expansion fails.
We can use a legal vector type that's just small enough via getLegalMaskAndStepVector, which is already used by expandVectorFindLastActive, expandCttzElts etc. This fixes crashes on RISC-V and also uses a more compact step vector. It also means we can handle scalable vectors.
We need to move the function into TargetLowering.cpp to use the getLegalMaskAndStepVector, but this brings it inline with the other expansions.
---
llvm/include/llvm/CodeGen/TargetLowering.h | 5 +++
.../SelectionDAG/LegalizeVectorOps.cpp | 44 +------------------
.../CodeGen/SelectionDAG/TargetLowering.cpp | 27 ++++++++++++
.../AArch64/neon-get-active-lane-mask.ll | 38 +++++++++-------
4 files changed, 54 insertions(+), 60 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index c08d0e53ec34b..6334955330c8d 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -5871,6 +5871,11 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
/// \returns The expansion result or SDValue() if it fails.
SDValue expandVectorFindLastActive(SDNode *N, SelectionDAG &DAG) const;
+ /// Expand GET_ACTIVE_LANE_MASK nodes
+ /// \param N Node to expand
+ /// \returns The expansion result or SDValue() if it fails.
+ SDValue expandGetActiveLaneMask(SDNode *N, SelectionDAG &DAG) const;
+
/// Expand LOOP_DEPENDENCE_MASK nodes
/// \param N Node to expand
/// \returns The expansion result or SDValue() if it fails.
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp
index edb8a971eef2b..3a5cfc00e1030 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp
@@ -135,7 +135,6 @@ class VectorLegalizer {
SDValue ExpandVSELECT(SDNode *Node);
SDValue ExpandVP_MERGE(SDNode *Node);
SDValue ExpandVP_REM(SDNode *Node);
- SDValue ExpandGET_ACTIVE_LANE_MASK(SDNode *N);
SDValue ExpandLOOP_DEPENDENCE_MASK(SDNode *N);
SDValue ExpandMaskedBinOp(SDNode *N);
SDValue ExpandSELECT(SDNode *Node);
@@ -1317,7 +1316,7 @@ void VectorLegalizer::Expand(SDNode *Node, SmallVectorImpl<SDValue> &Results) {
Results.push_back(TLI.expandCMP(Node, DAG));
return;
case ISD::GET_ACTIVE_LANE_MASK:
- if (SDValue R = ExpandGET_ACTIVE_LANE_MASK(Node))
+ if (SDValue R = TLI.expandGetActiveLaneMask(Node, DAG))
Results.push_back(R);
return;
case ISD::LOOP_DEPENDENCE_WAR_MASK:
@@ -1743,47 +1742,6 @@ SDValue VectorLegalizer::ExpandVP_REM(SDNode *Node) {
return DAG.getNode(ISD::SUB, DL, VT, Dividend, Mul);
}
-SDValue VectorLegalizer::ExpandGET_ACTIVE_LANE_MASK(SDNode *N) {
- SDLoc DL(N);
-
- SDValue Start = N->getOperand(0);
- SDValue End = N->getOperand(1);
- EVT VT = N->getValueType(0);
- EVT OpVT = Start.getValueType();
-
- if (VT.isScalableVector())
- return SDValue();
-
- // Try a promoted comparison type to simplify saturation.
- EVT PromoteVT = VT.changeVectorElementType(*DAG.getContext(), OpVT);
- if (TLI.isTypeLegal(PromoteVT) &&
- isUIntN(OpVT.getScalarSizeInBits(), VT.getVectorNumElements())) {
- SDValue StartV = DAG.getSplat(PromoteVT, DL, Start);
- SDValue Seq = DAG.getStepVector(DL, PromoteVT);
- Seq = DAG.getNode(ISD::UADDSAT, DL, PromoteVT, Seq, StartV);
-
- EVT MaskVT = TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(),
- PromoteVT);
- SDValue EndV = DAG.getSplat(PromoteVT, DL, End);
- SDValue Mask = DAG.getSetCC(DL, MaskVT, Seq, EndV, ISD::SETULT);
- return DAG.getBoolExtOrTrunc(Mask, DL, VT, PromoteVT);
- }
-
- // Is VT's element type big enough to hold all rebased indices?
- if (!isUIntN(VT.getScalarSizeInBits(), VT.getVectorNumElements()))
- return SDValue();
-
- // Rebase and saturate the termination value.
- SDValue Max = DAG.getConstant(maxUIntN(VT.getScalarSizeInBits()), DL, OpVT);
- End = DAG.getNode(ISD::USUBSAT, DL, OpVT, End, Start);
- End = DAG.getNode(ISD::UMIN, DL, OpVT, End, Max);
-
- // cmp <0, 1, 2, 3...>, End
- SDValue EndV = DAG.getSplat(VT, DL, End);
- SDValue StepVector = DAG.getStepVector(DL, VT);
- return DAG.getSetCC(DL, VT, StepVector, EndV, ISD::SETULT);
-}
-
SDValue VectorLegalizer::ExpandLOOP_DEPENDENCE_MASK(SDNode *N) {
return TLI.expandLoopDependenceMask(N, DAG);
}
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index f5402812659fe..2bf97fb58908a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -11258,6 +11258,33 @@ SDValue TargetLowering::expandVectorFindLastActive(SDNode *N,
return DAG.getZExtOrTrunc(HighestIdx, DL, N->getValueType(0));
}
+SDValue TargetLowering::expandGetActiveLaneMask(SDNode *N,
+ SelectionDAG &DAG) const {
+ SDLoc DL(N);
+
+ SDValue Start = N->getOperand(0);
+ SDValue End = N->getOperand(1);
+ EVT VT = N->getValueType(0);
+ EVT OpVT = Start.getValueType();
+
+ auto [Mask, StepVector] = getLegalMaskAndStepVector(
+ DAG.getPOISON(VT), /*ZeroIsPoison=*/false, DL, DAG);
+
+ // TODO: Handle the step vector needing to be split + widened.
+ assert(StepVector && "Step vector needs splitting");
+ assert(Mask.getValueType() == VT && "Step vector was widened");
+
+ // Rebase and saturate the termination value.
+ SDValue Max = DAG.getConstant(maxUIntN(StepVector.getScalarValueSizeInBits()),
+ DL, OpVT);
+ End = DAG.getNode(ISD::USUBSAT, DL, OpVT, End, Start);
+ End = DAG.getNode(ISD::UMIN, DL, OpVT, End, Max);
+
+ // cmp <0, 1, 2, 3...>, End
+ SDValue EndV = DAG.getSplat(StepVector.getValueType(), DL, End);
+ return DAG.getSetCC(DL, VT, StepVector, EndV, ISD::SETULT);
+}
+
SDValue TargetLowering::expandLoopDependenceMask(SDNode *N,
SelectionDAG &DAG) const {
SDLoc DL(N);
diff --git a/llvm/test/CodeGen/AArch64/neon-get-active-lane-mask.ll b/llvm/test/CodeGen/AArch64/neon-get-active-lane-mask.ll
index 36d8aaeb383f9..7a9d24659fc1b 100644
--- a/llvm/test/CodeGen/AArch64/neon-get-active-lane-mask.ll
+++ b/llvm/test/CodeGen/AArch64/neon-get-active-lane-mask.ll
@@ -91,13 +91,15 @@ define <8 x i1> @lane_mask_v8i1_i32(i32 %index, i32 %TC) {
define <4 x i1> @lane_mask_v4i1_i32(i32 %index, i32 %TC) {
; CHECK-LABEL: lane_mask_v4i1_i32:
; CHECK: // %bb.0:
+; CHECK-NEXT: subs w9, w1, w0
+; CHECK-NEXT: mov w8, #65535 // =0xffff
+; CHECK-NEXT: csel w9, wzr, w9, lo
+; CHECK-NEXT: cmp w9, w8
+; CHECK-NEXT: csel w8, w9, w8, lo
+; CHECK-NEXT: dup v0.4h, w8
; CHECK-NEXT: adrp x8, .LCPI3_0
-; CHECK-NEXT: dup v0.4s, w0
-; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI3_0]
-; CHECK-NEXT: uqadd v0.4s, v0.4s, v1.4s
-; CHECK-NEXT: dup v1.4s, w1
-; CHECK-NEXT: cmhi v0.4s, v1.4s, v0.4s
-; CHECK-NEXT: xtn v0.4h, v0.4s
+; CHECK-NEXT: ldr d1, [x8, :lo12:.LCPI3_0]
+; CHECK-NEXT: cmhi v0.4h, v0.4h, v1.4h
; CHECK-NEXT: ret
%active.lane.mask = call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 %index, i32 %TC)
ret <4 x i1> %active.lane.mask
@@ -106,12 +108,12 @@ define <4 x i1> @lane_mask_v4i1_i32(i32 %index, i32 %TC) {
define <2 x i1> @lane_mask_v2i1_i32(i32 %index, i32 %TC) {
; CHECK-LABEL: lane_mask_v2i1_i32:
; CHECK: // %bb.0:
+; CHECK-NEXT: subs w8, w1, w0
+; CHECK-NEXT: csel w8, wzr, w8, lo
+; CHECK-NEXT: dup v0.2s, w8
; CHECK-NEXT: mov x8, #4294967296 // =0x100000000
-; CHECK-NEXT: dup v0.2s, w0
; CHECK-NEXT: fmov d1, x8
-; CHECK-NEXT: uqadd v0.2s, v0.2s, v1.2s
-; CHECK-NEXT: dup v1.2s, w1
-; CHECK-NEXT: cmhi v0.2s, v1.2s, v0.2s
+; CHECK-NEXT: cmhi v0.2s, v0.2s, v1.2s
; CHECK-NEXT: ret
%active.lane.mask = call <2 x i1> @llvm.get.active.lane.mask.v2i1.i32(i32 %index, i32 %TC)
ret <2 x i1> %active.lane.mask
@@ -272,13 +274,15 @@ define <4 x i1> @lane_mask_v4i1_i64(i64 %index, i64 %TC) {
define <2 x i1> @lane_mask_v2i1_i64(i64 %index, i64 %TC) {
; CHECK-LABEL: lane_mask_v2i1_i64:
; CHECK: // %bb.0:
-; CHECK-NEXT: adrp x8, .LCPI10_0
-; CHECK-NEXT: dup v0.2d, x0
-; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI10_0]
-; CHECK-NEXT: uqadd v0.2d, v0.2d, v1.2d
-; CHECK-NEXT: dup v1.2d, x1
-; CHECK-NEXT: cmhi v0.2d, v1.2d, v0.2d
-; CHECK-NEXT: xtn v0.2s, v0.2d
+; CHECK-NEXT: subs x9, x1, x0
+; CHECK-NEXT: mov w8, #-1 // =0xffffffff
+; CHECK-NEXT: csel x9, xzr, x9, lo
+; CHECK-NEXT: cmp x9, x8
+; CHECK-NEXT: csel x8, x9, x8, lo
+; CHECK-NEXT: dup v0.2s, w8
+; CHECK-NEXT: mov x8, #4294967296 // =0x100000000
+; CHECK-NEXT: fmov d1, x8
+; CHECK-NEXT: cmhi v0.2s, v0.2s, v1.2s
; CHECK-NEXT: ret
%active.lane.mask = call <2 x i1> @llvm.get.active.lane.mask.v2i1.i64(i64 %index, i64 %TC)
ret <2 x i1> %active.lane.mask
More information about the llvm-commits
mailing list