[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:05:43 PDT 2026


https://github.com/Harry-Ramsey updated https://github.com/llvm/llvm-project/pull/207950

>From f4d46fbb3af30f874392f44996b0e05e0c38b89f 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 7ebe86930d68d470fa751a0050a67cf3e71ab830 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 f33d953eec747..151e6bd212f0d 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -27441,6 +27441,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) {
@@ -27479,17 +27516,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)
@@ -27503,7 +27547,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