[LegalizeDAG][X86][AMDGPU] Use ANY_EXTEND instead of ZERO_EXTEND when promoting ISD::CTTZ/CTTZ_ZERO_UNDEF.

[craig.topper]

Summary

For CTTZ we place a set bit just past where the non-promoted type
stopped so the extended bits won't be used for the count. For
CTTZ_ZERO_UNDEF we don't care what happens if no bits are set in
the original type and we end up counting into the extended bits.
So we can just use ANY_EXTEND for both cases.

This matches what is done in type legalization for these operations.
We make no effort to force the upper bits to zero.

[efriedma]

LGTM

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

void SelectionDAGLegalize::PromoteNode(SDNode *Node) {
   SDLoc dl(Node);
   SDValue Tmp1, Tmp2, Tmp3;
   switch (Node->getOpcode()) {
   case ISD::CTTZ:
   case ISD::CTTZ_ZERO_UNDEF:
   case ISD::CTLZ:
   case ISD::CTLZ_ZERO_UNDEF:
   case ISD::CTPOP:
-    // Zero extend the argument.
+    // Zero extend the argument unless its cttz, then use any_extend.
+    if (Node->getOpcode() == ISD::CTTZ ||
+        Node->getOpcode() == ISD::CTTZ_ZERO_UNDEF)
+      Tmp1 = DAG.getNode(ISD::ANY_EXTEND, dl, NVT, Node->getOperand(0));
+    else
       Tmp1 = DAG.getNode(ISD::ZERO_EXTEND, dl, NVT, Node->getOperand(0));
+
     if (Node->getOpcode() == ISD::CTTZ) {
       // The count is the same in the promoted type except if the original
       // value was zero.  This can be handled by setting the bit just off
       // the top of the original type.
       auto TopBit = APInt::getOneBitSet(NVT.getSizeInBits(),
                                         OVT.getSizeInBits());
       Tmp1 = DAG.getNode(ISD::OR, dl, NVT, Tmp1,
                          DAG.getConstant(TopBit, dl, NVT));

llvm/test/CodeGen/AMDGPU/cttz_zero_undef.ll

define amdgpu_kernel void @s_cttz_zero_undef_i64_with_select(i64 addrspace(1)* noalias %out, i64 %val) nounwind {
   %cttz_ret = icmp ne i64 %val, 0
   %ret = select i1 %cttz_ret, i64 %cttz, i64 32
   store i64 %cttz, i64 addrspace(1)* %out, align 4
   ret void
 }
 
 ; FUNC-LABEL: {{^}}v_cttz_zero_undef_i8_with_select:
 ; SI-NOSDWA: v_ffbl_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}
-; SI-SDWA: v_ffbl_b32_sdwa
+; SI-SDWA: v_ffbl_b32_e32
 ; EG: MEM_RAT MSKOR
 define amdgpu_kernel void @v_cttz_zero_undef_i8_with_select(i8 addrspace(1)* noalias %out, i8 addrspace(1)* nocapture readonly %arrayidx) nounwind {
   %val = load i8, i8 addrspace(1)* %arrayidx, align 1
   %cttz = tail call i8 @llvm.cttz.i8(i8 %val, i1 true) nounwind readnone
   %cttz_ret = icmp ne i8 %val, 0
   %ret = select i1 %cttz_ret, i8 %cttz, i8 32
   store i8 %ret, i8 addrspace(1)* %out, align 4
   ret void
 }
 
 ; FUNC-LABEL: {{^}}v_cttz_zero_undef_i16_with_select:
 ; SI-NOSDWA: v_ffbl_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}
-; SI-SDWA: v_ffbl_b32_sdwa
+; SI-SDWA: v_ffbl_b32_e32
 ; EG: MEM_RAT MSKOR
 define amdgpu_kernel void @v_cttz_zero_undef_i16_with_select(i16 addrspace(1)* noalias %out, i16 addrspace(1)* nocapture readonly %arrayidx) nounwind {
   %val = load i16, i16 addrspace(1)* %arrayidx, align 1
   %cttz = tail call i16 @llvm.cttz.i16(i16 %val, i1 true) nounwind readnone
   %cttz_ret = icmp ne i16 %val, 0
   %ret = select i1 %cttz_ret, i16 %cttz, i16 32
   store i16 %ret, i16 addrspace(1)* %out, align 4
   ret void
define amdgpu_kernel void @v_cttz_i32_sel_ne_bitwidth(i32 addrspace(1)* noalias %out, i32 addrspace(1)* nocapture readonly %arrayidx) nounwind {
   %sel = select i1 %cmp, i32 %ctlz, i32 -1
   store i32 %sel, i32 addrspace(1)* %out
   ret void
 }
 
 ; FUNC-LABEL: {{^}}v_cttz_i8_sel_eq_neg1:
 ; SI: {{buffer|flat}}_load_ubyte
 ; SI-NOSDWA: v_ffbl_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}
-; SI-SDWA: v_ffbl_b32_sdwa
+; SI-SDWA: v_ffbl_b32_e32
 ; EG: MEM_RAT MSKOR
 ; EG: FFBL_INT
  define amdgpu_kernel void @v_cttz_i8_sel_eq_neg1(i8 addrspace(1)* noalias %out, i8 addrspace(1)* nocapture readonly %arrayidx) nounwind {
   %val = load i8, i8 addrspace(1)* %arrayidx, align 1
   %ctlz = call i8 @llvm.cttz.i8(i8 %val, i1 false) nounwind readnone
   %cmp = icmp eq i8 %val, 0
   %sel = select i1 %cmp, i8 -1, i8 %ctlz
   store i8 %sel, i8 addrspace(1)* %out

llvm/test/CodeGen/X86/clz.ll

 declare i8 @llvm.ctlz.i8(i8, i1)
 declare i16 @llvm.ctlz.i16(i16, i1)
 declare i32 @llvm.ctlz.i32(i32, i1)
 declare i64 @llvm.ctlz.i64(i64, i1)
 
 define i8 @cttz_i8(i8 %x)  {
 ; X32-LABEL: cttz_i8:
 ; X32:       # %bb.0:
-; X32-NEXT:    movzbl {{[0-9]+}}(%esp), %eax
-; X32-NEXT:    bsfl %eax, %eax
+; X32-NEXT:    bsfl {{[0-9]+}}(%esp), %eax
 ; X32-NEXT:    # kill: def $al killed $al killed $eax
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: cttz_i8:
 ; X64:       # %bb.0:
-; X64-NEXT:    movzbl %dil, %eax
-; X64-NEXT:    bsfl %eax, %eax
+; X64-NEXT:    bsfl %edi, %eax
 ; X64-NEXT:    # kill: def $al killed $al killed $eax
 ; X64-NEXT:    retq
 ;
 ; X32-CLZ-LABEL: cttz_i8:
 ; X32-CLZ:       # %bb.0:
-; X32-CLZ-NEXT:    movzbl {{[0-9]+}}(%esp), %eax
-; X32-CLZ-NEXT:    tzcntl %eax, %eax
+; X32-CLZ-NEXT:    tzcntl {{[0-9]+}}(%esp), %eax
 ; X32-CLZ-NEXT:    # kill: def $al killed $al killed $eax
 ; X32-CLZ-NEXT:    retl
 ;
 ; X64-CLZ-LABEL: cttz_i8:
 ; X64-CLZ:       # %bb.0:
-; X64-CLZ-NEXT:    movzbl %dil, %eax
-; X64-CLZ-NEXT:    tzcntl %eax, %eax
+; X64-CLZ-NEXT:    tzcntl %edi, %eax
 ; X64-CLZ-NEXT:    # kill: def $al killed $al killed $eax
 ; X64-CLZ-NEXT:    retq
   %tmp = call i8 @llvm.cttz.i8( i8 %x, i1 true )
   ret i8 %tmp
 }
 
 define i16 @cttz_i16(i16 %x)  {
 ; X32-LABEL: cttz_i16:

 ; X64-NEXT:    retq
 ; X64-NEXT:  .LBB12_1:
 ; X64-NEXT:    movb $8, %al
 ; X64-NEXT:    # kill: def $al killed $al killed $eax
 ; X64-NEXT:    retq
 ;
 ; X32-CLZ-LABEL: cttz_i8_zero_test:
 ; X32-CLZ:       # %bb.0:
-; X32-CLZ-NEXT:    movzbl {{[0-9]+}}(%esp), %eax
-; X32-CLZ-NEXT:    orl $256, %eax # imm = 0x100
+; X32-CLZ-NEXT:    movl $256, %eax # imm = 0x100
+; X32-CLZ-NEXT:    orl {{[0-9]+}}(%esp), %eax
 ; X32-CLZ-NEXT:    tzcntl %eax, %eax
 ; X32-CLZ-NEXT:    # kill: def $al killed $al killed $eax
 ; X32-CLZ-NEXT:    retl
 ;
 ; X64-CLZ-LABEL: cttz_i8_zero_test:
 ; X64-CLZ:       # %bb.0:
-; X64-CLZ-NEXT:    movzbl %dil, %eax
-; X64-CLZ-NEXT:    orl $256, %eax # imm = 0x100
-; X64-CLZ-NEXT:    tzcntl %eax, %eax
+; X64-CLZ-NEXT:    orl $256, %edi # imm = 0x100
+; X64-CLZ-NEXT:    tzcntl %edi, %eax
 ; X64-CLZ-NEXT:    # kill: def $al killed $al killed $eax
 ; X64-CLZ-NEXT:    retq
   %tmp1 = call i8 @llvm.cttz.i8(i8 %n, i1 false)
   ret i8 %tmp1
 }
 
 ; Generate a test and branch to handle zero inputs because bsr/bsf are very slow.
 define i16 @cttz_i16_zero_test(i16 %n) {