Address review comments.

Author: Leon Clark

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

@@ -446,7 +446,9 @@ AMDGPUTargetLowering::AMDGPUTargetLowering(const TargetMachine &TM,
       {ISD::CTTZ, ISD::CTTZ_ZERO_UNDEF, ISD::CTLZ, ISD::CTLZ_ZERO_UNDEF},
       MVT::i64, Custom);
 
-  setOperationAction({ISD::CTLZ, ISD::CTLZ_ZERO_UNDEF}, MVT::i8, Custom);
+  for (auto VT : {MVT::i8, MVT::i16}) {
+    setOperationAction({ISD::CTLZ, ISD::CTLZ_ZERO_UNDEF}, VT, Custom);
+  }
 
   static const MVT::SimpleValueType VectorIntTypes[] = {
       MVT::v2i32, MVT::v3i32, MVT::v4i32, MVT::v5i32, MVT::v6i32, MVT::v7i32,
@@ -1402,7 +1404,8 @@ void AMDGPUTargetLowering::ReplaceNodeResults(SDNode *N,
     return;
   case ISD::CTLZ:
   case ISD::CTLZ_ZERO_UNDEF:
-    replaceCTLZResults(SDValue(N, 0u), DAG, Results);
+    if (auto Lowered = lowerCTLZResults(SDValue(N, 0u), DAG))
+      Results.push_back(Lowered);
     return;
   default:
     return;
@@ -3069,23 +3072,24 @@ static bool isCttzOpc(unsigned Opc) {
   return Opc == ISD::CTTZ || Opc == ISD::CTTZ_ZERO_UNDEF;
 }
 
-void AMDGPUTargetLowering::replaceCTLZResults(
-    SDValue Op, SelectionDAG &DAG, SmallVectorImpl<SDValue> &Results) const {
+SDValue AMDGPUTargetLowering::lowerCTLZResults(SDValue Op,
+                                               SelectionDAG &DAG) const {
   auto SL = SDLoc(Op);
   auto Arg = Op.getOperand(0u);
   auto ResultVT = Op.getValueType();
 
-  if (ResultVT != MVT::i8)
-    return;
+  if (!(ResultVT == MVT::i8 || ResultVT == MVT::i16))
+    return {};
 
   assert(isCtlzOpc(Op.getOpcode()));
   assert(ResultVT == Arg.getValueType());
 
-  auto SubVal = DAG.getConstant(24u, SL, MVT::i32);
+  auto const LeadingZeroes = 32u - ResultVT.getFixedSizeInBits();
+  auto SubVal = DAG.getConstant(LeadingZeroes, SL, MVT::i32);
   auto NewOp = DAG.getNode(ISD::ZERO_EXTEND, SL, MVT::i32, Arg);
   NewOp = DAG.getNode(Op.getOpcode(), SL, MVT::i32, NewOp);
   NewOp = DAG.getNode(ISD::SUB, SL, MVT::i32, NewOp, SubVal);
-  Results.push_back(DAG.getNode(ISD::TRUNCATE, SL, ResultVT, NewOp));
+  return DAG.getNode(ISD::TRUNCATE, SL, ResultVT, NewOp);
 }
 
 SDValue AMDGPUTargetLowering::LowerCTLZ_CTTZ(SDValue Op, SelectionDAG &DAG) const {

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.h

@@ -84,8 +84,7 @@ class AMDGPUTargetLowering : public TargetLowering {
                             SDNodeFlags Flags) const;
   SDValue lowerFEXP(SDValue Op, SelectionDAG &DAG) const;
 
-  void replaceCTLZResults(SDValue Op, SelectionDAG &DAG,
-                          SmallVectorImpl<SDValue> &Results) const;
+  SDValue lowerCTLZResults(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerCTLZ_CTTZ(SDValue Op, SelectionDAG &DAG) const;