address PR comments

Author: spall

llvm/lib/Target/DirectX/DXIL.td

@@ -553,7 +553,7 @@ def Rbits :  DXILOp<30, unary> {
   let attributes = [Attributes<DXIL1_0, [ReadNone]>];
 }
 
-def CBits :  DXILOp<31, unaryBits> {
+def CountBits :  DXILOp<31, unaryBits> {
   let Doc = "Returns the number of 1 bits in the specified value.";
   let arguments = [OverloadTy];
   let result = Int32Ty;

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -460,7 +460,7 @@ class OpLowerer {
     });
   }
 
-  [[nodiscard]] bool lowerCtpopToCBits(Function &F) {
+  [[nodiscard]] bool lowerCtpopToCountBits(Function &F) {
     IRBuilder<> &IRB = OpBuilder.getIRB();
     Type *Int32Ty = IRB.getInt32Ty();
 
@@ -471,13 +471,11 @@ class OpLowerer {
 
       Type *RetTy = Int32Ty;
       Type *FRT = F.getReturnType();
-      if (FRT->isVectorTy()) {
-        VectorType *VT = cast<VectorType>(FRT);
+      if (const auto *VT = dyn_cast<VectorType>(FRT))
         RetTy = VectorType::get(RetTy, VT);
-      }
 
       Expected<CallInst *> OpCall = OpBuilder.tryCreateOp(
-          dxil::OpCode::CBits, Args, CI->getName(), RetTy);
+          dxil::OpCode::CountBits, Args, CI->getName(), RetTy);
       if (Error E = OpCall.takeError())
         return E;
 
@@ -491,31 +489,36 @@ class OpLowerer {
       unsigned CastOp;
       if (FRT->isIntOrIntVectorTy(16))
         CastOp = Instruction::ZExt;
-      else // must be 64 bits
+      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;
+      }
 
       // It is correct to replace the ctpop with the dxil op and
-      // remove an existing cast iff the cast is the only usage of
-      // the ctpop
-      // can use hasOneUse instead of hasOneUser, because the user
-      // we care about should have one operand
-      if (CI->hasOneUse()) {
-        User *U = CI->user_back();
+      // remove all casts to i32
+      bool nonCastInstr = false;
+      for (User *User : make_early_inc_range(CI->users())) {
         Instruction *I;
-        if (isa<Instruction>(U) && (I = cast<Instruction>(U)) &&
+        if ((I = dyn_cast<Instruction>(User)) != NULL &&
             I->getOpcode() == CastOp && I->getType() == RetTy) {
           I->replaceAllUsesWith(*OpCall);
           I->eraseFromParent();
-          CI->eraseFromParent();
-          return Error::success();
-        }
+        } else
+          nonCastInstr = 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) {
+        Value *Cast =
+            IRB.CreateZExtOrTrunc(*OpCall, F.getReturnType(), "ctpop.cast");
+        CI->replaceAllUsesWith(Cast);
       }
 
-      // It is always correct to replace a ctpop with the dxil op and
-      // a cast
-      Value *Cast =
-          IRB.CreateZExtOrTrunc(*OpCall, F.getReturnType(), "ctpop.cast");
-      CI->replaceAllUsesWith(Cast);
       CI->eraseFromParent();
       return Error::success();
     });
@@ -550,7 +553,7 @@ class OpLowerer {
         HasErrors |= lowerTypedBufferStore(F);
         break;
       case Intrinsic::ctpop:
-        HasErrors |= lowerCtpopToCBits(F);
+        HasErrors |= lowerCtpopToCountBits(F);
         break;
       }
       Updated = true;

llvm/test/CodeGen/DirectX/countbits.ll

@@ -67,4 +67,4 @@ entry:
 declare i16 @llvm.ctpop.i16(i16)
 declare i32 @llvm.ctpop.i32(i32)
 declare i64 @llvm.ctpop.i64(i64)
-declare <4 x i32> @llvm.ctpop.v4i32(<4 x i32>)
+declare <4 x i32> @llvm.ctpop.v4i32(<4 x i32>)