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>

2 files changed, 20 insertions, 4 deletions

diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index 30fe4072b4..1a0cb9cae7 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/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::Mem(mem) => assert!(mem.num_bits == opnd_size),
+        X86Opnd::Reg(reg) => assert_eq!(reg.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!()
diff --git a/yjit/src/asm/x86_64/tests.rs b/yjit/src/asm/x86_64/tests.rs
index c6d49e084d..1cd005747d 100644
--- a/yjit/src/asm/x86_64/tests.rs
+++ b/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]
@@ -358,6 +359,14 @@ fn test_sub() {
 }
 
 #[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));
     check_bytes("6685c0", |cb| test(cb, AX, AX));