address latest pr comments + extra tests

Author: spall

clang/test/CodeGenHLSL/builtins/countbits.hlsl

@@ -10,6 +10,13 @@ uint test_countbits_ushort(uint16_t p0)
 {
 	return countbits(p0);
 }
+// CHECK-LABEL: test_countbits_short
+// CHECK: [[A:%.*]] = call i16 @llvm.ctpop.i16
+// CHECK-NEXT: sext i16 [[A]] to i32
+uint test_countbits_short(int16_t p0)
+{
+	return countbits(p0);
+}
 // CHECK-LABEL: test_countbits_ushort2
 // CHECK: [[A:%.*]] = call <2 x i16> @llvm.ctpop.v2i16
 // CHECK-NEXT: zext <2 x i16> [[A]] to <2 x i32>
@@ -39,6 +46,12 @@ uint test_countbits_uint(uint p0)
 {
 	return countbits(p0);
 }
+// CHECK-LABEL: test_countbits_int
+// CHECK: call i32 @llvm.ctpop.i32
+uint test_countbits_int(int p0)
+{
+	return countbits(p0);
+}
 // CHECK-LABEL: test_countbits_uint2
 // CHECK: call <2 x i32> @llvm.ctpop.v2i32
 uint2 test_countbits_uint2(uint2 p0)
@@ -65,6 +78,13 @@ uint test_countbits_long(uint64_t p0)
 {
 	return countbits(p0);
 }
+// CHECK-LABEL: test_countbits_slong
+// CHECK: [[A:%.*]] = call i64 @llvm.ctpop.i64
+// CHECK-NEXT: trunc i64 [[A]] to i32
+uint test_countbits_slong(int64_t p0)
+{
+	return countbits(p0);
+}
 // CHECK-LABEL: test_countbits_long2
 // CHECK: [[A:%.*]] = call <2 x i64> @llvm.ctpop.v2i64
 // CHECK-NEXT: trunc <2 x i64> [[A]] to <2 x i32>

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -532,33 +532,36 @@ class OpLowerer {
       }
 
       unsigned CastOp;
-      if (FRT->isIntOrIntVectorTy(16))
+      unsigned CastOp2;
+      if (FRT->isIntOrIntVectorTy(16)) {
         CastOp = Instruction::ZExt;
-      else { // must be 64 bits
+        CastOp2 = Instruction::SExt;
+      } else { // must be 64 bits
         assert(FRT->isIntOrIntVectorTy(64) &&
                "Currently only lowering 16, 32, or 64 bit ctpop to CountBits \
                 is supported.");
         CastOp = Instruction::Trunc;
+        CastOp2 = Instruction::Trunc;
       }
 
       // It is correct to replace the ctpop with the dxil op and
       // remove all casts to i32
-      bool nonCastInstr = false;
+      bool NeedsCast = false;
       for (User *User : make_early_inc_range(CI->users())) {
-        Instruction *I;
-        if ((I = dyn_cast<Instruction>(User)) != NULL &&
-            I->getOpcode() == CastOp && I->getType() == RetTy) {
+        Instruction *I = dyn_cast<Instruction>(User);
+        if (I && (I->getOpcode() == CastOp || I->getOpcode() == CastOp2) &&
+            I->getType() == RetTy) {
           I->replaceAllUsesWith(*OpCall);
           I->eraseFromParent();
         } else
-          nonCastInstr = true;
+          NeedsCast = true;
       }
 
       // It is correct to replace a ctpop with the dxil op and
       // a cast from i32 to the return type of the ctpop
       // the cast is emitted here if there is a non-cast to i32
       // instr which uses the ctpop
-      if (nonCastInstr) {
+      if (NeedsCast) {
         Value *Cast =
             IRB.CreateZExtOrTrunc(*OpCall, F.getReturnType(), "ctpop.cast");
         CI->replaceAllUsesWith(Cast);

llvm/test/CodeGen/DirectX/countbits.ll

@@ -20,6 +20,15 @@ entry:
   ret i32 %elt.zext
 }
 
+define noundef i32 @test_countbits_short3(i16 noundef %a) {
+entry:
+; CHECK: [[A:%.*]] = call i32 @dx.op.unaryBits.i16(i32 31, i16 %{{.*}})
+; CHECK-NEXT: ret i32 [[A]]
+  %elt.ctpop = call i16 @llvm.ctpop.i16(i16 %a)
+  %elt.sext = sext i16 %elt.ctpop to i32
+  ret i32 %elt.sext
+}
+
 define noundef i32 @test_countbits_int(i32 noundef %a) {
 entry:
 ; CHECK: [[A:%.*]] = call i32 @dx.op.unaryBits.i32(i32 31, i32 %{{.*}})