From 1dbb747d598280d57ee9f7f33793960e2e90535a Mon Sep 17 00:00:00 2001
From: bjorn3 <bjorn3@users.noreply.github.com>
Date: Fri, 17 Dec 2021 20:39:53 +0100
Subject: [PATCH] Fix popcnt for small integers

---
 cranelift/codegen/src/isa/x64/lower.rs | 27 +++++++++++++++++++++-----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/cranelift/codegen/src/isa/x64/lower.rs b/cranelift/codegen/src/isa/x64/lower.rs
index f75a9511bd..d0d0debfdf 100644
--- a/cranelift/codegen/src/isa/x64/lower.rs
+++ b/cranelift/codegen/src/isa/x64/lower.rs
@@ -2236,14 +2236,25 @@ fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
         Opcode::Popcnt => {
             let ty_tmp = ty.unwrap();
             if !ty_tmp.is_vector() {
-                let (ext_spec, ty) = match ctx.input_ty(insn, 0) {
-                    types::I8 | types::I16 => (Some(ExtSpec::ZeroExtendTo32), types::I32),
-                    a if a == types::I32 || a == types::I64 || a == types::I128 => (None, a),
-                    _ => unreachable!(),
-                };
+                let ty = ctx.input_ty(insn, 0);
 
                 if isa_flags.use_popcnt() {
                     match ty {
+                        types::I8 | types::I16 => {
+                            let src = RegMem::reg(extend_input_to_reg(
+                                ctx,
+                                inputs[0],
+                                ExtSpec::ZeroExtendTo32,
+                            ));
+                            let dst = get_output_reg(ctx, outputs[0]).only_reg().unwrap();
+                            ctx.emit(Inst::unary_rm_r(
+                                OperandSize::from_ty(types::I32),
+                                UnaryRmROpcode::Popcnt,
+                                src,
+                                dst,
+                            ));
+                            return Ok(());
+                        }
                         types::I32 | types::I64 => {
                             let src = input_to_reg_mem(ctx, inputs[0]);
                             let dst = get_output_reg(ctx, outputs[0]).only_reg().unwrap();
@@ -2299,6 +2310,12 @@ fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
                     }
                 }
 
+                let (ext_spec, ty) = match ty {
+                    types::I8 | types::I16 => (Some(ExtSpec::ZeroExtendTo32), types::I32),
+                    a if a == types::I32 || a == types::I64 || a == types::I128 => (None, a),
+                    _ => unreachable!(),
+                };
+
                 let (srcs, ty): (SmallVec<[RegMem; 2]>, Type) = if let Some(ext_spec) = ext_spec {
                     (
                         smallvec![RegMem::reg(extend_input_to_reg(ctx, inputs[0], ext_spec))],