[llvm] [DAG] canCreateUndefOrPoison - add handling for CTTZ/CTLZ_ZERO_UNDEF nodes (PR #146501)

[Simon Pilgrim via llvm-commits]

https://github.com/RKSimon created https://github.com/llvm/llvm-project/pull/146501

CTTZ/CTLZ_ZERO_UNDEF nodes can only create poison if the source value is zero - so check with isKnownNeverZero

Pulled out of #146361 and reapplied now that #146490 has landed.

>From 0f06eb82c83f05b1152e84e68a8eb805f0447fc3 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Tue, 1 Jul 2025 11:25:24 +0100
Subject: [PATCH] [DAG] canCreateUndefOrPoison - add handling for
 CTTZ/CTLZ_ZERO_UNDEF nodes

CTTZ/CTLZ_ZERO_UNDEF nodes can only create poison if the source value is zero - so check with isKnownNeverZero

Pulled out of #146361 and reapplied now that #146490 has landed.
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  6 +++++
 llvm/test/CodeGen/X86/freeze-unary.ll         | 22 +++++--------------
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index ad941a1964683..88df51bb3cfb8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5620,6 +5620,12 @@ bool SelectionDAG::canCreateUndefOrPoison(SDValue Op, const APInt &DemandedElts,
                                              PoisonOnly, Depth + 1) ||
            !getValidMaximumShiftAmount(Op, DemandedElts, Depth + 1);
 
+  case ISD::CTTZ_ZERO_UNDEF:
+  case ISD::CTLZ_ZERO_UNDEF:
+    // If the amount is zero then the result will be poison.
+    // TODO: Add isKnownNeverZero DemandedElts handling.
+    return !isKnownNeverZero(Op.getOperand(0), Depth + 1);
+
   case ISD::SCALAR_TO_VECTOR:
     // Check if we demand any upper (undef) elements.
     return !PoisonOnly && DemandedElts.ugt(1);
diff --git a/llvm/test/CodeGen/X86/freeze-unary.ll b/llvm/test/CodeGen/X86/freeze-unary.ll
index c85c831a25342..8602c385af834 100644
--- a/llvm/test/CodeGen/X86/freeze-unary.ll
+++ b/llvm/test/CodeGen/X86/freeze-unary.ll
@@ -179,21 +179,15 @@ define i32 @freeze_ctlz_undef_nonzero(i32 %a0) nounwind {
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    orl $1, %eax
-; X86-NEXT:    bsrl %eax, %ecx
-; X86-NEXT:    xorl $31, %ecx
-; X86-NEXT:    testl %eax, %eax
-; X86-NEXT:    movl $32, %eax
-; X86-NEXT:    cmovnel %ecx, %eax
+; X86-NEXT:    bsrl %eax, %eax
+; X86-NEXT:    xorl $31, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_ctlz_undef_nonzero:
 ; X64:       # %bb.0:
 ; X64-NEXT:    orl $1, %edi
-; X64-NEXT:    bsrl %edi, %ecx
-; X64-NEXT:    xorl $31, %ecx
-; X64-NEXT:    testl %edi, %edi
-; X64-NEXT:    movl $32, %eax
-; X64-NEXT:    cmovnel %ecx, %eax
+; X64-NEXT:    bsrl %edi, %eax
+; X64-NEXT:    xorl $31, %eax
 ; X64-NEXT:    retq
   %y = or i32 %a0, 1
   %x = call i32 @llvm.ctlz.i32(i32 %y, i1 -1)
@@ -250,17 +244,13 @@ define i32 @freeze_cttz_undef_nonzero(i32 %a0) nounwind {
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    orl $1, %eax
-; X86-NEXT:    bsfl %eax, %ecx
-; X86-NEXT:    movl $32, %eax
-; X86-NEXT:    cmovnel %ecx, %eax
+; X86-NEXT:    rep bsfl %eax, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_cttz_undef_nonzero:
 ; X64:       # %bb.0:
 ; X64-NEXT:    orl $1, %edi
-; X64-NEXT:    bsfl %edi, %ecx
-; X64-NEXT:    movl $32, %eax
-; X64-NEXT:    cmovnel %ecx, %eax
+; X64-NEXT:    rep bsfl %edi, %eax
 ; X64-NEXT:    retq
   %y = or i32 %a0, 1
   %x = call i32 @llvm.cttz.i32(i32 %y, i1 -1)