Fix YJIT backend to account for unsigned int immediates (#6789)

YJIT: x86_64: Fix cmp with number where sign bit is set

Before this commit, we were unconditionally treating unsigned ints as
signed ints when counting the number of bits required for representing
the immediate in machine code. When the size of the immediate matches
the size of the other operand, no sign extension happens, so this was
incorrect. `asm.cmp(opnd64, 0x8000_0000)` panicked even though it's
encodable as `CMP r/m32, imm32`. Large shape ids were impacted by this
issue.

Co-Authored-By: Aaron Patterson <tenderlove@ruby-lang.org>
Co-Authored-By: Alan Wu <alanwu@ruby-lang.org>

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>
Co-authored-by: Alan Wu <alanwu@ruby-lang.org>

yjit/src/asm/x86_64/mod.rs

@@ -555,8 +555,8 @@ fn write_rm_multi(cb: &mut CodeBlock, op_mem_reg8: u8, op_mem_reg_pref: u8, op_r
 
     // Check the size of opnd1
     match opnd1 {
-        X86Opnd::Reg(reg) => assert!(reg.num_bits == opnd_size),
+        X86Opnd::Reg(reg) => assert_eq!(reg.num_bits, opnd_size),
-        X86Opnd::Mem(mem) => assert!(mem.num_bits == opnd_size),
+        X86Opnd::Mem(mem) => assert_eq!(mem.num_bits, opnd_size),
         X86Opnd::Imm(imm) => assert!(imm.num_bits <= opnd_size),
         X86Opnd::UImm(uimm) => assert!(uimm.num_bits <= opnd_size),
         _ => ()
@@ -606,7 +606,14 @@ fn write_rm_multi(cb: &mut CodeBlock, op_mem_reg8: u8, op_mem_reg_pref: u8, op_r
         },
         // R/M + UImm
         (_, X86Opnd::UImm(uimm)) => {
-            let num_bits = imm_num_bits(uimm.value.try_into().unwrap());
+            // If the size of left hand operand equals the number of bits
+            // required to represent the right hand immediate, then we
+            // don't care about sign extension when calculating the immediate
+            let num_bits = if opnd0.num_bits() == uimm_num_bits(uimm.value) {
+                uimm_num_bits(uimm.value)
+            } else {
+                imm_num_bits(uimm.value.try_into().unwrap())
+            };
 
             if num_bits <= 8 {
                 // 8-bit immediate
@@ -625,7 +632,7 @@ fn write_rm_multi(cb: &mut CodeBlock, op_mem_reg8: u8, op_mem_reg_pref: u8, op_r
                 write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, op_ext_imm, &[op_mem_imm_lrg]);
                 cb.write_int(uimm.value, if opnd_size > 32 { 32 } else { opnd_size.into() });
             } else {
-                panic!("immediate value too large (num_bits={})", num_bits);
+                panic!("immediate value too large (num_bits={}, num={uimm:?})", num_bits);
             }
         },
         _ => unreachable!()

yjit/src/asm/x86_64/tests.rs

@@ -97,6 +97,7 @@ fn test_cmp() {
     check_bytes("39f9", |cb| cmp(cb, ECX, EDI));
     check_bytes("493b1424", |cb| cmp(cb, RDX, mem_opnd(64, R12, 0)));
     check_bytes("4883f802", |cb| cmp(cb, RAX, imm_opnd(2)));
+    check_bytes("81f900000080", |cb| cmp(cb, ECX, uimm_opnd(0x8000_0000)));
 }
 
 #[test]
@@ -357,6 +358,14 @@ fn test_sub() {
     check_bytes("4883e802", |cb| sub(cb, RAX, imm_opnd(2)));
 }
 
+#[test]
+#[should_panic]
+fn test_sub_uimm_too_large() {
+    // This immedaite becomes a different value after
+    // sign extension, so not safe to encode.
+    check_bytes("ff", |cb| sub(cb, RCX, uimm_opnd(0x8000_0000)));
+}
+
 #[test]
 fn test_test() {
     check_bytes("84c0", |cb| test(cb, AL, AL));

yjit/src/backend/x86_64/mod.rs

@@ -548,8 +548,24 @@ impl Assembler
 
                 // Compare
                 Insn::Cmp { left, right } => {
-                    let emitted = emit_64bit_immediate(cb, right);
+                    let num_bits = match right {
-                    cmp(cb, left.into(), emitted);
+                        Opnd::Imm(value) => Some(imm_num_bits(*value)),
+                        Opnd::UImm(value) => Some(uimm_num_bits(*value)),
+                        _ => None
+                    };
+
+                    // If the immediate is less than 64 bits (like 32, 16, 8), and the operand
+                    // sizes match, then we can represent it as an immediate in the instruction
+                    // without moving it to a register first.
+                    // IOW, 64 bit immediates must always be moved to a register
+                    // before comparisons, where other sizes may be encoded
+                    // directly in the instruction.
+                    if num_bits.is_some() && left.num_bits() == num_bits && num_bits.unwrap() < 64 {
+                        cmp(cb, left.into(), right.into());
+                    } else {
+                        let emitted = emit_64bit_immediate(cb, right);
+                        cmp(cb, left.into(), emitted);
+                    }
                 }
 
                 // Test and set flags
@@ -781,6 +797,30 @@ mod tests {
         assert_eq!(format!("{:x}", cb), "49bbffffffffffff00004c39d8");
     }
 
+    #[test]
+    fn test_emit_cmp_mem_16_bits_with_imm_16() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let shape_opnd = Opnd::mem(16, Opnd::Reg(RAX_REG), 6);
+
+        asm.cmp(shape_opnd, Opnd::UImm(0xF000));
+        asm.compile_with_num_regs(&mut cb, 0);
+
+        assert_eq!(format!("{:x}", cb), "6681780600f0");
+    }
+
+    #[test]
+    fn test_emit_cmp_mem_32_bits_with_imm_32() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let shape_opnd = Opnd::mem(32, Opnd::Reg(RAX_REG), 4);
+
+        asm.cmp(shape_opnd, Opnd::UImm(0xF000_0000));
+        asm.compile_with_num_regs(&mut cb, 0);
+
+        assert_eq!(format!("{:x}", cb), "817804000000f0");
+    }
+
     #[test]
     fn test_emit_or_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();