ZJIT: Introduce a better LIR printer (#14986)

8 files changed, 580 insertions, 53 deletions

diff --git a/zjit/src/asm/arm64/opnd.rs b/zjit/src/asm/arm64/opnd.rs
index 8246bea08e..667533ab93 100644
--- a/zjit/src/asm/arm64/opnd.rs
+++ b/zjit/src/asm/arm64/opnd.rs
@@ -1,4 +1,4 @@
-
+use std::fmt;
 
 /// This operand represents a register.
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
@@ -18,7 +18,7 @@ impl A64Reg {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub struct A64Mem
 {
     // Size in bits
@@ -42,7 +42,7 @@ impl A64Mem {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum A64Opnd
 {
     // Dummy operand
@@ -196,3 +196,75 @@ pub const W31: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 32, reg_no: 31 });
 // C argument registers
 pub const C_ARG_REGS: [A64Opnd; 4] = [X0, X1, X2, X3];
 pub const C_ARG_REGREGS: [A64Reg; 4] = [X0_REG, X1_REG, X2_REG, X3_REG];
+
+impl fmt::Display for A64Reg {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match A64Opnd::Reg(*self) {
+            X0  => write!(f, "x0"),
+            X1  => write!(f, "x1"),
+            X2  => write!(f, "x2"),
+            X3  => write!(f, "x3"),
+            X4  => write!(f, "x4"),
+            X5  => write!(f, "x5"),
+            X6  => write!(f, "x6"),
+            X7  => write!(f, "x7"),
+            X8  => write!(f, "x8"),
+            X9  => write!(f, "x9"),
+            X10 => write!(f, "x10"),
+            X11 => write!(f, "x11"),
+            X12 => write!(f, "x12"),
+            X13 => write!(f, "x13"),
+            X14 => write!(f, "x14"),
+            X15 => write!(f, "x15"),
+            X16 => write!(f, "x16"),
+            X17 => write!(f, "x17"),
+            X18 => write!(f, "x18"),
+            X19 => write!(f, "x19"),
+            X20 => write!(f, "x20"),
+            X21 => write!(f, "x21"),
+            X22 => write!(f, "x22"),
+            X23 => write!(f, "x23"),
+            X24 => write!(f, "x24"),
+            X25 => write!(f, "x25"),
+            X26 => write!(f, "x26"),
+            X27 => write!(f, "x27"),
+            X28 => write!(f, "x28"),
+            X29 => write!(f, "x29"),
+            X30 => write!(f, "x30"),
+            X31 => write!(f, "x31"),
+            W0  => write!(f, "w0"),
+            W1  => write!(f, "w1"),
+            W2  => write!(f, "w2"),
+            W3  => write!(f, "w3"),
+            W4  => write!(f, "w4"),
+            W5  => write!(f, "w5"),
+            W6  => write!(f, "w6"),
+            W7  => write!(f, "w7"),
+            W8  => write!(f, "w8"),
+            W9  => write!(f, "w9"),
+            W10 => write!(f, "w10"),
+            W11 => write!(f, "w11"),
+            W12 => write!(f, "w12"),
+            W13 => write!(f, "w13"),
+            W14 => write!(f, "w14"),
+            W15 => write!(f, "w15"),
+            W16 => write!(f, "w16"),
+            W17 => write!(f, "w17"),
+            W18 => write!(f, "w18"),
+            W19 => write!(f, "w19"),
+            W20 => write!(f, "w20"),
+            W21 => write!(f, "w21"),
+            W22 => write!(f, "w22"),
+            W23 => write!(f, "w23"),
+            W24 => write!(f, "w24"),
+            W25 => write!(f, "w25"),
+            W26 => write!(f, "w26"),
+            W27 => write!(f, "w27"),
+            W28 => write!(f, "w28"),
+            W29 => write!(f, "w29"),
+            W30 => write!(f, "w30"),
+            W31 => write!(f, "w31"),
+            _ => write!(f, "{self:?}"),
+        }
+    }
+}
diff --git a/zjit/src/asm/x86_64/mod.rs b/zjit/src/asm/x86_64/mod.rs
index 8077218c4a..9d3bf18dcd 100644
--- a/zjit/src/asm/x86_64/mod.rs
+++ b/zjit/src/asm/x86_64/mod.rs
@@ -47,7 +47,7 @@ pub struct X86Reg
     pub reg_no: u8,
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub struct X86Mem
 {
     // Size in bits
@@ -66,7 +66,7 @@ pub struct X86Mem
     pub disp: i32,
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum X86Opnd
 {
     // Dummy operand
@@ -1380,3 +1380,74 @@ pub fn xor(cb: &mut CodeBlock, opnd0: X86Opnd, opnd1: X86Opnd) {
         opnd1
     );
 }
+
+impl fmt::Display for X86Reg {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match X86Opnd::Reg(*self) {
+            RAX  => write!(f, "rax"),
+            RCX  => write!(f, "rcx"),
+            RDX  => write!(f, "rdx"),
+            RBX  => write!(f, "rbx"),
+            RSP  => write!(f, "rsp"),
+            RBP  => write!(f, "rbp"),
+            RSI  => write!(f, "rsi"),
+            RDI  => write!(f, "rdi"),
+            R8   => write!(f, "r8"),
+            R9   => write!(f, "r9"),
+            R10  => write!(f, "r10"),
+            R11  => write!(f, "r11"),
+            R12  => write!(f, "r12"),
+            R13  => write!(f, "r13"),
+            R14  => write!(f, "r14"),
+            R15  => write!(f, "r15"),
+            EAX  => write!(f, "eax"),
+            ECX  => write!(f, "ecx"),
+            EDX  => write!(f, "edx"),
+            EBX  => write!(f, "ebx"),
+            ESP  => write!(f, "esp"),
+            EBP  => write!(f, "ebp"),
+            ESI  => write!(f, "esi"),
+            EDI  => write!(f, "edi"),
+            R8D  => write!(f, "r8d"),
+            R9D  => write!(f, "r9d"),
+            R10D => write!(f, "r10d"),
+            R11D => write!(f, "r11d"),
+            R12D => write!(f, "r12d"),
+            R13D => write!(f, "r13d"),
+            R14D => write!(f, "r14d"),
+            R15D => write!(f, "r15d"),
+            AX   => write!(f, "ax"),
+            CX   => write!(f, "cx"),
+            DX   => write!(f, "dx"),
+            BX   => write!(f, "bx"),
+            BP   => write!(f, "bp"),
+            SI   => write!(f, "si"),
+            DI   => write!(f, "di"),
+            R8W  => write!(f, "r8w"),
+            R9W  => write!(f, "r9w"),
+            R10W => write!(f, "r10w"),
+            R11W => write!(f, "r11w"),
+            R12W => write!(f, "r12w"),
+            R13W => write!(f, "r13w"),
+            R14W => write!(f, "r14w"),
+            R15W => write!(f, "r15w"),
+            AL   => write!(f, "al"),
+            CL   => write!(f, "cl"),
+            DL   => write!(f, "dl"),
+            BL   => write!(f, "bl"),
+            SPL  => write!(f, "spl"),
+            BPL  => write!(f, "bpl"),
+            SIL  => write!(f, "sil"),
+            DIL  => write!(f, "dil"),
+            R8B  => write!(f, "r8b"),
+            R9B  => write!(f, "r9b"),
+            R10B => write!(f, "r10b"),
+            R11B => write!(f, "r11b"),
+            R12B => write!(f, "r12b"),
+            R13B => write!(f, "r13b"),
+            R14B => write!(f, "r14b"),
+            R15B => write!(f, "r15b"),
+            _ => write!(f, "{self:?}"),
+        }
+    }
+}
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
index 234f5ac059..d7f929a021 100644
--- a/zjit/src/backend/arm64/mod.rs
+++ b/zjit/src/backend/arm64/mod.rs
@@ -4,6 +4,7 @@ use crate::asm::{CodeBlock, Label};
 use crate::asm::arm64::*;
 use crate::cruby::*;
 use crate::backend::lir::*;
+use crate::options::asm_dump;
 use crate::stats::CompileError;
 use crate::virtualmem::CodePtr;
 use crate::cast::*;
@@ -11,6 +12,11 @@ use crate::cast::*;
 // Use the arm64 register type for this platform
 pub type Reg = A64Reg;
 
+/// Convert reg_no for MemBase::Reg into Reg, assuming it's a 64-bit register
+pub fn mem_base_reg(reg_no: u8) -> Reg {
+    Reg { num_bits: 64, reg_no }
+}
+
 // Callee-saved registers
 pub const CFP: Opnd = Opnd::Reg(X19_REG);
 pub const EC: Opnd = Opnd::Reg(X20_REG);
@@ -194,7 +200,7 @@ pub const ALLOC_REGS: &[Reg] = &[
 ];
 
 /// Special scratch registers for intermediate processing. They should be used only by
-/// [`Assembler::arm64_split_with_scratch_reg`] or [`Assembler::new_with_scratch_reg`].
+/// [`Assembler::arm64_scratch_split`] or [`Assembler::new_with_scratch_reg`].
 const SCRATCH0_OPND: Opnd = Opnd::Reg(X15_REG);
 const SCRATCH1_OPND: Opnd = Opnd::Reg(X17_REG);
 
@@ -683,7 +689,7 @@ impl Assembler {
     /// VRegs, most splits should happen in [`Self::arm64_split`]. However, some instructions
     /// need to be split with registers after `alloc_regs`, e.g. for `compile_side_exits`, so this
     /// splits them and uses scratch registers for it.
-    fn arm64_split_with_scratch_reg(self) -> Assembler {
+    fn arm64_scratch_split(self) -> Assembler {
         let mut asm = Assembler::new_with_asm(&self);
         asm.accept_scratch_reg = true;
         let mut iterator = self.insns.into_iter().enumerate().peekable();
@@ -971,6 +977,9 @@ impl Assembler {
         // For each instruction
         let mut insn_idx: usize = 0;
         while let Some(insn) = self.insns.get(insn_idx) {
+            // Dump Assembler with insn_idx if --zjit-dump-lir=panic is given
+            let _hook = AssemblerPanicHook::new(self, insn_idx);
+
             match insn {
                 Insn::Comment(text) => {
                     cb.add_comment(text);
@@ -1070,7 +1079,7 @@ impl Assembler {
                     }
                 },
                 Insn::Mul { left, right, out } => {
-                    // If the next instruction is JoMul with RShift created by arm64_split_with_scratch_reg
+                    // If the next instruction is JoMul with RShift created by arm64_scratch_split
                     match (self.insns.get(insn_idx + 1), self.insns.get(insn_idx + 2)) {
                         (Some(Insn::RShift { out: out_sign, opnd: out_opnd, shift: out_shift }), Some(Insn::JoMul(_))) => {
                             // Compute the high 64 bits
@@ -1081,7 +1090,7 @@ impl Assembler {
                             // so we do it after smulh
                             mul(cb, out.into(), left.into(), right.into());
 
-                            // Insert the shift instruction created by arm64_split_with_scratch_reg
+                            // Insert the shift instruction created by arm64_scratch_split
                             // to prepare the register that has the sign bit of the high 64 bits after mul.
                             asr(cb, out_sign.into(), out_opnd.into(), out_shift.into());
                             insn_idx += 1; // skip the next Insn::RShift
@@ -1382,12 +1391,12 @@ impl Assembler {
                     last_patch_pos = Some(cb.get_write_pos());
                 },
                 Insn::IncrCounter { mem, value } => {
-                    // Get the status register allocated by arm64_split_with_scratch_reg
+                    // Get the status register allocated by arm64_scratch_split
                     let Some(Insn::Cmp {
                         left: status_reg @ Opnd::Reg(_),
                         right: Opnd::UImm(_) | Opnd::Imm(_),
                     }) = self.insns.get(insn_idx + 1) else {
-                        panic!("arm64_split_with_scratch_reg should add Cmp after IncrCounter: {:?}", self.insns.get(insn_idx + 1));
+                        panic!("arm64_scratch_split should add Cmp after IncrCounter: {:?}", self.insns.get(insn_idx + 1));
                     };
 
                     // Attempt to increment a counter
@@ -1451,13 +1460,21 @@ impl Assembler {
     pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Result<(CodePtr, Vec<CodePtr>), CompileError> {
         // The backend is allowed to use scratch registers only if it has not accepted them so far.
         let use_scratch_reg = !self.accept_scratch_reg;
+        asm_dump!(self, init);
 
         let asm = self.arm64_split();
+        asm_dump!(asm, split);
+
         let mut asm = asm.alloc_regs(regs)?;
+        asm_dump!(asm, alloc_regs);
+
         // We put compile_side_exits after alloc_regs to avoid extending live ranges for VRegs spilled on side exits.
         asm.compile_side_exits();
+        asm_dump!(asm, compile_side_exits);
+
         if use_scratch_reg {
-            asm = asm.arm64_split_with_scratch_reg();
+            asm = asm.arm64_scratch_split();
+            asm_dump!(asm, scratch_split);
         }
 
         // Create label instances in the code block
@@ -1529,6 +1546,42 @@ mod tests {
     }
 
     #[test]
+    fn test_lir_string() {
+        use crate::hir::SideExitReason;
+
+        let mut asm = Assembler::new();
+        asm.stack_base_idx = 1;
+
+        let label = asm.new_label("bb0");
+        asm.write_label(label.clone());
+        asm.push_insn(Insn::Comment("bb0(): foo@/tmp/a.rb:1".into()));
+        asm.frame_setup(JIT_PRESERVED_REGS);
+
+        let val64 = asm.add(CFP, Opnd::UImm(64));
+        asm.store(Opnd::mem(64, SP, 0x10), val64);
+        let side_exit = Target::SideExit { reason: SideExitReason::Interrupt, pc: 0 as _, stack: vec![], locals: vec![], label: None };
+        asm.push_insn(Insn::Joz(val64, side_exit));
+
+        let val32 = asm.sub(Opnd::Value(Qtrue), Opnd::Imm(1));
+        asm.store(Opnd::mem(64, EC, 0x10).with_num_bits(32), val32.with_num_bits(32));
+        asm.cret(val64);
+
+        asm.frame_teardown(JIT_PRESERVED_REGS);
+        assert_disasm_snapshot!(lir_string(&mut asm), @r"
+        bb0:
+          # bb0(): foo@/tmp/a.rb:1
+          FrameSetup 1, x19, x21, x20
+          v0 = Add x19, 0x40
+          Store [x21 + 0x10], v0
+          Joz Exit(Interrupt), v0
+          v1 = Sub Value(0x14), Imm(1)
+          Store Mem32[x20 + 0x10], VReg32(v1)
+          CRet v0
+          FrameTeardown x19, x21, x20
+        ");
+    }
+
+    #[test]
     fn test_mul_with_immediate() {
         let (mut asm, mut cb) = setup_asm();
 
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index 72ebeaf11b..e5707beb51 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -1,10 +1,13 @@
 use std::collections::HashMap;
 use std::fmt;
 use std::mem::take;
+use std::panic;
+use std::rc::Rc;
+use std::sync::Arc;
 use crate::codegen::local_size_and_idx_to_ep_offset;
 use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, vm_stack_canary};
 use crate::hir::SideExitReason;
-use crate::options::{debug, get_option, TraceExits};
+use crate::options::{TraceExits, debug, dump_lir_option, get_option};
 use crate::cruby::VALUE;
 use crate::stats::{exit_counter_ptr, exit_counter_ptr_for_opcode, side_exit_counter, CompileError};
 use crate::virtualmem::CodePtr;
@@ -12,6 +15,7 @@ use crate::asm::{CodeBlock, Label};
 use crate::state::rb_zjit_record_exit_stack;
 
 pub use crate::backend::current::{
+    mem_base_reg,
     Reg,
     EC, CFP, SP,
     NATIVE_STACK_PTR, NATIVE_BASE_PTR,
@@ -42,6 +46,28 @@ pub struct Mem
     pub num_bits: u8,
 }
 
+impl fmt::Display for Mem {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if self.num_bits != 64 {
+            write!(f, "Mem{}", self.num_bits)?;
+        }
+        write!(f, "[")?;
+        match self.base {
+            MemBase::Reg(reg_no) => write!(f, "{}", mem_base_reg(reg_no))?,
+            MemBase::VReg(idx) => write!(f, "v{idx}")?,
+        }
+        if self.disp != 0 {
+            let sign = if self.disp > 0 { '+' } else { '-' };
+            write!(f, " {sign} ")?;
+            if self.disp.abs() >= 10 {
+                write!(f, "0x")?;
+            }
+            write!(f, "{:x}", self.disp.abs())?;
+        }
+        write!(f, "]")
+    }
+}
+
 impl fmt::Debug for Mem {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         write!(fmt, "Mem{}[{:?}", self.num_bits, self.base)?;
@@ -73,6 +99,25 @@ pub enum Opnd
     Reg(Reg),           // Machine register
 }
 
+impl fmt::Display for Opnd {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use Opnd::*;
+        match self {
+            None => write!(f, "None"),
+            Value(VALUE(value)) if *value < 10 => write!(f, "Value({value:x})"),
+            Value(VALUE(value)) => write!(f, "Value(0x{value:x})"),
+            VReg { idx, num_bits } if *num_bits == 64 => write!(f, "v{idx}"),
+            VReg { idx, num_bits } => write!(f, "VReg{num_bits}(v{idx})"),
+            Imm(value) if value.abs() < 10 => write!(f, "Imm({value:x})"),
+            Imm(value) => write!(f, "Imm(0x{value:x})"),
+            UImm(value) if *value < 10 => write!(f, "{value:x}"),
+            UImm(value) => write!(f, "0x{value:x}"),
+            Mem(mem) => write!(f, "{mem}"),
+            Reg(reg) => write!(f, "{reg}"),
+        }
+    }
+}
+
 impl fmt::Debug for Opnd {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         use Opnd::*;
@@ -291,9 +336,10 @@ impl From<CodePtr> for Target {
     }
 }
 
-type PosMarkerFn = Box<dyn Fn(CodePtr, &CodeBlock)>;
+type PosMarkerFn = Rc<dyn Fn(CodePtr, &CodeBlock)>;
 
 /// ZJIT Low-level IR instruction
+#[derive(Clone)]
 pub enum Insn {
     /// Add two operands together, and return the result as a new operand.
     Add { left: Opnd, right: Opnd, out: Opnd },
@@ -793,8 +839,6 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::CPop { .. } |
             Insn::CPopAll |
             Insn::CPushAll |
-            Insn::FrameSetup { .. } |
-            Insn::FrameTeardown { .. } |
             Insn::PadPatchPoint |
             Insn::PosMarker(_) => None,
 
@@ -860,14 +904,29 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
                 }
             },
             Insn::ParallelMov { moves } => {
-                if self.idx < moves.len() {
-                    let opnd = &moves[self.idx].1;
+                if self.idx < moves.len() * 2 {
+                    let move_idx = self.idx / 2;
+                    let opnd = if self.idx % 2 == 0 {
+                        &moves[move_idx].0
+                    } else {
+                        &moves[move_idx].1
+                    };
                     self.idx += 1;
                     Some(opnd)
                 } else {
                     None
                 }
             },
+            Insn::FrameSetup { preserved, .. } |
+            Insn::FrameTeardown { preserved } => {
+                if self.idx < preserved.len() {
+                    let opnd = &preserved[self.idx];
+                    self.idx += 1;
+                    Some(opnd)
+                } else {
+                    None
+                }
+            }
         }
     }
 }
@@ -1016,8 +1075,13 @@ impl<'a> InsnOpndMutIterator<'a> {
                 }
             },
             Insn::ParallelMov { moves } => {
-                if self.idx < moves.len() {
-                    let opnd = &mut moves[self.idx].1;
+                if self.idx < moves.len() * 2 {
+                    let move_idx = self.idx / 2;
+                    let opnd = if self.idx % 2 == 0 {
+                        &mut moves[move_idx].0
+                    } else {
+                        &mut moves[move_idx].1
+                    };
                     self.idx += 1;
                     Some(opnd)
                 } else {
@@ -1166,6 +1230,7 @@ const ASSEMBLER_INSNS_CAPACITY: usize = 256;
 
 /// Object into which we assemble instructions to be
 /// optimized and lowered
+#[derive(Clone)]
 pub struct Assembler {
     pub(super) insns: Vec<Insn>,
 
@@ -1714,6 +1779,93 @@ impl Assembler
     }
 }
 
+const BOLD_BEGIN: &str = "\x1b[1m";
+const BOLD_END: &str = "\x1b[22m";
+
+/// Return a result of fmt::Display for Assembler without escape sequence
+pub fn lir_string(asm: &Assembler) -> String {
+    format!("{asm}").replace(BOLD_BEGIN, "").replace(BOLD_END, "")
+}
+
+impl fmt::Display for Assembler {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        for insn in self.insns.iter() {
+            match insn {
+                Insn::Comment(comment) => {
+                    writeln!(f, "    {BOLD_BEGIN}# {comment}{BOLD_END}")?;
+                }
+                Insn::Label(target) => {
+                    let &Target::Label(Label(label_idx)) = target else {
+                        panic!("unexpected target for Insn::Label: {target:?}");
+                    };
+                    writeln!(f, "  {}:", self.label_names[label_idx])?;
+                }
+                _ => {
+                    write!(f, "    ")?;
+
+                    // Print output operand if any
+                    if let Some(out) = insn.out_opnd() {
+                        write!(f, "{out} = ")?;
+                    }
+
+                    write!(f, "{}", insn.op())?;
+
+                    // Show slot_count for FrameSetup
+                    if let Insn::FrameSetup { slot_count, preserved } = insn {
+                        write!(f, " {slot_count}")?;
+                        if !preserved.is_empty() {
+                            write!(f, ",")?;
+                        }
+                    }
+
+                    // Print target
+                    if let Some(target) = insn.target() {
+                        match target {
+                            Target::CodePtr(code_ptr) => write!(f, " {code_ptr:?}")?,
+                            Target::Label(Label(label_idx)) => write!(f, " {}", self.label_names[*label_idx])?,
+                            Target::SideExit { reason, .. } => write!(f, " Exit({reason})")?,
+                        }
+                    }
+
+                    // Print list of operands
+                    if let Some(Target::SideExit { .. }) = insn.target() {
+                        // If the instruction has a SideExit, avoid using opnd_iter(), which has stack/locals.
+                        // Here, only handle instructions that have both Opnd and Target.
+                        match insn {
+                            Insn::Joz(opnd, _) |
+                            Insn::Jonz(opnd, _) |
+                            Insn::LeaJumpTarget { out: opnd, target: _ } => {
+                                write!(f, ", {opnd}")?;
+                            }
+                            _ => {}
+                        }
+                    } else if let Insn::ParallelMov { moves } = insn {
+                        // Print operands with a special syntax for ParallelMov
+                        let mut moves_iter = moves.iter();
+                        if let Some((first_dst, first_src)) = moves_iter.next() {
+                            write!(f, " {first_dst} <- {first_src}")?;
+                        }
+                        for (dst, src) in moves_iter {
+                            write!(f, ", {dst} <- {src}")?;
+                        }
+                    } else {
+                        let mut opnd_iter = insn.opnd_iter();
+                        if let Some(first_opnd) = opnd_iter.next() {
+                            write!(f, " {first_opnd}")?;
+                        }
+                        for opnd in opnd_iter {
+                            write!(f, ", {opnd}")?;
+                        }
+                    }
+
+                    write!(f, "\n")?;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
 impl fmt::Debug for Assembler {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         writeln!(fmt, "Assembler")?;
@@ -1777,8 +1929,8 @@ impl Assembler {
         self.push_insn(Insn::CCall {
             fptr,
             opnds,
-            start_marker: Some(Box::new(start_marker)),
-            end_marker: Some(Box::new(end_marker)),
+            start_marker: Some(Rc::new(start_marker)),
+            end_marker: Some(Rc::new(end_marker)),
             out,
         });
         out
@@ -2028,7 +2180,7 @@ impl Assembler {
     }
 
     pub fn pos_marker(&mut self, marker_fn: impl Fn(CodePtr, &CodeBlock) + 'static) {
-        self.push_insn(Insn::PosMarker(Box::new(marker_fn)));
+        self.push_insn(Insn::PosMarker(Rc::new(marker_fn)));
     }
 
     #[must_use]
@@ -2092,7 +2244,7 @@ impl Assembler {
 /// when not dumping disassembly.
 macro_rules! asm_comment {
     ($asm:expr, $($fmt:tt)*) => {
-        if $crate::options::get_option!(dump_disasm) {
+        if $crate::options::get_option!(dump_disasm) || $crate::options::get_option!(dump_lir).is_some() {
             $asm.push_insn(crate::backend::lir::Insn::Comment(format!($($fmt)*)));
         }
     };
@@ -2108,6 +2260,64 @@ macro_rules! asm_ccall {
 }
 pub(crate) use asm_ccall;
 
+// Allow moving Assembler to panic hooks. Since we take the VM lock on compilation,
+// no other threads should reference the same Assembler instance.
+unsafe impl Send for Insn {}
+unsafe impl Sync for Insn {}
+
+/// Dump Assembler with insn_idx on panic. Restore the original panic hook on drop.
+pub struct AssemblerPanicHook {
+    /// Original panic hook before AssemblerPanicHook is installed.
+    prev_hook: Box<dyn Fn(&panic::PanicHookInfo<'_>) + Sync + Send + 'static>,
+}
+
+impl AssemblerPanicHook {
+    pub fn new(asm: &Assembler, insn_idx: usize) -> Option<Arc<Self>> {
+        // Install a panic hook if --zjit-dump-lir=panic is specified.
+        if dump_lir_option!(panic) {
+            // Wrap prev_hook with Arc to share it among the new hook and Self to be dropped.
+            let prev_hook = panic::take_hook();
+            let panic_hook_ref = Arc::new(Self { prev_hook });
+            let weak_hook = Arc::downgrade(&panic_hook_ref);
+
+            // Install a new hook to dump Assembler with insn_idx
+            let asm = asm.clone();
+            panic::set_hook(Box::new(move |panic_info| {
+                if let Some(panic_hook) = weak_hook.upgrade() {
+                    // Dump Assembler, highlighting the insn_idx line
+                    Self::dump_asm(&asm, insn_idx);
+
+                    // Call the previous panic hook
+                    (panic_hook.prev_hook)(panic_info);
+                }
+            }));
+
+            Some(panic_hook_ref)
+        } else {
+            None
+        }
+    }
+
+    /// Dump Assembler, highlighting the insn_idx line
+    fn dump_asm(asm: &Assembler, insn_idx: usize) {
+        println!("Failed to compile LIR at insn_idx={insn_idx}:");
+        for (idx, line) in lir_string(asm).split('\n').enumerate() {
+            if idx == insn_idx && line.starts_with("  ") {
+                println!("{BOLD_BEGIN}=>{}{BOLD_END}", &line["  ".len()..]);
+            } else {
+                println!("{line}");
+            }
+        }
+    }
+}
+
+impl Drop for AssemblerPanicHook {
+    fn drop(&mut self) {
+        // Restore the original hook
+        panic::set_hook(std::mem::replace(&mut self.prev_hook, Box::new(|_| {})));
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/zjit/src/backend/tests.rs b/zjit/src/backend/tests.rs
index 52547cb31f..6e62b3068d 100644
--- a/zjit/src/backend/tests.rs
+++ b/zjit/src/backend/tests.rs
@@ -62,10 +62,8 @@ fn test_alloc_regs() {
 }
 
 fn setup_asm() -> (Assembler, CodeBlock) {
-    (
-        Assembler::new(),
-        CodeBlock::new_dummy()
-    )
+    rb_zjit_prepare_options(); // for get_option! on asm.compile
+    (Assembler::new(), CodeBlock::new_dummy())
 }
 
 // Test full codegen pipeline
@@ -87,7 +85,6 @@ fn test_compile()
 fn test_mov_mem2mem()
 {
     let (mut asm, mut cb) = setup_asm();
-    rb_zjit_prepare_options(); // for asm_comment
 
     asm_comment!(asm, "check that comments work too");
     asm.mov(Opnd::mem(64, SP, 0), Opnd::mem(64, SP, 8));
@@ -181,7 +178,6 @@ fn test_c_call()
     }
 
     let (mut asm, mut cb) = setup_asm();
-    rb_zjit_prepare_options(); // for asm.compile
 
     let ret_val = asm.ccall(
         dummy_c_fun as *const u8,
@@ -196,11 +192,10 @@ fn test_c_call()
 
 #[test]
 fn test_alloc_ccall_regs() {
-    let mut asm = Assembler::new();
+    let (mut asm, mut cb) = setup_asm();
     let out1 = asm.ccall(std::ptr::null::<u8>(), vec![]);
     let out2 = asm.ccall(std::ptr::null::<u8>(), vec![out1]);
     asm.mov(EC, out2);
-    let mut cb = CodeBlock::new_dummy();
     asm.compile_with_regs(&mut cb, Assembler::get_alloc_regs()).unwrap();
 }
 
diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
index 2590ceaf7d..a2e27028e6 100644
--- a/zjit/src/backend/x86_64/mod.rs
+++ b/zjit/src/backend/x86_64/mod.rs
@@ -7,10 +7,16 @@ use crate::virtualmem::CodePtr;
 use crate::cruby::*;
 use crate::backend::lir::*;
 use crate::cast::*;
+use crate::options::asm_dump;
 
 // Use the x86 register type for this platform
 pub type Reg = X86Reg;
 
+/// Convert reg_no for MemBase::Reg into Reg, assuming it's a 64-bit GP register
+pub fn mem_base_reg(reg_no: u8) -> Reg {
+    Reg { num_bits: 64, reg_type: RegType::GP, reg_no }
+}
+
 // Callee-saved registers
 pub const CFP: Opnd = Opnd::Reg(R13_REG);
 pub const EC: Opnd = Opnd::Reg(R12_REG);
@@ -96,7 +102,7 @@ pub const ALLOC_REGS: &[Reg] = &[
 ];
 
 /// Special scratch register for intermediate processing. It should be used only by
-/// [`Assembler::x86_split_with_scratch_reg`] or [`Assembler::new_with_scratch_reg`].
+/// [`Assembler::x86_scratch_split`] or [`Assembler::new_with_scratch_reg`].
 const SCRATCH0_OPND: Opnd = Opnd::Reg(R11_REG);
 
 impl Assembler {
@@ -384,7 +390,7 @@ impl Assembler {
     /// for VRegs, most splits should happen in [`Self::x86_split`]. However, some instructions
     /// need to be split with registers after `alloc_regs`, e.g. for `compile_side_exits`, so
     /// this splits them and uses scratch registers for it.
-    pub fn x86_split_with_scratch_reg(self) -> Assembler {
+    pub fn x86_scratch_split(self) -> Assembler {
         /// For some instructions, we want to be able to lower a 64-bit operand
         /// without requiring more registers to be available in the register
         /// allocator. So we just use the SCRATCH0_OPND register temporarily to hold
@@ -490,7 +496,7 @@ impl Assembler {
                 // Handle various operand combinations for spills on compile_side_exits.
                 &mut Insn::Store { dest, src } => {
                     let Opnd::Mem(Mem { num_bits, .. }) = dest else {
-                        panic!("Unexpected Insn::Store destination in x86_split_with_scratch_reg: {dest:?}");
+                        panic!("Unexpected Insn::Store destination in x86_scratch_split: {dest:?}");
                     };
 
                     let src = match src {
@@ -525,7 +531,7 @@ impl Assembler {
                             asm.load_into(SCRATCH0_OPND, src);
                             SCRATCH0_OPND
                         }
-                        src @ (Opnd::None | Opnd::VReg { .. }) => panic!("Unexpected source operand during x86_split_with_scratch_reg: {src:?}"),
+                        src @ (Opnd::None | Opnd::VReg { .. }) => panic!("Unexpected source operand during x86_scratch_split: {src:?}"),
                     };
                     asm.store(dest, src);
                 }
@@ -587,6 +593,9 @@ impl Assembler {
         // For each instruction
         let mut insn_idx: usize = 0;
         while let Some(insn) = self.insns.get(insn_idx) {
+            // Dump Assembler with insn_idx if --zjit-dump-lir=panic is given
+            let _hook = AssemblerPanicHook::new(self, insn_idx);
+
             match insn {
                 Insn::Comment(text) => {
                     cb.add_comment(text);
@@ -949,13 +958,21 @@ impl Assembler {
     pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Result<(CodePtr, Vec<CodePtr>), CompileError> {
         // The backend is allowed to use scratch registers only if it has not accepted them so far.
         let use_scratch_regs = !self.accept_scratch_reg;
+        asm_dump!(self, init);
 
         let asm = self.x86_split();
+        asm_dump!(asm, split);
+
         let mut asm = asm.alloc_regs(regs)?;
+        asm_dump!(asm, alloc_regs);
+
         // We put compile_side_exits after alloc_regs to avoid extending live ranges for VRegs spilled on side exits.
         asm.compile_side_exits();
+        asm_dump!(asm, compile_side_exits);
+
         if use_scratch_regs {
-            asm = asm.x86_split_with_scratch_reg();
+            asm = asm.x86_scratch_split();
+            asm_dump!(asm, scratch_split);
         }
 
         // Create label instances in the code block
@@ -981,13 +998,54 @@ impl Assembler {
 mod tests {
     use insta::assert_snapshot;
     use crate::assert_disasm_snapshot;
+    use crate::options::rb_zjit_prepare_options;
     use super::*;
 
+    const BOLD_BEGIN: &str = "\x1b[1m";
+    const BOLD_END: &str = "\x1b[22m";
+
     fn setup_asm() -> (Assembler, CodeBlock) {
+        rb_zjit_prepare_options(); // for get_option! on asm.compile
         (Assembler::new(), CodeBlock::new_dummy())
     }
 
     #[test]
+    fn test_lir_string() {
+        use crate::hir::SideExitReason;
+
+        let mut asm = Assembler::new();
+        asm.stack_base_idx = 1;
+
+        let label = asm.new_label("bb0");
+        asm.write_label(label.clone());
+        asm.push_insn(Insn::Comment("bb0(): foo@/tmp/a.rb:1".into()));
+        asm.frame_setup(JIT_PRESERVED_REGS);
+
+        let val64 = asm.add(CFP, Opnd::UImm(64));
+        asm.store(Opnd::mem(64, SP, 0x10), val64);
+        let side_exit = Target::SideExit { reason: SideExitReason::Interrupt, pc: 0 as _, stack: vec![], locals: vec![], label: None };
+        asm.push_insn(Insn::Joz(val64, side_exit));
+
+        let val32 = asm.sub(Opnd::Value(Qtrue), Opnd::Imm(1));
+        asm.store(Opnd::mem(64, EC, 0x10).with_num_bits(32), val32.with_num_bits(32));
+        asm.cret(val64);
+
+        asm.frame_teardown(JIT_PRESERVED_REGS);
+        assert_disasm_snapshot!(lir_string(&mut asm), @r"
+        bb0:
+          # bb0(): foo@/tmp/a.rb:1
+          FrameSetup 1, r13, rbx, r12
+          v0 = Add r13, 0x40
+          Store [rbx + 0x10], v0
+          Joz Exit(Interrupt), v0
+          v1 = Sub Value(0x14), Imm(1)
+          Store Mem32[r12 + 0x10], VReg32(v1)
+          CRet v0
+          FrameTeardown r13, rbx, r12
+        ");
+    }
+
+    #[test]
     #[ignore]
     fn test_emit_add_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
@@ -1596,7 +1654,7 @@ mod tests {
         assert!(imitation_heap_value.heap_object_p());
         asm.store(Opnd::mem(VALUE_BITS, SP, 0), imitation_heap_value.into());
 
-        asm = asm.x86_split_with_scratch_reg();
+        asm = asm.x86_scratch_split();
         let gc_offsets = asm.x86_emit(&mut cb).unwrap();
         assert_eq!(1, gc_offsets.len(), "VALUE source operand should be reported as gc offset");
 
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 2ac7ca348a..db9d6a14e3 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -184,10 +184,6 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function_ptr: CodePtr) -> Result
     asm.frame_teardown(lir::JIT_PRESERVED_REGS);
     asm.cret(C_RET_OPND);
 
-    if get_option!(dump_lir) {
-        println!("LIR:\nJIT entry for {}:\n{:?}", iseq_name(iseq), asm);
-    }
-
     let (code_ptr, gc_offsets) = asm.compile(cb)?;
     assert!(gc_offsets.is_empty());
     if get_option!(perf) {
@@ -256,6 +252,10 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Resul
     // Compile each basic block
     let reverse_post_order = function.rpo();
     for &block_id in reverse_post_order.iter() {
+        // Write a label to jump to the basic block
+        let label = jit.get_label(&mut asm, block_id);
+        asm.write_label(label);
+
         let block = function.block(block_id);
         asm_comment!(
             asm, "{block_id}({}): {}",
@@ -263,10 +263,6 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Resul
             iseq_get_location(iseq, block.insn_idx),
         );
 
-        // Write a label to jump to the basic block
-        let label = jit.get_label(&mut asm, block_id);
-        asm.write_label(label);
-
         // Compile all parameters
         for (idx, &insn_id) in block.params().enumerate() {
             match function.find(insn_id) {
@@ -293,10 +289,6 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Resul
         asm.pad_patch_point();
     }
 
-    if get_option!(dump_lir) {
-        println!("LIR:\nfn {}:\n{:?}", iseq_name(iseq), asm);
-    }
-
     // Generate code if everything can be compiled
     let result = asm.compile(cb);
     if let Ok((start_ptr, _)) = result {
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
index f4a52e1ccd..44f4dbc7a4 100644
--- a/zjit/src/options.rs
+++ b/zjit/src/options.rs
@@ -67,7 +67,7 @@ pub struct Options {
     pub dump_hir_graphviz: Option<std::path::PathBuf>,
 
     /// Dump low-level IR
-    pub dump_lir: bool,
+    pub dump_lir: Option<HashSet<DumpLIR>>,
 
     /// Dump all compiled machine code.
     pub dump_disasm: bool,
@@ -101,7 +101,7 @@ impl Default for Options {
             dump_hir_init: None,
             dump_hir_opt: None,
             dump_hir_graphviz: None,
-            dump_lir: false,
+            dump_lir: None,
             dump_disasm: false,
             trace_side_exits: None,
             trace_side_exits_sample_interval: 0,
@@ -150,6 +150,55 @@ pub enum DumpHIR {
     Debug,
 }
 
+/// --zjit-dump-lir values. Using snake_case to stringify the exact filter value.
+#[allow(non_camel_case_types)]
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub enum DumpLIR {
+    /// Dump the initial LIR
+    init,
+    /// Dump LIR after {arch}_split
+    split,
+    /// Dump LIR after alloc_regs
+    alloc_regs,
+    /// Dump LIR after compile_side_exits
+    compile_side_exits,
+    /// Dump LIR after {arch}_scratch_split
+    scratch_split,
+    /// Dump LIR at panic on {arch}_emit
+    panic,
+}
+
+/// All compiler stages for --zjit-dump-lir=all. This does NOT include DumpLIR::panic.
+const DUMP_LIR_ALL: &[DumpLIR] = &[
+    DumpLIR::init,
+    DumpLIR::split,
+    DumpLIR::alloc_regs,
+    DumpLIR::compile_side_exits,
+    DumpLIR::scratch_split,
+];
+
+/// Macro to dump LIR if --zjit-dump-lir is specified
+macro_rules! asm_dump {
+    ($asm:expr, $target:ident) => {
+        if crate::options::dump_lir_option!($target) {
+            println!("LIR {}:\n{}", stringify!($target), $asm);
+        }
+    };
+}
+pub(crate) use asm_dump;
+
+/// Macro to check if a particular dump_lir option is enabled
+macro_rules! dump_lir_option {
+    ($target:ident) => {
+        if let Some(crate::options::Options { dump_lir: Some(dump_lirs), .. }) = unsafe { crate::options::OPTIONS.as_ref() } {
+            dump_lirs.contains(&crate::options::DumpLIR::$target)
+        } else {
+            false
+        }
+    };
+}
+pub(crate) use dump_lir_option;
+
 /// Macro to get an option value by name
 macro_rules! get_option {
     // Unsafe is ok here because options are initialized
@@ -301,7 +350,34 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             options.dump_hir_graphviz = Some(opt_val);
         }
 
-        ("dump-lir", "") => options.dump_lir = true,
+        ("dump-lir", "") => options.dump_lir = Some(HashSet::from([DumpLIR::init])),
+        ("dump-lir", filters) => {
+            let mut dump_lirs = HashSet::new();
+            for filter in filters.split(',') {
+                let dump_lir = match filter {
+                    "all" => {
+                        for &dump_lir in DUMP_LIR_ALL {
+                            dump_lirs.insert(dump_lir);
+                        }
+                        continue;
+                    }
+                    "init" => DumpLIR::init,
+                    "split" => DumpLIR::split,
+                    "alloc_regs" => DumpLIR::alloc_regs,
+                    "compile_side_exits" => DumpLIR::compile_side_exits,
+                    "scratch_split" => DumpLIR::scratch_split,
+                    "panic" => DumpLIR::panic,
+                    _ => {
+                        let valid_options = DUMP_LIR_ALL.iter().map(|opt| format!("{opt:?}")).collect::<Vec<_>>().join(", ");
+                        eprintln!("invalid --zjit-dump-lir option: '{filter}'");
+                        eprintln!("valid --zjit-dump-lir options: all, {}, panic", valid_options);
+                        return None;
+                    }
+                };
+                dump_lirs.insert(dump_lir);
+            }
+            options.dump_lir = Some(dump_lirs);
+        }
 
         ("dump-disasm", "") => options.dump_disasm = true,