address comments + cleanup

Author: AlexMaclean

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

@@ -5113,7 +5113,8 @@ void SelectionDAGLegalize::PromoteNode(SDNode *Node) {
                           DAG.getConstant(NVT.getSizeInBits() -
                                           OVT.getSizeInBits(), dl, NVT));
     }
-    Results.push_back(DAG.getNode(ISD::TRUNCATE, dl, OVT, Tmp1));
+    Results.push_back(
+        DAG.getNode(ISD::TRUNCATE, dl, OVT, Tmp1, SDNodeFlags::NoWrap));
     break;
   }
   case ISD::CTLZ_ZERO_UNDEF: {

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -767,8 +767,10 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setOperationAction({ISD::ABS, ISD::SMIN, ISD::SMAX, ISD::UMIN, ISD::UMAX},
                      {MVT::i16, MVT::i32, MVT::i64}, Legal);
 
+  setOperationAction({ISD::CTPOP, ISD::CTLZ, ISD::CTLZ_ZERO_UNDEF}, MVT::i16,
+                     Promote);
   setOperationAction({ISD::CTPOP, ISD::CTLZ}, MVT::i32, Legal);
-  setOperationAction({ISD::CTPOP, ISD::CTLZ}, {MVT::i16, MVT::i64}, Custom);
+  setOperationAction({ISD::CTPOP, ISD::CTLZ}, MVT::i64, Custom);
 
   setI16x2OperationAction(ISD::ABS, MVT::v2i16, Legal, Custom);
   setI16x2OperationAction(ISD::SMIN, MVT::v2i16, Legal, Custom);
@@ -2743,40 +2745,17 @@ static SDValue LowerIntrinsicVoid(SDValue Op, SelectionDAG &DAG) {
   return Op;
 }
 
-static SDValue lowerCTPOP(SDValue Op, SelectionDAG &DAG) {
+// In PTX 64-bit CTLZ and CTPOP are supported, but they return a 32-bit value.
+// Lower these into a node returning the correct type which is zero-extended
+// back to the correct size.
+static SDValue lowerCTLZCTPOP(SDValue Op, SelectionDAG &DAG) {
   SDValue V = Op->getOperand(0);
-  SDLoc DL(Op);
-
-  if (V.getValueType() == MVT::i16) {
-    SDValue Zext = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i32, V);
-    SDValue CT = DAG.getNode(ISD::CTPOP, DL, MVT::i32, Zext);
-    return DAG.getNode(ISD::TRUNCATE, DL, MVT::i16, CT, SDNodeFlags::NoWrap);
-  }
-  if (V.getValueType() == MVT::i64) {
-    SDValue CT = DAG.getNode(ISD::CTPOP, DL, MVT::i32, V);
-    return DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, CT);
-  }
-  llvm_unreachable("Unexpected CTPOP type to legalize");
-}
+  assert(V.getValueType() == MVT::i64 &&
+         "Unexpected CTLZ/CTPOP type to legalize");
 
-static SDValue lowerCTLZ(SDValue Op, SelectionDAG &DAG) {
-  SDValue V = Op->getOperand(0);
   SDLoc DL(Op);
-
-  if (V.getValueType() == MVT::i16) {
-    SDValue Zext = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i32, V);
-    SDValue CT = DAG.getNode(ISD::CTLZ, DL, MVT::i32, Zext);
-    SDValue Sub =
-        DAG.getNode(ISD::ADD, DL, MVT::i32, CT,
-                    DAG.getConstant(APInt(32, -16, true), DL, MVT::i32),
-                    SDNodeFlags::NoSignedWrap);
-    return DAG.getNode(ISD::TRUNCATE, DL, MVT::i16, Sub, SDNodeFlags::NoWrap);
-  }
-  if (V.getValueType() == MVT::i64) {
-    SDValue CT = DAG.getNode(ISD::CTLZ, DL, MVT::i32, V);
-    return DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, CT);
-  }
-  llvm_unreachable("Unexpected CTLZ type to legalize");
+  SDValue CT = DAG.getNode(Op->getOpcode(), DL, MVT::i32, V);
+  return DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, CT, SDNodeFlags::NonNeg);
 }
 
 SDValue
@@ -2865,9 +2844,8 @@ NVPTXTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
     // Used only for bf16 on SM80, where we select fma for non-ftz operation
     return PromoteBinOpIfF32FTZ(Op, DAG);
   case ISD::CTPOP:
-    return lowerCTPOP(Op, DAG);
   case ISD::CTLZ:
-    return lowerCTLZ(Op, DAG);
+    return lowerCTLZCTPOP(Op, DAG);
 
   default:
     llvm_unreachable("Custom lowering not defined for operation");

llvm/test/CodeGen/NVPTX/ctlz.ll

@@ -130,8 +130,8 @@ define i16 @myctlz_ret16_2(i16 %a) {
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.u16 %r1, [myctlz_ret16_2_param_0];
-; CHECK-NEXT:    clz.b32 %r2, %r1;
-; CHECK-NEXT:    add.s32 %r3, %r2, -16;
+; CHECK-NEXT:    shl.b32 %r2, %r1, 16;
+; CHECK-NEXT:    clz.b32 %r3, %r2;
 ; CHECK-NEXT:    st.param.b32 [func_retval0], %r3;
 ; CHECK-NEXT:    ret;
   %val = call i16 @llvm.ctlz.i16(i16 %a, i1 true) readnone

llvm/test/CodeGen/NVPTX/intrinsics.ll

@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
-; RUN: llc < %s -mtriple=nvptx -mcpu=sm_60 | FileCheck %s
-; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_60 | FileCheck %s
+; RUN: llc < %s -mtriple=nvptx -mcpu=sm_60 | FileCheck %s --check-prefixes=CHECK,CHECK32
+; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_60 | FileCheck %s --check-prefixes=CHECK,CHECK64
 ; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -mtriple=nvptx -mcpu=sm_60 | %ptxas-verify %}
 ; RUN: %if ptxas %{ llc < %s -mtriple=nvptx64 -mcpu=sm_60 | %ptxas-verify %}
 
@@ -140,6 +140,28 @@ define i32 @test_popc64_trunc(i64 %a) {
 ; llvm.ctpop.i16 is implemenented by converting to i32, running popc.b32, and
 ; then converting back to i16.
 define void @test_popc16(i16 %a, ptr %b) {
+; CHECK32-LABEL: test_popc16(
+; CHECK32:       {
+; CHECK32-NEXT:    .reg .b32 %r<4>;
+; CHECK32-EMPTY:
+; CHECK32-NEXT:  // %bb.0:
+; CHECK32-NEXT:    ld.param.u16 %r1, [test_popc16_param_0];
+; CHECK32-NEXT:    popc.b32 %r2, %r1;
+; CHECK32-NEXT:    ld.param.u32 %r3, [test_popc16_param_1];
+; CHECK32-NEXT:    st.u16 [%r3], %r2;
+; CHECK32-NEXT:    ret;
+;
+; CHECK64-LABEL: test_popc16(
+; CHECK64:       {
+; CHECK64-NEXT:    .reg .b32 %r<3>;
+; CHECK64-NEXT:    .reg .b64 %rd<2>;
+; CHECK64-EMPTY:
+; CHECK64-NEXT:  // %bb.0:
+; CHECK64-NEXT:    ld.param.u16 %r1, [test_popc16_param_0];
+; CHECK64-NEXT:    popc.b32 %r2, %r1;
+; CHECK64-NEXT:    ld.param.u64 %rd1, [test_popc16_param_1];
+; CHECK64-NEXT:    st.u16 [%rd1], %r2;
+; CHECK64-NEXT:    ret;
   %val = call i16 @llvm.ctpop.i16(i16 %a)
   store i16 %val, ptr %b
   ret void