YJIT: fix 32 and 16 bit register store (#6840)

* Fix 32 and 16 bit register store in YJIT

Co-Authored-By: Takashi Kokubun <takashikkbn@gmail.com>

* Remove an unnecessary diff

* Reuse an rm_num_bits result

* Use u16::MAX instead

* Update the link

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

* Just use sturh for 16 bits

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

3 files changed, 45 insertions, 2 deletions

diff --git a/yjit/src/asm/arm64/inst/load_store.rs b/yjit/src/asm/arm64/inst/load_store.rs
index d08398380a..b5c8a3c294 100644
--- a/yjit/src/asm/arm64/inst/load_store.rs
+++ b/yjit/src/asm/arm64/inst/load_store.rs
@@ -118,6 +118,12 @@ impl LoadStore {
     pub fn stur(rt: u8, rn: u8, imm9: i16, num_bits: u8) -> Self {
         Self { rt, rn, idx: Index::None, imm9, opc: Opc::STR, size: num_bits.into() }
     }
+
+    /// STURH (store register, halfword, unscaled)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/STURH--Store-Register-Halfword--unscaled--?lang=en
+    pub fn sturh(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, idx: Index::None, imm9, opc: Opc::STR, size: Size::Size16 }
+    }
 }
 
 /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Loads-and-Stores?lang=en
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index d463613943..04660e580e 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -903,7 +903,7 @@ pub fn strh_post(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
 pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
         (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
-            assert!(rt.num_bits == rn.num_bits, "Expected registers to be the same size");
+            assert!(rn.num_bits == 32 || rn.num_bits == 64);
             assert!(mem_disp_fits_bits(rn.disp), "Expected displacement to be 9 bits or less");
 
             LoadStore::stur(rt.reg_no, rn.base_reg_no, rn.disp as i16, rt.num_bits).into()
@@ -914,6 +914,21 @@ pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// STURH - store a value in a register at a memory address
+pub fn sturh(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert!(rn.num_bits == 16);
+            assert!(mem_disp_fits_bits(rn.disp), "Expected displacement to be 9 bits or less");
+
+            LoadStore::sturh(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to stur instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// SUB - subtract rm from rn, put the result in rd, don't update flags
 pub fn sub(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
     let bytes: [u8; 4] = match (rd, rn, rm) {
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index a34b844c8f..eb096ce677 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -821,7 +821,11 @@ impl Assembler
                     // the Arm64 assembler works, the register that is going to
                     // be stored is first and the address is second. However in
                     // our IR we have the address first and the register second.
-                    stur(cb, src.into(), dest.into());
+                    match dest.rm_num_bits() {
+                        64 | 32 => stur(cb, src.into(), dest.into()),
+                        16 => sturh(cb, src.into(), dest.into()),
+                        num_bits => panic!("unexpected dest num_bits: {} (src: {:#?}, dest: {:#?})", num_bits, src, dest),
+                    }
                 },
                 Insn::Load { opnd, out } |
                 Insn::LoadInto { opnd, dest: out } => {
@@ -1379,6 +1383,24 @@ mod tests {
     }
 
     #[test]
+    fn test_16_bit_register_store_some_number() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let shape_opnd = Opnd::mem(16, Opnd::Reg(X0_REG), 0);
+        asm.store(shape_opnd, Opnd::UImm(4097));
+        asm.compile_with_num_regs(&mut cb, 2);
+    }
+
+    #[test]
+    fn test_32_bit_register_store_some_number() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let shape_opnd = Opnd::mem(32, Opnd::Reg(X0_REG), 6);
+        asm.store(shape_opnd, Opnd::UImm(4097));
+        asm.compile_with_num_regs(&mut cb, 2);
+    }
+
+    #[test]
     fn test_emit_xor() {
         let (mut asm, mut cb) = setup_asm();