[llvm-branch-commits] [llvm] [AArch64] Lower mask_beforefirst to brkb for NEON types on SVE (PR #232255)
Luke Lau via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Oct 15 03:13:49 PDT 2026
https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/232255
When SVE is available we can use lower fixed length mask_beforefirst calls to brkb, see the discussion in https://github.com/llvm/llvm-project/pull/230163#pullrequestreview-5469079195
Fixed vectors in streaming mode are left out of this PR for now.
>From 346439eb5b31cb806af2dd71fdf5d8b3f3629978 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Thu, 15 Oct 2026 13:30:15 +0800
Subject: [PATCH] [AArch64] Lower mask_beforefirst to brkb for NEON types on
SVE
When SVE is available we can use lower fixed length mask_beforefirst calls to brkb, see the discussion in https://github.com/llvm/llvm-project/pull/230163#pullrequestreview-5469079195
Fixed vectors in streaming mode are left out of this PR for now.
---
.../Target/AArch64/AArch64ISelLowering.cpp | 30 +++++++++-
llvm/lib/Target/AArch64/AArch64ISelLowering.h | 1 +
.../CodeGen/AArch64/mask-beforefirst-sve.ll | 55 ++++++-------------
3 files changed, 46 insertions(+), 40 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index f5f66db3cbd280..6dbd30af19ceaf 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1580,7 +1580,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
setOperationAction(ISD::GET_ACTIVE_LANE_MASK, MVT::nxv32i1, Custom);
for (auto VT : {MVT::v16i8, MVT::v8i8, MVT::v4i16, MVT::v2i32})
- setOperationAction(ISD::GET_ACTIVE_LANE_MASK, VT, Custom);
+ setOperationAction({ISD::GET_ACTIVE_LANE_MASK, ISD::MASK_BEFOREFIRST}, VT,
+ Custom);
for (auto VT : {MVT::v8f16, MVT::v4f32, MVT::v2f64})
setOperationAction(ISD::FMA, VT, Custom);
@@ -8996,6 +8997,8 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
return LowerVECTOR_INTERLEAVE(Op, DAG);
case ISD::GET_ACTIVE_LANE_MASK:
return LowerGET_ACTIVE_LANE_MASK(Op, DAG);
+ case ISD::MASK_BEFOREFIRST:
+ return LowerMASK_BEFOREFIRST(Op, DAG);
case ISD::VECTOR_MATCH:
return LowerVectorMatch(Op, DAG);
case ISD::LRINT:
@@ -35624,6 +35627,31 @@ AArch64TargetLowering::LowerGET_ACTIVE_LANE_MASK(SDValue Op,
return DAG.getExtractSubvector(DL, VT, MaskAsInt, 0);
}
+SDValue AArch64TargetLowering::LowerMASK_BEFOREFIRST(SDValue Op,
+ SelectionDAG &DAG) const {
+ assert(Subtarget->isSVEorStreamingSVEAvailable() &&
+ "Lowering fixed length mask_beforefirst requires SVE!");
+
+ SDLoc DL(Op);
+ EVT VT = Op.getValueType();
+ assert(VT.isFixedLengthVector() && "Expected fixed length vector type!");
+
+ SDValue Pg = getPredicateForFixedLengthVector(DAG, DL, VT);
+ SDValue Mask = convertFixedMaskToScalableVector(Op.getOperand(0), DAG);
+
+ // brkb only supports .b forms, so cast to make sure all our p regs match.
+ SDValue Brk = DAG.getNode(
+ ISD::INTRINSIC_WO_CHAIN, DL, MVT::nxv16i1,
+ DAG.getTargetConstant(Intrinsic::aarch64_sve_brkb_z, DL, MVT::i64),
+ getSVEPredicateBitCast(MVT::nxv16i1, Pg, DAG),
+ DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, MVT::nxv16i1, Mask));
+ Brk = DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, Pg.getValueType(), Brk);
+
+ EVT ContainerVT = getContainerForFixedLengthVector(DAG, VT);
+ SDValue MaskAsInt = DAG.getNode(ISD::SIGN_EXTEND, DL, ContainerVT, Brk);
+ return convertFromScalableVector(DAG, VT, MaskAsInt);
+}
+
static unsigned getSVEOpcodeForFPToInt(unsigned Opc) {
switch (Opc) {
case ISD::FP_TO_UINT_SAT:
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index a017e0508b57cf..59bd6fc4ee3e9d 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -781,6 +781,7 @@ class AArch64TargetLowering : public TargetLowering {
SDValue LowerVECTOR_HISTOGRAM(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerPARTIAL_REDUCE_MLA(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerGET_ACTIVE_LANE_MASK(SDValue Op, SelectionDAG &DAG) const;
+ SDValue LowerMASK_BEFOREFIRST(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerDIV(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerMUL(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerVectorSRA_SRL_SHL(SDValue Op, SelectionDAG &DAG) const;
diff --git a/llvm/test/CodeGen/AArch64/mask-beforefirst-sve.ll b/llvm/test/CodeGen/AArch64/mask-beforefirst-sve.ll
index 793a7a7303594c..9f9d9029ee8df7 100644
--- a/llvm/test/CodeGen/AArch64/mask-beforefirst-sve.ll
+++ b/llvm/test/CodeGen/AArch64/mask-beforefirst-sve.ll
@@ -41,22 +41,14 @@ define <vscale x 2 x i1> @nxv2i1(<vscale x 2 x i1> %m) {
ret <vscale x 2 x i1> %x
}
-; TODO: Use brkb for fixed length vectors.
-
define <16 x i1> @v16i1(<16 x i1> %m) {
; CHECK-LABEL: v16i1:
; CHECK: // %bb.0:
; CHECK-NEXT: shl v0.16b, v0.16b, #7
-; CHECK-NEXT: adrp x8, .LCPI4_0
-; CHECK-NEXT: mov w9, #16 // =0x10
-; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI4_0]
-; CHECK-NEXT: cmlt v0.16b, v0.16b, #0
-; CHECK-NEXT: and v0.16b, v0.16b, v1.16b
-; CHECK-NEXT: umaxv b0, v0.16b
-; CHECK-NEXT: fmov w8, s0
-; CHECK-NEXT: sub w8, w9, w8, uxtb
-; CHECK-NEXT: whilelo p0.b, xzr, x8
-; CHECK-NEXT: mov z0.b, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT: ptrue p0.b, vl16
+; CHECK-NEXT: cmpne p1.b, p0/z, z0.b, #0
+; CHECK-NEXT: brkb p1.b, p0/z, p1.b
+; CHECK-NEXT: mov z0.b, p1/z, #-1 // =0xffffffffffffffff
; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT: ret
%x = call <16 x i1> @llvm.mask.beforefirst(<16 x i1> %m)
@@ -67,15 +59,10 @@ define <8 x i1> @v8i1(<8 x i1> %m) {
; CHECK-LABEL: v8i1:
; CHECK: // %bb.0:
; CHECK-NEXT: shl v0.8b, v0.8b, #7
-; CHECK-NEXT: index z1.b, #8, #-1
-; CHECK-NEXT: mov w9, #8 // =0x8
-; CHECK-NEXT: cmlt v0.8b, v0.8b, #0
-; CHECK-NEXT: and v0.8b, v0.8b, v1.8b
-; CHECK-NEXT: umaxv b0, v0.8b
-; CHECK-NEXT: fmov w8, s0
-; CHECK-NEXT: sub w8, w9, w8, uxtb
-; CHECK-NEXT: whilelo p0.b, xzr, x8
-; CHECK-NEXT: mov z0.b, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT: ptrue p0.b, vl8
+; CHECK-NEXT: cmpne p1.b, p0/z, z0.b, #0
+; CHECK-NEXT: brkb p1.b, p0/z, p1.b
+; CHECK-NEXT: mov z0.b, p1/z, #-1 // =0xffffffffffffffff
; CHECK-NEXT: // kill: def $d0 killed $d0 killed $z0
; CHECK-NEXT: ret
%x = call <8 x i1> @llvm.mask.beforefirst(<8 x i1> %m)
@@ -86,15 +73,10 @@ define <4 x i1> @v4i1(<4 x i1> %m) {
; CHECK-LABEL: v4i1:
; CHECK: // %bb.0:
; CHECK-NEXT: shl v0.4h, v0.4h, #15
-; CHECK-NEXT: index z1.h, #4, #-1
-; CHECK-NEXT: mov w9, #4 // =0x4
-; CHECK-NEXT: cmlt v0.4h, v0.4h, #0
-; CHECK-NEXT: and v0.8b, v0.8b, v1.8b
-; CHECK-NEXT: umaxv h0, v0.4h
-; CHECK-NEXT: fmov w8, s0
-; CHECK-NEXT: sub w8, w9, w8, uxth
-; CHECK-NEXT: whilelo p0.h, xzr, x8
-; CHECK-NEXT: mov z0.h, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT: ptrue p0.h, vl4
+; CHECK-NEXT: cmpne p1.h, p0/z, z0.h, #0
+; CHECK-NEXT: brkb p1.b, p0/z, p1.b
+; CHECK-NEXT: mov z0.h, p1/z, #-1 // =0xffffffffffffffff
; CHECK-NEXT: // kill: def $d0 killed $d0 killed $z0
; CHECK-NEXT: ret
%x = call <4 x i1> @llvm.mask.beforefirst(<4 x i1> %m)
@@ -105,15 +87,10 @@ define <2 x i1> @v2i1(<2 x i1> %m) {
; CHECK-LABEL: v2i1:
; CHECK: // %bb.0:
; CHECK-NEXT: shl v0.2s, v0.2s, #31
-; CHECK-NEXT: index z1.s, #2, #-1
-; CHECK-NEXT: mov w9, #2 // =0x2
-; CHECK-NEXT: cmlt v0.2s, v0.2s, #0
-; CHECK-NEXT: and v0.8b, v0.8b, v1.8b
-; CHECK-NEXT: umaxp v0.2s, v0.2s, v0.2s
-; CHECK-NEXT: fmov w8, s0
-; CHECK-NEXT: sub w8, w9, w8
-; CHECK-NEXT: whilelo p0.s, xzr, x8
-; CHECK-NEXT: mov z0.s, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT: ptrue p0.s, vl2
+; CHECK-NEXT: cmpne p1.s, p0/z, z0.s, #0
+; CHECK-NEXT: brkb p1.b, p0/z, p1.b
+; CHECK-NEXT: mov z0.s, p1/z, #-1 // =0xffffffffffffffff
; CHECK-NEXT: // kill: def $d0 killed $d0 killed $z0
; CHECK-NEXT: ret
%x = call <2 x i1> @llvm.mask.beforefirst(<2 x i1> %m)
More information about the llvm-branch-commits
mailing list