[llvm] [X86] Fold prefix masked load/store copy (PR #207900)
Jaeuk Lee via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 6 22:50:08 PDT 2026
https://github.com/skku970412 created https://github.com/llvm/llvm-project/pull/207900
## Summary
Fold a conservative AVX-512 masked load/store copy pattern with a constant
true-prefix mask into an equivalent unmasked vector load/store over the active
prefix bytes.
This targets non-power-of-two vector cases such as:
- `<5 x i64>` with four active lanes
- `<9 x i32>` with eight active lanes
- `<17 x i8>` with sixteen active lanes
- `<33 x i8>` with thirty-two active lanes
Before this patch, these cases materialized an AVX-512 mask and emitted masked
memory operations. After the patch, they use the equivalent 128-bit or 256-bit
unmasked load/store when the inactive lanes are not observable.
Fixes #205032.
## Implementation
The transform is intentionally narrow and X86-specific. It runs from the X86
SelectionDAG masked-store combine when all of the following are true:
- AVX-512 is available.
- A `masked.load` feeds the matching `masked.store`.
- The loaded value and load chain each have only the store use.
- The load and store use the same mask SDValue.
- The mask is a constant true-prefix mask with power-of-two active lane count.
- The load passthru is `undef` in SelectionDAG.
- The load/store are simple, unindexed, non-volatile, non-atomic operations.
- The load is non-extending and non-expanding.
- The store is non-truncating and non-compressing.
- The active byte span is 128, 256, or 512 bits.
- The replacement load/store type is legal for X86.
The combine also handles masks after vector type legalization, where a boolean
mask may appear as `vXi1 = bitcast Constant<iX>`.
## Tests
Added `llvm/test/CodeGen/X86/avx512-non-pow2-masked-load-store.ll`.
Positive coverage:
- `v5i64` -> 256-bit unmasked load/store
- `v9i32` -> 256-bit unmasked load/store
- `v17i8` -> 128-bit unmasked load/store
- `v33i8` -> 256-bit unmasked load/store
Negative coverage:
- non-prefix mask
- prefix mask whose active lane count is not a power of two
- non-undef passthru
- load result used by two stores
## Local Verification
```bash
ninja -C build -j20 llc
build/bin/llvm-lit -q \
llvm/test/CodeGen/X86/avx512-non-pow2-masked-load-store.ll \
llvm/test/CodeGen/X86/masked_load.ll \
llvm/test/CodeGen/X86/masked_store.ll \
llvm/test/CodeGen/X86/masked_loadstore_split.ll \
llvm/test/CodeGen/X86/avx512-masked-memop-64-32.ll \
llvm/test/CodeGen/X86/avx512-masked_memop-16-8.ll
ninja -C build -j20 check-llvm-codegen-x86
```
`check-llvm-codegen-x86` result:
```text
Total Discovered Tests: 5533
Unsupported : 2
Passed : 5516
Expectedly Failed: 15
```
## Tool Usage
Assisted-by: OpenAI Codex
Codex was used for code navigation, testcase exploration, initial patch
drafting, and local test orchestration. I reviewed the implementation, tests,
legality constraints, and generated code before submission.
>From 52f49f58d14f99d4021d4487a89602f81245977c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EC=9D=B4=EC=9E=AC=EC=9A=B1?=
<126692701+skku970412 at users.noreply.github.com>
Date: Tue, 7 Jul 2026 14:49:34 +0900
Subject: [PATCH] [X86] Fold prefix masked load/store copy
---
llvm/lib/Target/X86/X86ISelLowering.cpp | 158 ++++++++++++++++++
.../X86/avx512-non-pow2-masked-load-store.ll | 120 +++++++++++++
2 files changed, 278 insertions(+)
create mode 100644 llvm/test/CodeGen/X86/avx512-non-pow2-masked-load-store.ll
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 9209134a055c6..dcce76551c451 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -54089,6 +54089,103 @@ static int getOneTrueElt(SDValue V) {
return TrueIndex;
}
+static std::optional<unsigned>
+getPow2PrefixMaskNumElts(ArrayRef<uint8_t> MaskBits) {
+ unsigned ActiveElts = 0;
+ while (ActiveElts != MaskBits.size() && MaskBits[ActiveElts])
+ ++ActiveElts;
+
+ for (unsigned I = ActiveElts, E = MaskBits.size(); I != E; ++I)
+ if (MaskBits[I])
+ return std::nullopt;
+
+ // The one-active-lane case is handled by the scalar masked load/store
+ // combines. This fold is for a narrower vector memory operation.
+ if (ActiveElts <= 1 || ActiveElts == MaskBits.size() ||
+ !isPowerOf2_32(ActiveElts))
+ return std::nullopt;
+
+ return ActiveElts;
+}
+
+/// Return the active lane count for a constant boolean mask with a true prefix
+/// of power-of-two length and false tail lanes. Return std::nullopt otherwise.
+static std::optional<unsigned> getPow2PrefixMaskNumElts(SDValue V) {
+ EVT VT = V.getValueType();
+ if (!VT.isFixedLengthVector() || VT.getVectorElementType() != MVT::i1)
+ return std::nullopt;
+
+ unsigned NumElts = VT.getVectorNumElements();
+ SmallVector<uint8_t, 64> MaskBits;
+
+ if (auto *BV = dyn_cast<BuildVectorSDNode>(V)) {
+ MaskBits.reserve(NumElts);
+ for (unsigned I = 0; I != NumElts; ++I) {
+ SDValue Op = BV->getOperand(I);
+ if (Op.isUndef())
+ return std::nullopt;
+ auto *C = dyn_cast<ConstantSDNode>(Op);
+ if (!C)
+ return std::nullopt;
+ MaskBits.push_back(!C->isZero());
+ }
+ return getPow2PrefixMaskNumElts(MaskBits);
+ }
+
+ // Legalization may turn a boolean vector mask into a bitcast from an integer
+ // constant, for example v8i1 bitcast(i8 15).
+ if (V.getOpcode() != ISD::BITCAST)
+ return std::nullopt;
+
+ auto *C = dyn_cast<ConstantSDNode>(V.getOperand(0));
+ if (!C)
+ return std::nullopt;
+
+ const APInt &RawMask = C->getAPIntValue();
+ if (RawMask.getBitWidth() < NumElts)
+ return std::nullopt;
+
+ MaskBits.reserve(NumElts);
+ for (unsigned I = 0; I != NumElts; ++I)
+ MaskBits.push_back(RawMask[I]);
+
+ return getPow2PrefixMaskNumElts(MaskBits);
+}
+
+static std::optional<MVT>
+getLegalMaskedMemCopyVT(EVT EltVT, unsigned ActiveElts,
+ const TargetLowering &TLI) {
+ if (!EltVT.isSimple())
+ return std::nullopt;
+
+ MVT EltMVT = EltVT.getSimpleVT();
+ if (EltMVT == MVT::i1 ||
+ EltMVT.getStoreSizeInBits() != EltMVT.getSizeInBits())
+ return std::nullopt;
+
+ unsigned ActiveBits = ActiveElts * EltMVT.getStoreSizeInBits();
+ if (ActiveBits != 128 && ActiveBits != 256 && ActiveBits != 512)
+ return std::nullopt;
+
+ MVT SameEltVT = MVT::getVectorVT(EltMVT, ActiveElts);
+ if (TLI.isTypeLegal(SameEltVT))
+ return SameEltVT;
+
+ // For a masked-load-to-masked-store copy, the loaded vector value is not
+ // otherwise observed. If the same element type is not legal, use any legal
+ // integer vector type with the same byte width.
+ for (MVT IntVT : {MVT::i32, MVT::i64, MVT::i8, MVT::i16}) {
+ unsigned EltBits = IntVT.getSizeInBits();
+ if (ActiveBits % EltBits != 0)
+ continue;
+ MVT VecVT = MVT::getVectorVT(IntVT, ActiveBits / EltBits);
+ if (TLI.isTypeLegal(VecVT))
+ return VecVT;
+ }
+
+ return std::nullopt;
+}
+
/// Given a masked memory load/store operation, return true if it has one mask
/// bit set. If it has one mask bit set, then also return the memory address of
/// the scalar element to load/store, the vector index to insert/extract that
@@ -54309,6 +54406,64 @@ static SDValue reduceMaskedStoreToScalarStore(MaskedStoreSDNode *MS,
Alignment, MS->getMemOperand()->getFlags());
}
+static SDValue combinePrefixMaskedLoadStore(MaskedStoreSDNode *MS,
+ SelectionDAG &DAG,
+ const X86Subtarget &Subtarget) {
+ if (!Subtarget.hasAVX512() || !MS->isUnindexed() || !MS->isSimple() ||
+ MS->isTruncatingStore() || MS->isCompressingStore() ||
+ MS->isNonTemporal())
+ return SDValue();
+
+ SDValue StoredVal = MS->getValue();
+ if (StoredVal.getOpcode() != ISD::MLOAD || !StoredVal.hasOneUse())
+ return SDValue();
+
+ auto *ML = cast<MaskedLoadSDNode>(StoredVal);
+ if (!ML->isUnindexed() || !ML->isSimple() || ML->isExpandingLoad() ||
+ ML->getExtensionType() != ISD::NON_EXTLOAD || ML->isNonTemporal() ||
+ !ML->getPassThru().isUndef())
+ return SDValue();
+
+ // Keep the first patch intentionally narrow: the store must be chained
+ // directly after the load, and the load result must only feed this store.
+ SDValue LoadChain(ML, 1);
+ if (MS->getChain() != LoadChain || !LoadChain.hasOneUse())
+ return SDValue();
+
+ SDValue Mask = MS->getMask();
+ if (ML->getMask() != Mask)
+ return SDValue();
+
+ std::optional<unsigned> ActiveElts = getPow2PrefixMaskNumElts(Mask);
+ if (!ActiveElts)
+ return SDValue();
+
+ EVT MemVT = ML->getMemoryVT();
+ EVT VT = StoredVal.getValueType();
+ if (!VT.isFixedLengthVector() || !MemVT.isFixedLengthVector() ||
+ MS->getMemoryVT() != MemVT ||
+ MemVT.getVectorElementType() != VT.getVectorElementType() ||
+ MemVT.getVectorNumElements() < *ActiveElts)
+ return SDValue();
+
+ const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+ std::optional<MVT> NarrowVT =
+ getLegalMaskedMemCopyVT(MemVT.getVectorElementType(), *ActiveElts, TLI);
+ if (!NarrowVT)
+ return SDValue();
+
+ SDLoc DL(MS);
+ SDValue NarrowLoad =
+ DAG.getLoad(*NarrowVT, DL, ML->getChain(), ML->getBasePtr(),
+ ML->getPointerInfo(), ML->getBaseAlign(),
+ ML->getMemOperand()->getFlags(), ML->getAAInfo(),
+ ML->getRanges());
+
+ return DAG.getStore(NarrowLoad.getValue(1), DL, NarrowLoad, MS->getBasePtr(),
+ MS->getPointerInfo(), MS->getBaseAlign(),
+ MS->getMemOperand()->getFlags(), MS->getAAInfo());
+}
+
static SDValue combineMaskedStore(SDNode *N, SelectionDAG &DAG,
TargetLowering::DAGCombinerInfo &DCI,
const X86Subtarget &Subtarget) {
@@ -54316,6 +54471,9 @@ static SDValue combineMaskedStore(SDNode *N, SelectionDAG &DAG,
if (Mst->isCompressingStore())
return SDValue();
+ if (SDValue NarrowStore = combinePrefixMaskedLoadStore(Mst, DAG, Subtarget))
+ return NarrowStore;
+
if (SDValue ScalarStore = reduceMaskedStoreToScalarStore(Mst, DAG, Subtarget))
return ScalarStore;
diff --git a/llvm/test/CodeGen/X86/avx512-non-pow2-masked-load-store.ll b/llvm/test/CodeGen/X86/avx512-non-pow2-masked-load-store.ll
new file mode 100644
index 0000000000000..5bb1a5ad42a9d
--- /dev/null
+++ b/llvm/test/CodeGen/X86/avx512-non-pow2-masked-load-store.ll
@@ -0,0 +1,120 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -mattr=+avx512f,+avx512bw,+avx512vl | FileCheck %s
+
+; These cases should avoid AVX-512 mask materialization and masked memory
+; operations for a constant true-prefix mask. They are equivalent to an
+; unmasked power-of-two-width load/store over the active prefix:
+; v5i64 -> v4i64, v9i32 -> v8i32, v17i8 -> v16i8, v33i8 -> v32i8.
+
+define void @load_store_masked_v5i64(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_masked_v5i64:
+; CHECK: # %bb.0:
+; CHECK-NEXT: vmovups (%rdi), %ymm0
+; CHECK-NEXT: vmovups %ymm0, (%rsi)
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <5 x i64> @llvm.masked.load.v5i64.p0(ptr %src, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>, <5 x i64> poison)
+ call void @llvm.masked.store.v5i64.p0(<5 x i64> %val, ptr %dst, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+define void @load_store_masked_v9i32(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_masked_v9i32:
+; CHECK: # %bb.0:
+; CHECK-NEXT: vmovups (%rdi), %ymm0
+; CHECK-NEXT: vmovups %ymm0, (%rsi)
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <9 x i32> @llvm.masked.load.v9i32.p0(ptr %src, i32 1, <9 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false>, <9 x i32> poison)
+ call void @llvm.masked.store.v9i32.p0(<9 x i32> %val, ptr %dst, i32 1, <9 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+define void @load_store_masked_v17i8(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_masked_v17i8:
+; CHECK: # %bb.0:
+; CHECK-NEXT: vmovups (%rdi), %xmm0
+; CHECK-NEXT: vmovups %xmm0, (%rsi)
+; CHECK-NEXT: retq
+ %val = call <17 x i8> @llvm.masked.load.v17i8.p0(ptr %src, i32 1, <17 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false>, <17 x i8> poison)
+ call void @llvm.masked.store.v17i8.p0(<17 x i8> %val, ptr %dst, i32 1, <17 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+define void @load_store_masked_v33i8(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_masked_v33i8:
+; CHECK: # %bb.0:
+; CHECK-NEXT: vmovups (%rdi), %ymm0
+; CHECK-NEXT: vmovups %ymm0, (%rsi)
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <33 x i8> @llvm.masked.load.v33i8.p0(ptr %src, i32 1, <33 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false>, <33 x i8> poison)
+ call void @llvm.masked.store.v33i8.p0(<33 x i8> %val, ptr %dst, i32 1, <33 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+define void @load_store_non_prefix_mask_v5i64(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_non_prefix_mask_v5i64:
+; CHECK: # %bb.0:
+; CHECK-NEXT: movb $13, %al
+; CHECK-NEXT: kmovd %eax, %k1
+; CHECK-NEXT: vmovdqu64 (%rdi), %zmm0 {%k1} {z}
+; CHECK-NEXT: vmovdqu64 %zmm0, (%rsi) {%k1}
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <5 x i64> @llvm.masked.load.v5i64.p0(ptr %src, i32 1, <5 x i1> <i1 true, i1 false, i1 true, i1 true, i1 false>, <5 x i64> poison)
+ call void @llvm.masked.store.v5i64.p0(<5 x i64> %val, ptr %dst, i32 1, <5 x i1> <i1 true, i1 false, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+define void @load_store_prefix3_mask_v5i64(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_prefix3_mask_v5i64:
+; CHECK: # %bb.0:
+; CHECK-NEXT: movb $7, %al
+; CHECK-NEXT: kmovd %eax, %k1
+; CHECK-NEXT: vmovdqu64 (%rdi), %zmm0 {%k1} {z}
+; CHECK-NEXT: vmovdqu64 %zmm0, (%rsi) {%k1}
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <5 x i64> @llvm.masked.load.v5i64.p0(ptr %src, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 false, i1 false>, <5 x i64> poison)
+ call void @llvm.masked.store.v5i64.p0(<5 x i64> %val, ptr %dst, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 false, i1 false>)
+ ret void
+}
+
+define void @load_store_passthru_zero_v5i64(ptr %src, ptr %dst) {
+; CHECK-LABEL: load_store_passthru_zero_v5i64:
+; CHECK: # %bb.0:
+; CHECK-NEXT: movb $15, %al
+; CHECK-NEXT: kmovd %eax, %k1
+; CHECK-NEXT: vmovdqu64 (%rdi), %zmm0 {%k1} {z}
+; CHECK-NEXT: vmovdqu64 %zmm0, (%rsi) {%k1}
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <5 x i64> @llvm.masked.load.v5i64.p0(ptr %src, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>, <5 x i64> zeroinitializer)
+ call void @llvm.masked.store.v5i64.p0(<5 x i64> %val, ptr %dst, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+define void @load_store_two_uses_v5i64(ptr %src, ptr %dst0, ptr %dst1) {
+; CHECK-LABEL: load_store_two_uses_v5i64:
+; CHECK: # %bb.0:
+; CHECK-NEXT: movb $15, %al
+; CHECK-NEXT: kmovd %eax, %k1
+; CHECK-NEXT: vmovdqu64 (%rdi), %zmm0 {%k1} {z}
+; CHECK-NEXT: vmovdqu64 %zmm0, (%rsi) {%k1}
+; CHECK-NEXT: vmovdqu64 %zmm0, (%rdx) {%k1}
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %val = call <5 x i64> @llvm.masked.load.v5i64.p0(ptr %src, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>, <5 x i64> poison)
+ call void @llvm.masked.store.v5i64.p0(<5 x i64> %val, ptr %dst0, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>)
+ call void @llvm.masked.store.v5i64.p0(<5 x i64> %val, ptr %dst1, i32 1, <5 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false>)
+ ret void
+}
+
+declare <5 x i64> @llvm.masked.load.v5i64.p0(ptr, i32 immarg, <5 x i1>, <5 x i64>)
+declare void @llvm.masked.store.v5i64.p0(<5 x i64>, ptr, i32 immarg, <5 x i1>)
+declare <9 x i32> @llvm.masked.load.v9i32.p0(ptr, i32 immarg, <9 x i1>, <9 x i32>)
+declare void @llvm.masked.store.v9i32.p0(<9 x i32>, ptr, i32 immarg, <9 x i1>)
+declare <17 x i8> @llvm.masked.load.v17i8.p0(ptr, i32 immarg, <17 x i1>, <17 x i8>)
+declare void @llvm.masked.store.v17i8.p0(<17 x i8>, ptr, i32 immarg, <17 x i1>)
+declare <33 x i8> @llvm.masked.load.v33i8.p0(ptr, i32 immarg, <33 x i1>, <33 x i8>)
+declare void @llvm.masked.store.v33i8.p0(<33 x i8>, ptr, i32 immarg, <33 x i1>)
More information about the llvm-commits
mailing list