[llvm] [X86] Keep bitcasted mask zero-tests in k-registers (PR #207908)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 6 23:57:17 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-x86
Author: Jaeuk Lee (skku970412)
<details>
<summary>Changes</summary>
# [X86] Keep bitcasted mask zero-tests in k-registers
## Summary
This improves an AVX-512 mask lowering case where a scalar zero-test around a
bitcasted `<N x i1>` value forced the compare result out of a k-register and
into a GPR.
The handled pattern is:
```llvm
%mask.bits = bitcast <N x i1> %cmp to iN
%and = and iN %mask.bits, %scalar.mask
%eq = icmp eq iN %and, 0
```
Before this change, the selected code copied the compare result to a GPR and
used scalar `test`. For a representative `v16i1` case:
```asm
vpcmpneqb %xmm1, %xmm0, %k0
kmovd %k0, %ecx
testw %di, %cx
```
After this change, X86 lowering rebuilds the scalar mask as a k-register mask,
creates a mask-domain `and`, and tests it with `kortest`. This lets instruction
selection fold the `and` into a masked compare when possible:
```asm
kmovd %edi, %k1
vpcmpneqb %xmm1, %xmm0, %k0 {%k1}
kortestw %k0, %k0
```
## Scope
The transform is intentionally narrow:
- only X86 AVX-512 lowering
- only fixed legal mask types `v8i1`, `v16i1`, `v32i1`, and `v64i1`
- only `icmp eq/ne ..., 0`
- only non-constant scalar masks
Constant one-use scalar masks such as `127` are left on the existing scalar
immediate-test path because materializing such constants in k-registers is not
clearly better without a separate constant-mask materialization improvement.
Relates to #<!-- -->206722.
## Tests
Added `llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll`.
Coverage includes:
- `v8i1`, `v16i1`, and `v32i1` variable scalar masks
- `eq` and `ne` zero-tests
- AVX512F-only `v16i1` fallback
- a constant-mask negative case that remains scalar
Commands run:
```bash
ninja -C ../build -j20 llc
../build/bin/llvm-lit -sv llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll
../build/bin/llvm-lit -sv \
llvm/test/CodeGen/X86/avx512-mask-op.ll \
llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll
../build/bin/llvm-lit -sv llvm/test/CodeGen/X86
git diff --check
```
Full X86 CodeGen 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 test execution. I reviewed the implementation, tests, generated
assembly, and legality constraints before submission.
---
Full diff: https://github.com/llvm/llvm-project/pull/207908.diff
2 Files Affected:
- (modified) llvm/lib/Target/X86/X86ISelLowering.cpp (+49-6)
- (added) llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll (+117)
``````````diff
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 9209134a055c6..b5d7b1bbb89d0 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -24980,15 +24980,60 @@ static SDValue EmitAVX512Test(SDValue Op0, SDValue Op1, ISD::CondCode CC,
SDValue &X86CC) {
assert((CC == ISD::SETEQ || CC == ISD::SETNE) && "Unsupported ISD::CondCode");
+ auto IsLegalTestVT = [&](MVT VT) {
+ return (Subtarget.hasAVX512() && VT == MVT::v16i1) ||
+ (Subtarget.hasDQI() && VT == MVT::v8i1) ||
+ (Subtarget.hasBWI() && (VT == MVT::v32i1 || VT == MVT::v64i1));
+ };
+
+ auto IsKTestableVT = [&](MVT VT) {
+ return (Subtarget.hasDQI() && (VT == MVT::v8i1 || VT == MVT::v16i1)) ||
+ (Subtarget.hasBWI() && (VT == MVT::v32i1 || VT == MVT::v64i1));
+ };
+
+ // Match scalar tests of a bitcasted mask:
+ // (and (bitcast vXi1 X), Y) == 0
+ // as KTEST/KORTEST so the vXi1 mask does not have to be copied to a GPR.
+ if (isNullConstant(Op1) && Op0.getOpcode() == ISD::AND && Op0.hasOneUse()) {
+ auto MatchScalarAnd = [&](SDValue Mask, SDValue ScalarMask) -> SDValue {
+ if (Mask.getOpcode() != ISD::BITCAST || !Mask.hasOneUse() ||
+ !ScalarMask.getValueType().isScalarInteger())
+ return SDValue();
+
+ // Keep one-use immediate masks on the scalar TEST path unless we know
+ // materializing them in a k-register is profitable.
+ if (isa<ConstantSDNode>(ScalarMask))
+ return SDValue();
+
+ SDValue MaskOp = Mask.getOperand(0);
+ MVT VT = MaskOp.getSimpleValueType();
+ if (!VT.isVectorOf(MVT::i1) || !IsLegalTestVT(VT))
+ return SDValue();
+
+ EVT IntVT = EVT::getIntegerVT(*DAG.getContext(),
+ VT.getVectorNumElements());
+ ScalarMask = DAG.getZExtOrTrunc(ScalarMask, dl, IntVT);
+ ScalarMask = DAG.getBitcast(VT, ScalarMask);
+
+ X86::CondCode X86Cond = CC == ISD::SETEQ ? X86::COND_E : X86::COND_NE;
+ X86CC = DAG.getTargetConstant(X86Cond, dl, MVT::i8);
+ SDValue And = DAG.getNode(ISD::AND, dl, VT, MaskOp, ScalarMask);
+ return DAG.getNode(X86ISD::KORTEST, dl, MVT::i32, And, And);
+ };
+
+ if (SDValue Test = MatchScalarAnd(Op0.getOperand(0), Op0.getOperand(1)))
+ return Test;
+ if (SDValue Test = MatchScalarAnd(Op0.getOperand(1), Op0.getOperand(0)))
+ return Test;
+ }
+
// Must be a bitcast from vXi1.
if (Op0.getOpcode() != ISD::BITCAST)
return SDValue();
Op0 = Op0.getOperand(0);
MVT VT = Op0.getSimpleValueType();
- if (!(Subtarget.hasAVX512() && VT == MVT::v16i1) &&
- !(Subtarget.hasDQI() && VT == MVT::v8i1) &&
- !(Subtarget.hasBWI() && (VT == MVT::v32i1 || VT == MVT::v64i1)))
+ if (!IsLegalTestVT(VT))
return SDValue();
X86::CondCode X86Cond;
@@ -25002,9 +25047,7 @@ static SDValue EmitAVX512Test(SDValue Op0, SDValue Op1, ISD::CondCode CC,
// If the input is an AND, we can combine it's operands into the KTEST.
bool KTestable = false;
- if (Subtarget.hasDQI() && (VT == MVT::v8i1 || VT == MVT::v16i1))
- KTestable = true;
- if (Subtarget.hasBWI() && (VT == MVT::v32i1 || VT == MVT::v64i1))
+ if (IsKTestableVT(VT))
KTestable = true;
if (!isNullConstant(Op1))
KTestable = false;
diff --git a/llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll b/llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll
new file mode 100644
index 0000000000000..9cc90b97f20f7
--- /dev/null
+++ b/llvm/test/CodeGen/X86/avx512-mask-bitcast-test.ll
@@ -0,0 +1,117 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512bw,+avx512vl,+avx512dq | FileCheck %s
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512f | FileCheck %s --check-prefix=AVX512F
+
+define i32 @test_v8i1_scalar_mask(<8 x i32> %a, <8 x i32> %b, i8 %mask) {
+; CHECK-LABEL: test_v8i1_scalar_mask:
+; CHECK: # %bb.0:
+; CHECK-NEXT: kmovd %edi, %k1
+; CHECK-NEXT: vpcmpneqd %ymm1, %ymm0, %k0 {%k1}
+; CHECK-NEXT: xorl %eax, %eax
+; CHECK-NEXT: kortestb %k0, %k0
+; CHECK-NEXT: sete %al
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %cmp = icmp ne <8 x i32> %a, %b
+ %bc = bitcast <8 x i1> %cmp to i8
+ %and = and i8 %bc, %mask
+ %eq = icmp eq i8 %and, 0
+ %ret = zext i1 %eq to i32
+ ret i32 %ret
+}
+
+define i32 @test_v16i1_scalar_mask(<2 x i64> %a, <2 x i64> %b, i16 %mask) {
+; CHECK-LABEL: test_v16i1_scalar_mask:
+; CHECK: # %bb.0:
+; CHECK-NEXT: kmovd %edi, %k1
+; CHECK-NEXT: vpcmpneqb %xmm1, %xmm0, %k0 {%k1}
+; CHECK-NEXT: xorl %eax, %eax
+; CHECK-NEXT: kortestw %k0, %k0
+; CHECK-NEXT: sete %al
+; CHECK-NEXT: retq
+ %va = bitcast <2 x i64> %a to <16 x i8>
+ %vb = bitcast <2 x i64> %b to <16 x i8>
+ %cmp = icmp ne <16 x i8> %va, %vb
+ %bc = bitcast <16 x i1> %cmp to i16
+ %and = and i16 %bc, %mask
+ %eq = icmp eq i16 %and, 0
+ %ret = zext i1 %eq to i32
+ ret i32 %ret
+}
+
+define i32 @test_v16i1_scalar_mask_ne(<2 x i64> %a, <2 x i64> %b,
+ i16 %mask) {
+; CHECK-LABEL: test_v16i1_scalar_mask_ne:
+; CHECK: # %bb.0:
+; CHECK-NEXT: kmovd %edi, %k1
+; CHECK-NEXT: vpcmpneqb %xmm1, %xmm0, %k0 {%k1}
+; CHECK-NEXT: xorl %eax, %eax
+; CHECK-NEXT: kortestw %k0, %k0
+; CHECK-NEXT: setne %al
+; CHECK-NEXT: retq
+ %va = bitcast <2 x i64> %a to <16 x i8>
+ %vb = bitcast <2 x i64> %b to <16 x i8>
+ %cmp = icmp ne <16 x i8> %va, %vb
+ %bc = bitcast <16 x i1> %cmp to i16
+ %and = and i16 %bc, %mask
+ %ne = icmp ne i16 %and, 0
+ %ret = zext i1 %ne to i32
+ ret i32 %ret
+}
+
+define i32 @test_v32i1_scalar_mask(<32 x i16> %a, <32 x i16> %b, i32 %mask) {
+; CHECK-LABEL: test_v32i1_scalar_mask:
+; CHECK: # %bb.0:
+; CHECK-NEXT: kmovd %edi, %k1
+; CHECK-NEXT: vpcmpneqw %zmm1, %zmm0, %k0 {%k1}
+; CHECK-NEXT: xorl %eax, %eax
+; CHECK-NEXT: kortestd %k0, %k0
+; CHECK-NEXT: sete %al
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+ %cmp = icmp ne <32 x i16> %a, %b
+ %bc = bitcast <32 x i1> %cmp to i32
+ %and = and i32 %bc, %mask
+ %eq = icmp eq i32 %and, 0
+ %ret = zext i1 %eq to i32
+ ret i32 %ret
+}
+
+; Keep constant masks on the scalar immediate-test path for now. Materializing
+; a one-use immediate mask in a k-register is not clearly better.
+define i32 @test_v16i1_constant_mask(<2 x i64> %a, <2 x i64> %b) {
+; CHECK-LABEL: test_v16i1_constant_mask:
+; CHECK: # %bb.0:
+; CHECK-NEXT: vpcmpneqb %xmm1, %xmm0, %k0
+; CHECK-NEXT: kmovd %k0, %ecx
+; CHECK-NEXT: xorl %eax, %eax
+; CHECK-NEXT: testb $127, %cl
+; CHECK-NEXT: sete %al
+; CHECK-NEXT: retq
+ %va = bitcast <2 x i64> %a to <16 x i8>
+ %vb = bitcast <2 x i64> %b to <16 x i8>
+ %cmp = icmp ne <16 x i8> %va, %vb
+ %bc = bitcast <16 x i1> %cmp to i16
+ %and = and i16 %bc, 127
+ %eq = icmp eq i16 %and, 0
+ %ret = zext i1 %eq to i32
+ ret i32 %ret
+}
+
+define i32 @test_v16i1_scalar_mask_avx512f(<16 x i32> %a, <16 x i32> %b,
+ i16 %mask) {
+; AVX512F-LABEL: test_v16i1_scalar_mask_avx512f:
+; AVX512F: # %bb.0:
+; AVX512F-NEXT: kmovw %edi, %k1
+; AVX512F-NEXT: vpcmpneqd %zmm1, %zmm0, %k0 {%k1}
+; AVX512F-NEXT: xorl %eax, %eax
+; AVX512F-NEXT: kortestw %k0, %k0
+; AVX512F-NEXT: sete %al
+; AVX512F-NEXT: vzeroupper
+; AVX512F-NEXT: retq
+ %cmp = icmp ne <16 x i32> %a, %b
+ %bc = bitcast <16 x i1> %cmp to i16
+ %and = and i16 %bc, %mask
+ %eq = icmp eq i16 %and, 0
+ %ret = zext i1 %eq to i32
+ ret i32 %ret
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/207908
More information about the llvm-commits
mailing list