[llvm] [AArch64][ISel] Enable masked interleaved stores for splat values (PR #207950)
Harry Ramsey via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 13 02:00:59 PDT 2026
https://github.com/Harry-Ramsey updated https://github.com/llvm/llvm-project/pull/207950
>From 00c02082cb283d4e14968009c16b645dc1a35f11 Mon Sep 17 00:00:00 2001
From: Harry Ramsey <harry.ramsey at arm.com>
Date: Mon, 6 Jul 2026 15:55:59 +0000
Subject: [PATCH 1/2] [AArch64][ISel] Enable masked interleaved stores for
splat values
Enable masked interleaved store combine to recognise splat values and
split them into SVE component vectors. This allows the existing stN
lowering path to handle masked stores where the stored value is a splat
rather than an explicit vector.interleave result.
---
.../scalable_masked_interleaved_stores.ll | 72 +++++++++++++++++++
1 file changed, 72 insertions(+)
diff --git a/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll b/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll
index d3cd9bf08cc0a..bdc25bba711e1 100644
--- a/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll
+++ b/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll
@@ -282,3 +282,75 @@ define void @foo_st2_nxv8i8_trunc(<vscale x 4 x i1> %mask, <vscale x 4 x i16> %v
call void @llvm.masked.store.nxv8i8.p0(<vscale x 8 x i8> %trunc.value, ptr %p, i32 1, <vscale x 8 x i1> %interleaved.mask)
ret void
}
+
+define void @foo_st2_nxv16i8_zeroinitializer(<vscale x 16 x i1> %mask, ptr %p) {
+; CHECK-LABEL: foo_st2_nxv16i8_zeroinitializer:
+; CHECK: // %bb.0:
+; CHECK-NEXT: movi v0.2d, #0000000000000000
+; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
+; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
+; CHECK-NEXT: st1b { z0.b }, p1, [x0, #1, mul vl]
+; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: ret
+ %interleaved.mask = call <vscale x 32 x i1> @llvm.vector.interleave2.nxv32i1(<vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask)
+ call void @llvm.masked.store.nxv32i8.p0(<vscale x 32 x i8> zeroinitializer, ptr %p, i32 1, <vscale x 32 x i1> %interleaved.mask)
+ ret void
+}
+
+define void @foo_st4_nxv16i8_zeroinitializer(<vscale x 16 x i1> %mask, ptr %p) {
+; CHECK-LABEL: foo_st4_nxv16i8_zeroinitializer:
+; CHECK: // %bb.0:
+; CHECK-NEXT: movi v0.2d, #0000000000000000
+; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
+; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
+; CHECK-NEXT: zip2 p2.b, p1.b, p1.b
+; CHECK-NEXT: zip1 p1.b, p1.b, p1.b
+; CHECK-NEXT: zip2 p3.b, p0.b, p0.b
+; CHECK-NEXT: st1b { z0.b }, p2, [x0, #3, mul vl]
+; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
+; CHECK-NEXT: st1b { z0.b }, p1, [x0, #2, mul vl]
+; CHECK-NEXT: st1b { z0.b }, p3, [x0, #1, mul vl]
+; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: ret
+ %interleaved.mask = call <vscale x 64 x i1> @llvm.vector.interleave4.nxv64i1(<vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask)
+ call void @llvm.masked.store.nxv64i8.p0(<vscale x 64 x i8> zeroinitializer, ptr %p, i32 1, <vscale x 64 x i1> %interleaved.mask)
+ ret void
+}
+
+define void @foo_st2_nxv16i8_splat(<vscale x 16 x i1> %mask, ptr %p) {
+; CHECK-LABEL: foo_st2_nxv16i8_splat:
+; CHECK: // %bb.0:
+; CHECK-NEXT: mov z0.b, #1 // =0x1
+; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
+; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
+; CHECK-NEXT: st1b { z0.b }, p1, [x0, #1, mul vl]
+; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: ret
+ %base = insertelement <vscale x 32 x i8> poison, i8 1, i32 0
+ %interleaved.value = shufflevector <vscale x 32 x i8> %base, <vscale x 32 x i8> poison, <vscale x 32 x i32> zeroinitializer
+ %interleaved.mask = call <vscale x 32 x i1> @llvm.vector.interleave2.nxv32i1(<vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask)
+ call void @llvm.masked.store.nxv32i8.p0(<vscale x 32 x i8> %interleaved.value, ptr %p, i32 1, <vscale x 32 x i1> %interleaved.mask)
+ ret void
+}
+
+define void @foo_st4_nxv16i8_splat(<vscale x 16 x i1> %mask, ptr %p) {
+; CHECK-LABEL: foo_st4_nxv16i8_splat:
+; CHECK: // %bb.0:
+; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
+; CHECK-NEXT: mov z0.b, #1 // =0x1
+; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
+; CHECK-NEXT: zip2 p2.b, p1.b, p1.b
+; CHECK-NEXT: zip1 p1.b, p1.b, p1.b
+; CHECK-NEXT: zip2 p3.b, p0.b, p0.b
+; CHECK-NEXT: st1b { z0.b }, p2, [x0, #3, mul vl]
+; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
+; CHECK-NEXT: st1b { z0.b }, p1, [x0, #2, mul vl]
+; CHECK-NEXT: st1b { z0.b }, p3, [x0, #1, mul vl]
+; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: ret
+ %base = insertelement <vscale x 64 x i8> poison, i8 1, i32 0
+ %interleaved.value = shufflevector <vscale x 64 x i8> %base, <vscale x 64 x i8> poison, <vscale x 64 x i32> zeroinitializer
+ %interleaved.mask = call <vscale x 64 x i1> @llvm.vector.interleave4.nxv64i1(<vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask)
+ call void @llvm.masked.store.nxv64i8.p0(<vscale x 64 x i8> %interleaved.value, ptr %p, i32 1, <vscale x 64 x i1> %interleaved.mask)
+ ret void
+}
>From 7d31a5f9e39e387503b94ca0cc422b3ba9050704 Mon Sep 17 00:00:00 2001
From: Harry Ramsey <harry.ramsey at arm.com>
Date: Mon, 6 Jul 2026 15:56:09 +0000
Subject: [PATCH 2/2] fixup! [AArch64][ISel] Enable masked interleaved stores
for splat values
---
.../Target/AArch64/AArch64ISelLowering.cpp | 59 ++++++++++++++++---
.../scalable_masked_interleaved_stores.ll | 40 ++++---------
2 files changed, 63 insertions(+), 36 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 7df8aa045dbf9..67c5aa045b866 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -27248,6 +27248,43 @@ isSequentialConcatOfVectorInterleave(SDNode *N, SmallVectorImpl<SDValue> &Ops) {
return true;
}
+static bool isSplatVectorInterleaveOps(SelectionDAG &DAG, SDLoc DL,
+ SDValue WideValue,
+ SmallVectorImpl<SDValue> &Ops) {
+ EVT WideVT = WideValue.getValueType();
+ if (!WideVT.isScalableVector())
+ return false;
+
+ SDValue SplatValue = DAG.getSplatValue(WideValue);
+ if (!SplatValue)
+ return false;
+
+ TypeSize WideSize = WideVT.getSizeInBits();
+ if (!WideSize.isKnownMultipleOf(128))
+ return false;
+
+ unsigned NumParts = WideSize.getKnownMinValue() / 128;
+ if (NumParts != 2 && NumParts != 3 && NumParts != 4)
+ return false;
+
+ ElementCount WideEC = WideVT.getVectorElementCount();
+ if (!WideEC.isKnownMultipleOf(NumParts))
+ return false;
+
+ EVT SubVecTy = EVT::getVectorVT(*DAG.getContext(),
+ WideVT.getVectorElementType(),
+ WideEC.divideCoefficientBy(NumParts));
+
+ if (SubVecTy.getSizeInBits().getKnownMinValue() != 128 ||
+ !DAG.getTargetLoweringInfo().isTypeLegal(SubVecTy))
+ return false;
+
+ SDValue NarrowSplat =
+ DAG.getNode(ISD::SPLAT_VECTOR, DL, SubVecTy, SplatValue);
+ Ops.append(NumParts, NarrowSplat);
+ return true;
+}
+
static SDValue getNarrowMaskForInterleavedOps(SelectionDAG &DAG, SDLoc &DL,
SDValue WideMask,
unsigned RequiredNumParts) {
@@ -27286,17 +27323,24 @@ static SDValue performInterleavedMaskedStoreCombine(
MaskedStoreSDNode *MST = cast<MaskedStoreSDNode>(N);
SDValue WideValue = MST->getValue();
- // Bail out if the stored value has an unexpected number of uses, since we'll
- // have to perform manual interleaving and may as well just use normal masked
- // stores. Also, discard masked stores that are truncating or indexed.
- if (!WideValue.hasOneUse() || !ISD::isNormalMaskedStore(MST) ||
- !MST->isSimple() || !MST->getOffset().isUndef())
+ // Discard masked stores that are truncating or indexed.
+ if (!ISD::isNormalMaskedStore(MST) || !MST->isSimple() ||
+ !MST->getOffset().isUndef())
return SDValue();
+ SDLoc DL(N);
SmallVector<SDValue, 4> ValueInterleaveOps;
- if (!isSequentialConcatOfVectorInterleave(WideValue.getNode(),
- ValueInterleaveOps))
+ if (isSequentialConcatOfVectorInterleave(WideValue.getNode(),
+ ValueInterleaveOps)) {
+ // Bail out if the stored value has an unexpected number of uses, since
+ // we'll have to perform manual interleaving and may as well just use normal
+ // masked stores.
+ if (!WideValue.hasOneUse())
+ return SDValue();
+ } else if (!isSplatVectorInterleaveOps(DAG, DL, WideValue,
+ ValueInterleaveOps)) {
return SDValue();
+ }
unsigned NumParts = ValueInterleaveOps.size();
if (NumParts != 2 && NumParts != 4)
@@ -27310,7 +27354,6 @@ static SDValue performInterleavedMaskedStoreCombine(
!DAG.getTargetLoweringInfo().isTypeLegal(SubVecTy))
return SDValue();
- SDLoc DL(N);
SDValue NarrowMask =
getNarrowMaskForInterleavedOps(DAG, DL, MST->getMask(), NumParts);
if (!NarrowMask)
diff --git a/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll b/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll
index bdc25bba711e1..87833306571c2 100644
--- a/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll
+++ b/llvm/test/CodeGen/AArch64/scalable_masked_interleaved_stores.ll
@@ -287,10 +287,8 @@ define void @foo_st2_nxv16i8_zeroinitializer(<vscale x 16 x i1> %mask, ptr %p) {
; CHECK-LABEL: foo_st2_nxv16i8_zeroinitializer:
; CHECK: // %bb.0:
; CHECK-NEXT: movi v0.2d, #0000000000000000
-; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
-; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
-; CHECK-NEXT: st1b { z0.b }, p1, [x0, #1, mul vl]
-; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: mov z1.d, z0.d
+; CHECK-NEXT: st2b { z0.b, z1.b }, p0, [x0]
; CHECK-NEXT: ret
%interleaved.mask = call <vscale x 32 x i1> @llvm.vector.interleave2.nxv32i1(<vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask)
call void @llvm.masked.store.nxv32i8.p0(<vscale x 32 x i8> zeroinitializer, ptr %p, i32 1, <vscale x 32 x i1> %interleaved.mask)
@@ -301,16 +299,10 @@ define void @foo_st4_nxv16i8_zeroinitializer(<vscale x 16 x i1> %mask, ptr %p) {
; CHECK-LABEL: foo_st4_nxv16i8_zeroinitializer:
; CHECK: // %bb.0:
; CHECK-NEXT: movi v0.2d, #0000000000000000
-; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
-; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
-; CHECK-NEXT: zip2 p2.b, p1.b, p1.b
-; CHECK-NEXT: zip1 p1.b, p1.b, p1.b
-; CHECK-NEXT: zip2 p3.b, p0.b, p0.b
-; CHECK-NEXT: st1b { z0.b }, p2, [x0, #3, mul vl]
-; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
-; CHECK-NEXT: st1b { z0.b }, p1, [x0, #2, mul vl]
-; CHECK-NEXT: st1b { z0.b }, p3, [x0, #1, mul vl]
-; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: mov z1.d, z0.d
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: mov z3.d, z0.d
+; CHECK-NEXT: st4b { z0.b - z3.b }, p0, [x0]
; CHECK-NEXT: ret
%interleaved.mask = call <vscale x 64 x i1> @llvm.vector.interleave4.nxv64i1(<vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask, <vscale x 16 x i1> %mask)
call void @llvm.masked.store.nxv64i8.p0(<vscale x 64 x i8> zeroinitializer, ptr %p, i32 1, <vscale x 64 x i1> %interleaved.mask)
@@ -321,10 +313,8 @@ define void @foo_st2_nxv16i8_splat(<vscale x 16 x i1> %mask, ptr %p) {
; CHECK-LABEL: foo_st2_nxv16i8_splat:
; CHECK: // %bb.0:
; CHECK-NEXT: mov z0.b, #1 // =0x1
-; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
-; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
-; CHECK-NEXT: st1b { z0.b }, p1, [x0, #1, mul vl]
-; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: mov z1.d, z0.d
+; CHECK-NEXT: st2b { z0.b, z1.b }, p0, [x0]
; CHECK-NEXT: ret
%base = insertelement <vscale x 32 x i8> poison, i8 1, i32 0
%interleaved.value = shufflevector <vscale x 32 x i8> %base, <vscale x 32 x i8> poison, <vscale x 32 x i32> zeroinitializer
@@ -336,17 +326,11 @@ define void @foo_st2_nxv16i8_splat(<vscale x 16 x i1> %mask, ptr %p) {
define void @foo_st4_nxv16i8_splat(<vscale x 16 x i1> %mask, ptr %p) {
; CHECK-LABEL: foo_st4_nxv16i8_splat:
; CHECK: // %bb.0:
-; CHECK-NEXT: zip2 p1.b, p0.b, p0.b
; CHECK-NEXT: mov z0.b, #1 // =0x1
-; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
-; CHECK-NEXT: zip2 p2.b, p1.b, p1.b
-; CHECK-NEXT: zip1 p1.b, p1.b, p1.b
-; CHECK-NEXT: zip2 p3.b, p0.b, p0.b
-; CHECK-NEXT: st1b { z0.b }, p2, [x0, #3, mul vl]
-; CHECK-NEXT: zip1 p0.b, p0.b, p0.b
-; CHECK-NEXT: st1b { z0.b }, p1, [x0, #2, mul vl]
-; CHECK-NEXT: st1b { z0.b }, p3, [x0, #1, mul vl]
-; CHECK-NEXT: st1b { z0.b }, p0, [x0]
+; CHECK-NEXT: mov z1.d, z0.d
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: mov z3.d, z0.d
+; CHECK-NEXT: st4b { z0.b - z3.b }, p0, [x0]
; CHECK-NEXT: ret
%base = insertelement <vscale x 64 x i8> poison, i8 1, i32 0
%interleaved.value = shufflevector <vscale x 64 x i8> %base, <vscale x 64 x i8> poison, <vscale x 64 x i32> zeroinitializer
More information about the llvm-commits
mailing list