[SelectionDAGLegalize] Improve promotion of CTLZ

[jonpa]

Summary

When CTLZ is first promoted from i16 to i32 by type-legalization, and then again from i32 to i64 during Legalization, two unfolded subtractions of constants will remain in output.

This was discussed on llvm-dev, see http://lists.llvm.org/pipermail/llvm-dev/2019-February/thread.html#130115. Sanjay pointed out that the reason for this missed optimization is that DAGCombiner::visitTRUNCATE() is restricted to handle this pre-legalize only.

As an alternative to enabling the DAGCombiner post-legalize also, this is a patch that addresses this issue directly by detecting this case during promotion of CTLZ and propagating the constant to the present add (of negative constant) instead of always creating a new sub.

Does this look reasonable? This is NFC on SystemZ/SPEC, but improves tests both on SystemZ and AMDGPU.

[spatel]

Adding more potential reviewers for legalization changes. I'm not sure how much combining/optimization we want to do in here.

[efriedma]

Generally we should prefer to perform combines in DAGCombine in cases where it's straightforward. (There are a few scattered optimizations in legalization, but mostly for cases that can't simplify after legalization, like libcalls.)

[spatel]

@jonpa - IIRC, some targets infinite loop if we always allow the generic trunc fold in DAGCombiner after legalization, but you might be able to constrain it using the existing TLI hooks:
IsDesirableToPromoteOp()
isTypeDesirableForOp()

[arsenm]

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

4155

Can you do the sub in APInt and avoid getSExtValue?

[jonpa]

Generally we should prefer to perform combines in DAGCombine in cases where it's straightforward.

Given this, it seems reasonable to abandon this patch for https://reviews.llvm.org/D58521, which enables the DAGCombiner post-legalize which then handles this case.

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

if (Node->getOpcode() == ISD::CTTZ) {
       Tmp1 = DAG.getNode(ISD::OR, dl, NVT, Tmp1,
                          DAG.getConstant(TopBit, dl, NVT));
     }
     // Perform the larger operation. For CTPOP and CTTZ_ZERO_UNDEF, this is
     // already the correct result.
     Tmp1 = DAG.getNode(Node->getOpcode(), dl, NVT, Tmp1);
     if (Node->getOpcode() == ISD::CTLZ ||
         Node->getOpcode() == ISD::CTLZ_ZERO_UNDEF) {
+      unsigned Bits = NVT.getSizeInBits() - OVT.getSizeInBits();
+      // If type-legalization has already promoted the CTLZ once, replace the
+      // add with an updated constant to avoid the extra unfolded add.  This
+      // is needed as long as DAGCombiner::visitTRUNCATE() will not help this
+      // case post-legalize.
+      SDValue NewAdd;
+      if (Node->hasOneUse() && Node->use_begin()->getOpcode() == ISD::ADD) {
+        SDNode *Add = *Node->use_begin();
+        if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(Add->getOperand(1))) {
+          NewAdd = DAG.getNode(ISD::ADD, dl, OVT, Add->getOperand(0),
+                          DAG.getConstant(C->getSExtValue() - Bits, dl, OVT));
+          ReplaceNode(SDValue(Add, 0), NewAdd);
+        }
+      }
+      if (!NewAdd.getNode())
         // Tmp1 = Tmp1 - (sizeinbits(NVT) - sizeinbits(Old VT))
         Tmp1 = DAG.getNode(ISD::SUB, dl, NVT, Tmp1,
-                          DAG.getConstant(NVT.getSizeInBits() -
-                                          OVT.getSizeInBits(), dl, NVT));
+                           DAG.getConstant(Bits, dl, NVT));
     }
     Results.push_back(DAG.getNode(ISD::TRUNCATE, dl, OVT, Tmp1));
     break;
   case ISD::BITREVERSE:
   case ISD::BSWAP: {
     unsigned DiffBits = NVT.getSizeInBits() - OVT.getSizeInBits();
     Tmp1 = DAG.getNode(ISD::ZERO_EXTEND, dl, NVT, Node->getOperand(0));
     Tmp1 = DAG.getNode(Node->getOpcode(), dl, NVT, Tmp1);

test/CodeGen/AMDGPU/ctlz.ll

 ; SI-DAG: v_ffbh_u32_e32 [[FFBH:v[0-9]+]], [[VAL]]
 ; VI-DAG: v_ffbh_u32_sdwa [[FFBH:v[0-9]+]], [[VAL]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_0
 ; SI: v_cmp_ne_u32_e32 vcc, 0, [[VAL]]
 ; VI: v_cmp_ne_u16_e32 vcc, 0, [[VAL]]
 
 ; GCN: v_cndmask_b32_e32 [[SELECT:v[0-9]+]], 32, [[FFBH]], vcc
 
 ; SI: v_subrev_i32_e32 [[RESULT:v[0-9]+]], vcc, 24, [[SELECT]]
-; VI: v_add_u32_e32 [[RESULT:v[0-9]+]], vcc, -16, [[SELECT]]
+; VI: v_u16_e32 [[RESULT:v[0-9]+]], 24, [[SELECT]]
 ; GCN: buffer_store_byte [[RESULT]],
 ; GCN: s_endpgm
 define amdgpu_kernel void @v_ctlz_i8(i8 addrspace(1)* noalias %out, i8 addrspace(1)* noalias %valptr) nounwind {
   %val = load i8, i8 addrspace(1)* %valptr
   %ctlz = call i8 @llvm.ctlz.i8(i8 %val, i1 false) nounwind readnone
   store i8 %ctlz, i8 addrspace(1)* %out
   ret void
 }

test/CodeGen/SystemZ/scalar-ctlz.ll

 ; RUN: llc < %s -mtriple=s390x-linux-gnu -mcpu=z13 | FileCheck %s
-;
-; FIXME: two consecutive immediate adds not fused in i16/i8 functions.
 
 declare i64 @llvm.ctlz.i64(i64, i1)
 declare i32 @llvm.ctlz.i32(i32, i1)
 declare i16 @llvm.ctlz.i16(i16, i1)
 declare i8 @llvm.ctlz.i8(i8, i1)
 
 define i64 @f0(i64 %arg) {
 ; CHECK-LABEL: f0:

 }
 
 define i16 @f4(i16 %arg) {
 ; CHECK-LABEL: f4:
 ; CHECK-LABEL: %bb.0:
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: llghr %r0, %r2
 ; CHECK-NEXT: flogr %r2, %r0
-; CHECK-NEXT: aghi  %r2, -32
-; CHECK-NEXT: ahi   %r2, -16
+; CHECK-NEXT: ahi   %r2, -48
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: br %r14
   %1 = tail call i16 @llvm.ctlz.i16(i16 %arg, i1 false)
   ret i16 %1
 }
 
 define i16 @f5(i16 %arg) {
 ; CHECK-LABEL: f5:
 ; CHECK-LABEL: %bb.0:
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: llghr %r0, %r2
 ; CHECK-NEXT: flogr %r2, %r0
-; CHECK-NEXT: aghi  %r2, -32
-; CHECK-NEXT: ahi   %r2, -16
+; CHECK-NEXT: ahi   %r2, -48
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: br %r14
   %1 = tail call i16 @llvm.ctlz.i16(i16 %arg, i1 true)
   ret i16 %1
 }
 
 define i8 @f6(i8 %arg) {
 ; CHECK-LABEL: f6:
 ; CHECK-LABEL: %bb.0:
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: llgcr %r0, %r2
 ; CHECK-NEXT: flogr %r2, %r0
-; CHECK-NEXT: aghi  %r2, -32
-; CHECK-NEXT: ahi   %r2, -24
+; CHECK-NEXT: ahi   %r2, -56
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: br %r14
   %1 = tail call i8 @llvm.ctlz.i8(i8 %arg, i1 false)
   ret i8 %1
 }
 
 define i8 @f7(i8 %arg) {
 ; CHECK-LABEL: f7:
 ; CHECK-LABEL: %bb.0:
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: llgcr %r0, %r2
 ; CHECK-NEXT: flogr %r2, %r0
-; CHECK-NEXT: aghi  %r2, -32
-; CHECK-NEXT: ahi   %r2, -24
+; CHECK-NEXT: ahi   %r2, -56
 ; CHECK-NEXT: # kill
 ; CHECK-NEXT: br %r14
   %1 = tail call i8 @llvm.ctlz.i8(i8 %arg, i1 true)
   ret i8 %1
 }