[llvm] [SelectionDAG] Preserve CTTZ_ELTS lane count during expansion (PR #217982)

Oscar Priego via llvm-commits llvm-commits at lists.llvm.org
Sat Aug 22 14:48:44 PDT 2026


https://github.com/Opriego updated https://github.com/llvm/llvm-project/pull/217982

>From 1acab00142e4ae69d810344140bed69fec34fa70 Mon Sep 17 00:00:00 2001
From: Oscar Priego Verdugo <oscar.priegov at gmail.com>
Date: Fri, 21 Aug 2026 10:57:35 -0600
Subject: [PATCH] [SelectionDAG] Preserve CTTZ_ELTS lane count during expansion

CTTZ_ELTS expansion derived VL from its legalized auxiliary step vector. On X86 with AVX512F, the step vector for a semantic v4i1 mask is widened from v4i8 to v16i8, causing an all-zero mask to return 16 instead of 4.

Derive VL from the CTTZ_ELTS operand's ElementCount instead. Auxiliary type legalization can still widen the step vector without changing the logical lane domain of the operation.
---
 .../CodeGen/SelectionDAG/TargetLowering.cpp   | 10 +-
 llvm/test/CodeGen/X86/intrinsic-cttz-elts.ll  | 95 +++++++++++++++++++
 2 files changed, 100 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index cca7cdad0e2c8..459da7c9bd33b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -13859,10 +13859,11 @@ SDValue TargetLowering::expandVECTOR_COMPRESS(SDNode *Node,
 SDValue TargetLowering::expandCttzElts(SDNode *Node, SelectionDAG &DAG) const {
   SDLoc DL(Node);
   EVT VT = Node->getValueType(0);
+  SDValue Op = Node->getOperand(0);
+  ElementCount EC = Op.getValueType().getVectorElementCount();
 
   bool ZeroIsPoison = Node->getOpcode() == ISD::CTTZ_ELTS_ZERO_POISON;
-  auto [Mask, StepVec] =
-      getLegalMaskAndStepVector(Node->getOperand(0), ZeroIsPoison, DL, DAG);
+  auto [Mask, StepVec] = getLegalMaskAndStepVector(Op, ZeroIsPoison, DL, DAG);
 
   // No legal step vector: split mask in half and recombine results.
   // LoNumElts uses the non-poison CTTZ_ELTS so its result is well-defined
@@ -13870,7 +13871,7 @@ SDValue TargetLowering::expandCttzElts(SDNode *Node, SelectionDAG &DAG) const {
   // Result: (ResLo != LoNumElts) ? ResLo : (LoNumElts + ResHi)
   if (!StepVec) {
     EVT ResVT = Node->getValueType(0);
-    auto [MaskLo, MaskHi] = DAG.SplitVector(Node->getOperand(0), DL);
+    auto [MaskLo, MaskHi] = DAG.SplitVector(Op, DL);
     SDValue LoNumElts = DAG.getElementCount(
         DL, ResVT, MaskLo.getValueType().getVectorElementCount());
     SDValue ResLo = DAG.getNode(ISD::CTTZ_ELTS, DL, ResVT, MaskLo);
@@ -13893,8 +13894,7 @@ SDValue TargetLowering::expandCttzElts(SDNode *Node, SelectionDAG &DAG) const {
   if (getTypeAction(StepVT.getSimpleVT()) == TypePromoteInteger)
     StepVT = getTypeToTransformTo(*DAG.getContext(), StepVT);
 
-  SDValue VL =
-      DAG.getElementCount(DL, StepVT, StepVecVT.getVectorElementCount());
+  SDValue VL = DAG.getElementCount(DL, StepVT, EC);
   SDValue SplatVL = DAG.getSplat(StepVecVT, DL, VL);
   StepVec = DAG.getNode(ISD::SUB, DL, StepVecVT, SplatVL, StepVec);
   SDValue Zeroes = DAG.getConstant(0, DL, StepVecVT);
diff --git a/llvm/test/CodeGen/X86/intrinsic-cttz-elts.ll b/llvm/test/CodeGen/X86/intrinsic-cttz-elts.ll
index 65231c484db98..f8cc45c1b5bda 100644
--- a/llvm/test/CodeGen/X86/intrinsic-cttz-elts.ll
+++ b/llvm/test/CodeGen/X86/intrinsic-cttz-elts.ll
@@ -1,4 +1,5 @@
 ; RUN: llc -mtriple=x86_64-unknown-unknown < %s | FileCheck %s
+; RUN: llc -mtriple=x86_64-unknown-unknown -mattr=+avx512f < %s | FileCheck %s --check-prefix=AVX512
 
 define i8 @ctz_v8i16(<8 x i16> %a) {
 ; CHECK-LABEL: .LCPI0_0:
@@ -29,6 +30,26 @@ define i8 @ctz_v8i16(<8 x i16> %a) {
 ; CHECK-NEXT:    subl %ecx, %eax
 ; CHECK-NEXT:    # kill: def $al killed $al killed $eax
 ; CHECK-NEXT:    retq
+;
+; AVX512-LABEL: ctz_v8i16:
+; AVX512:       # %bb.0:
+; AVX512-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512-NEXT:    vpcmpeqw %xmm1, %xmm0, %xmm0
+; AVX512-NEXT:    vpternlogq {{.*#+}} zmm0 = ~zmm0
+; AVX512-NEXT:    vpmovsxwq %xmm0, %zmm0
+; AVX512-NEXT:    vptestmq %zmm0, %zmm0, %k1
+; AVX512-NEXT:    vpternlogd {{.*#+}} zmm0 {%k1} {z} = -1
+; AVX512-NEXT:    vpmovdb %zmm0, %xmm0
+; AVX512-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; AVX512-NEXT:    vpternlogq {{.*#+}} zmm0 = ~zmm0
+; AVX512-NEXT:    vpsrlw $8, %xmm0, %xmm1
+; AVX512-NEXT:    vpminub %xmm1, %xmm0, %xmm0
+; AVX512-NEXT:    vphminposuw %xmm0, %xmm0
+; AVX512-NEXT:    vmovd %xmm0, %eax
+; AVX512-NEXT:    addb $9, %al
+; AVX512-NEXT:    # kill: def $al killed $al killed $eax
+; AVX512-NEXT:    vzeroupper
+; AVX512-NEXT:    retq
   %res = call i8 @llvm.experimental.cttz.elts.i8.v8i16(<8 x i16> %a, i1 0)
   ret i8 %res
 }
@@ -62,6 +83,26 @@ define i16 @ctz_v4i32(<4 x i32> %a) {
 ; CHECK-NEXT:    subl %ecx, %eax
 ; CHECK-NEXT:    # kill: def $ax killed $ax killed $eax
 ; CHECK-NEXT:    retq
+;
+; AVX512-LABEL: ctz_v4i32:
+; AVX512:       # %bb.0:
+; AVX512-NEXT:    # kill: def $xmm0 killed $xmm0 def $zmm0
+; AVX512-NEXT:    vptestmd %zmm0, %zmm0, %k0
+; AVX512-NEXT:    kshiftlw $12, %k0, %k0
+; AVX512-NEXT:    kshiftrw $12, %k0, %k1
+; AVX512-NEXT:    vpternlogd {{.*#+}} zmm0 {%k1} {z} = -1
+; AVX512-NEXT:    vpmovdb %zmm0, %xmm0
+; AVX512-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; AVX512-NEXT:    vpternlogq {{.*#+}} zmm0 = ~zmm0
+; AVX512-NEXT:    vpsrlw $8, %xmm0, %xmm1
+; AVX512-NEXT:    vpminub %xmm1, %xmm0, %xmm0
+; AVX512-NEXT:    vphminposuw %xmm0, %xmm0
+; AVX512-NEXT:    vmovd %xmm0, %eax
+; AVX512-NEXT:    addb $5, %al
+; AVX512-NEXT:    movzbl %al, %eax
+; AVX512-NEXT:    # kill: def $ax killed $ax killed $eax
+; AVX512-NEXT:    vzeroupper
+; AVX512-NEXT:    retq
   %res = call i16 @llvm.experimental.cttz.elts.i16.v4i32(<4 x i32> %a, i1 0)
   ret i16 %res
 }
@@ -97,9 +138,63 @@ define i8 @ctz_v8i16_poison(<8 x i16> %a) {
 ; CHECK-NEXT:    subl %ecx, %eax
 ; CHECK-NEXT:    # kill: def $al killed $al killed $eax
 ; CHECK-NEXT:    retq
+;
+; AVX512-LABEL: ctz_v8i16_poison:
+; AVX512:       # %bb.0:
+; AVX512-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512-NEXT:    vpcmpeqw %xmm1, %xmm0, %xmm0
+; AVX512-NEXT:    vpternlogq {{.*#+}} zmm0 = ~zmm0
+; AVX512-NEXT:    vpmovsxwq %xmm0, %zmm0
+; AVX512-NEXT:    vptestmq %zmm0, %zmm0, %k1
+; AVX512-NEXT:    vpternlogd {{.*#+}} zmm0 {%k1} {z} = -1
+; AVX512-NEXT:    vpmovdb %zmm0, %xmm0
+; AVX512-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; AVX512-NEXT:    vpternlogq {{.*#+}} zmm0 = ~zmm0
+; AVX512-NEXT:    vpsrlw $8, %xmm0, %xmm1
+; AVX512-NEXT:    vpminub %xmm1, %xmm0, %xmm0
+; AVX512-NEXT:    vphminposuw %xmm0, %xmm0
+; AVX512-NEXT:    vmovd %xmm0, %eax
+; AVX512-NEXT:    addb $9, %al
+; AVX512-NEXT:    # kill: def $al killed $al killed $eax
+; AVX512-NEXT:    vzeroupper
+; AVX512-NEXT:    retq
   %res = call i8 @llvm.experimental.cttz.elts.i8.v8i16(<8 x i16> %a, i1 1)
   ret i8 %res
 }
 
+define i32 @ctz_zero_v4i1() {
+; CHECK-LABEL: ctz_zero_v4i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    movl $4, %eax
+; CHECK-NEXT:    retq
+;
+; AVX512-LABEL: ctz_zero_v4i1:
+; AVX512:       # %bb.0:
+; AVX512-NEXT:    vphminposuw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0
+; AVX512-NEXT:    vmovd %xmm0, %eax
+; AVX512-NEXT:    addb $5, %al
+; AVX512-NEXT:    movzbl %al, %eax
+; AVX512-NEXT:    retq
+  %res = call i32 @llvm.experimental.cttz.elts.i32.v4i1(<4 x i1> zeroinitializer, i1 false)
+  ret i32 %res
+}
+
+define i32 @ctz_zero_v8i1() {
+; CHECK-LABEL: ctz_zero_v8i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    movl $8, %eax
+; CHECK-NEXT:    retq
+;
+; AVX512-LABEL: ctz_zero_v8i1:
+; AVX512:       # %bb.0:
+; AVX512-NEXT:    vphminposuw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0
+; AVX512-NEXT:    vmovd %xmm0, %eax
+; AVX512-NEXT:    addb $9, %al
+; AVX512-NEXT:    movzbl %al, %eax
+; AVX512-NEXT:    retq
+  %res = call i32 @llvm.experimental.cttz.elts.i32.v8i1(<8 x i1> zeroinitializer, i1 false)
+  ret i32 %res
+}
+
 declare i8 @llvm.experimental.cttz.elts.i8.v8i16(<8 x i16>, i1)
 declare i16 @llvm.experimental.cttz.elts.i16.v4i32(<4 x i32>, i1)



More information about the llvm-commits mailing list