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..e80ed537a2f6c 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: @@ -101,5 +102,25 @@ define i8 @ctz_v8i16_poison(<8 x i16> %a) { ret i8 %res } +define i32 @ctz_zero_v4i1() { +; CHECK-LABEL: ctz_zero_v4i1: +; CHECK: movl $4, %eax +; CHECK-NEXT: retq +; AVX512-LABEL: ctz_zero_v4i1: +; AVX512: addb $5, %al + %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: movl $8, %eax +; CHECK-NEXT: retq +; AVX512-LABEL: ctz_zero_v8i1: +; AVX512: addb $9, %al + %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)