[clang] [llvm] [HLSL] Re-implement countbits with the correct return type (PR #113189)

[Justin Bogner via cfe-commits]

================
@@ -461,6 +461,67 @@ class OpLowerer {
     });
   }
 
+  [[nodiscard]] bool lowerCtpopToCBits(Function &F) {
+    IRBuilder<> &IRB = OpBuilder.getIRB();
+    Type *Int32Ty = IRB.getInt32Ty();
+
+    return replaceFunction(F, [&](CallInst *CI) -> Error {
+      IRB.SetInsertPoint(CI);
+      SmallVector<Value *> Args;
+      Args.append(CI->arg_begin(), CI->arg_end());
+
+      Type *RetTy = Int32Ty;
+      Type *FRT = F.getReturnType();
+      if (FRT->isVectorTy()) {
+        VectorType *VT = cast<VectorType>(FRT);
+        RetTy = VectorType::get(RetTy, VT);
+      }
+
+      Expected<CallInst *> OpCall = OpBuilder.tryCreateOp(
+          dxil::OpCode::CBits, Args, CI->getName(), RetTy);
+      if (Error E = OpCall.takeError())
+        return E;
+
+      // If the result type is 32 bits we can do a direct replacement.
+      if (FRT->isIntOrIntVectorTy(32)) {
+        CI->replaceAllUsesWith(*OpCall);
+        CI->eraseFromParent();
+        return Error::success();
+      }
+
+      unsigned CastOp;
+      if (FRT->isIntOrIntVectorTy(16))
+        CastOp = Instruction::ZExt;
+      else // must be 64 bits
+        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
----------------
bogner wrote:

Technically this would be fine as long as all of the uses are the same type of cast, I think. This looks conservatively correct though.

https://github.com/llvm/llvm-project/pull/113189