1 files changed, 36 insertions, 0 deletions
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index f4e10e32484..499067364b0 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -619,6 +619,42 @@ LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
   switch (MI.getOpcode()) {
   default:
     return UnableToLegalize;
+  case TargetOpcode::G_CTTZ:
+  case TargetOpcode::G_CTTZ_ZERO_UNDEF:
+  case TargetOpcode::G_CTLZ:
+  case TargetOpcode::G_CTLZ_ZERO_UNDEF:
+  case TargetOpcode::G_CTPOP: {
+    // First ZEXT the input.
+    auto MIBSrc = MIRBuilder.buildZExt(WideTy, MI.getOperand(1).getReg());
+    LLT CurTy = MRI.getType(MI.getOperand(0).getReg());
+    if (MI.getOpcode() == TargetOpcode::G_CTTZ) {
+      // The count is the same in the larger type except if the original
+      // value was zero.  This can be handled by setting the bit just off
+      // the top of the original type.
+      auto TopBit =
+          APInt::getOneBitSet(WideTy.getSizeInBits(), CurTy.getSizeInBits());
+      MIBSrc = MIRBuilder.buildInstr(
+          TargetOpcode::G_OR, WideTy, MIBSrc,
+          MIRBuilder.buildConstant(WideTy, TopBit.getSExtValue()));
+    }
+    // Perform the operation at the larger size.
+    auto MIBNewOp = MIRBuilder.buildInstr(MI.getOpcode(), WideTy, MIBSrc);
+    // This is already the correct result for CTPOP and CTTZs
+    if (MI.getOpcode() == TargetOpcode::G_CTLZ ||
+        MI.getOpcode() == TargetOpcode::G_CTLZ_ZERO_UNDEF) {
+      // The correct result is NewOp - (Difference in widety and current ty).
+      unsigned SizeDiff = WideTy.getSizeInBits() - CurTy.getSizeInBits();
+      MIBNewOp =
+          MIRBuilder.buildInstr(TargetOpcode::G_SUB, WideTy, MIBNewOp,
+                                MIRBuilder.buildConstant(WideTy, SizeDiff));
+    }
+    auto &TII = *MI.getMF()->getSubtarget().getInstrInfo();
+    // Make the original instruction a trunc now, and update it's source.
+    MI.setDesc(TII.get(TargetOpcode::G_TRUNC));
+    MI.getOperand(1).setReg(MIBNewOp->getOperand(0).getReg());
+    MIRBuilder.recordInsertion(&MI);
+    return Legalized;
+  }
 
   case TargetOpcode::G_ADD:
   case TargetOpcode::G_AND: