Add --yjit-dump-disasm to dump every compiled code (https://github.com/Shopify/ruby/pull/430)

* Add --yjit-dump-disasm to dump every compiled code

* Just use get_option

* Carve out disasm_from_addr

* Avoid push_str with format!

* Share the logic through asm.compile

* This seems to negatively impact the compilation speed

yjit/src/asm/mod.rs

@@ -57,6 +57,10 @@ pub struct CodeBlock {
     #[cfg(feature = "asm_comments")]
     asm_comments: BTreeMap<usize, Vec<String>>,
 
+    // True for OutlinedCb
+    #[cfg(feature = "disasm")]
+    pub outlined: bool,
+
     // Set if the CodeBlock is unable to output some instructions,
     // for example, when there is not enough space or when a jump
     // target is too far away.
@@ -65,7 +69,7 @@ pub struct CodeBlock {
 
 impl CodeBlock {
     /// Make a new CodeBlock
-    pub fn new(mem_block: VirtualMem) -> Self {
+    pub fn new(mem_block: VirtualMem, outlined: bool) -> Self {
         Self {
             mem_size: mem_block.virtual_region_size(),
             mem_block,
@@ -75,6 +79,8 @@ impl CodeBlock {
             label_refs: Vec::new(),
             #[cfg(feature = "asm_comments")]
             asm_comments: BTreeMap::new(),
+            #[cfg(feature = "disasm")]
+            outlined,
             dropped_bytes: false,
         }
     }
@@ -282,7 +288,7 @@ impl CodeBlock {
         let mem_start: *const u8 = alloc.mem_start();
         let virt_mem = VirtualMem::new(alloc, 1, mem_start as *mut u8, mem_size);
 
-        Self::new(virt_mem)
+        Self::new(virt_mem, false)
     }
 }
 

yjit/src/backend/ir.rs

@@ -10,6 +10,7 @@ use crate::cruby::{VALUE};
 use crate::virtualmem::{CodePtr};
 use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits};
 use crate::core::{Context, Type, TempMapping};
+use crate::options::*;
 
 #[cfg(target_arch = "x86_64")]
 use crate::backend::x86_64::*;
@@ -1075,11 +1076,25 @@ impl Assembler
     ///       compiling multiple blocks at a time?
     pub fn compile(self, cb: &mut CodeBlock) -> Vec<u32>
     {
+        #[cfg(feature = "disasm")]
+        let start_addr = cb.get_write_ptr().raw_ptr();
+
         let alloc_regs = Self::get_alloc_regs();
-        self.compile_with_regs(cb, alloc_regs)
+        let gc_offsets = self.compile_with_regs(cb, alloc_regs);
+
+        #[cfg(feature = "disasm")]
+        if get_option!(dump_disasm) && !cb.outlined {
+            use crate::disasm::disasm_addr_range;
+            let last_ptr = cb.get_write_ptr();
+            let disasm = disasm_addr_range(cb, start_addr, last_ptr.raw_ptr() as usize - start_addr as usize);
+            if disasm.len() > 0 {
+                println!("{disasm}");
+            }
+        }
+        gc_offsets
     }
 
-    /// Compile with a limited number of registers
+    /// Compile with a limited number of registers. Used only for unit tests.
     pub fn compile_with_num_regs(self, cb: &mut CodeBlock, num_regs: usize) -> Vec<u32>
     {
         let mut alloc_regs = Self::get_alloc_regs();

yjit/src/codegen.rs

@@ -622,9 +622,13 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, iseq: IseqPtr, insn_idx: u32) -> O
     cb.align_pos(64);
 
     let code_ptr = cb.get_write_ptr();
-    add_comment(cb, "yjit entry");
 
     let mut asm = Assembler::new();
+    if get_option!(dump_disasm) {
+        asm.comment(&format!("YJIT entry: {}", iseq_get_location(iseq)));
+    } else {
+        asm.comment("YJIT entry");
+    }
 
     asm.frame_setup();
 
@@ -748,6 +752,11 @@ pub fn gen_single_block(
     // Create a backend assembler instance
     let mut asm = Assembler::new();
 
+    #[cfg(feature = "disasm")]
+    if get_option!(dump_disasm) {
+        asm.comment(&format!("Block: {} (ISEQ offset: {})", iseq_get_location(blockid.iseq), blockid.idx));
+    }
+
     // For each instruction to compile
     // NOTE: could rewrite this loop with a std::iter::Iterator
     while insn_idx < iseq_size {
@@ -6049,8 +6058,8 @@ impl CodegenGlobals {
                 half_size
             );
 
-            let cb = CodeBlock::new(first_half);
+            let cb = CodeBlock::new(first_half, false);
-            let ocb = OutlinedCb::wrap(CodeBlock::new(second_half));
+            let ocb = OutlinedCb::wrap(CodeBlock::new(second_half, true));
 
             (cb, ocb)
         };

yjit/src/core.rs

@@ -1510,6 +1510,7 @@ fn regenerate_branch(cb: &mut CodeBlock, branch: &mut Branch) {
     cb.set_write_ptr(branch.start_addr.unwrap());
 
     let mut asm = Assembler::new();
+    asm.comment("regenerate_branch");
 
     (branch.gen_fn)(
         &mut asm,

yjit/src/disasm.rs

@@ -1,6 +1,9 @@
 use crate::core::*;
 use crate::cruby::*;
 use crate::yjit::yjit_enabled_p;
+use crate::asm::CodeBlock;
+use crate::codegen::CodePtr;
+use std::fmt::Write;
 
 /// Primitive called in yjit.rb
 /// Produce a string representing the disassembly for an ISEQ
@@ -36,7 +39,7 @@ pub extern "C" fn rb_yjit_disasm_iseq(_ec: EcPtr, _ruby_self: VALUE, iseqw: VALU
 
 #[cfg(feature = "disasm")]
 pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u32, end_idx: u32) -> String {
-        let mut out = String::from("");
+    let mut out = String::from("");
 
     // Get a list of block versions generated for this iseq
     let mut block_list = get_iseq_block_list(iseq);
@@ -67,26 +70,6 @@ pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u32, end_idx: u32) -> St
         total_code_size += blockref.borrow().code_size();
     }
 
-    // Initialize capstone
-    use capstone::prelude::*;
-
-    #[cfg(target_arch = "x86_64")]
-    let mut cs = Capstone::new()
-        .x86()
-        .mode(arch::x86::ArchMode::Mode64)
-        .syntax(arch::x86::ArchSyntax::Intel)
-        .build()
-        .unwrap();
-
-    #[cfg(target_arch = "aarch64")]
-    let mut cs = Capstone::new()
-        .arm64()
-        .mode(arch::arm64::ArchMode::Arm)
-        .detail(true)
-        .build()
-        .unwrap();
-    cs.set_skipdata(true);
-
     out.push_str(&format!("NUM BLOCK VERSIONS: {}\n", block_list.len()));
     out.push_str(&format!(
         "TOTAL INLINE CODE SIZE: {} bytes\n",
@@ -115,19 +98,7 @@ pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u32, end_idx: u32) -> St
             out.push_str(&format!("== {:=<60}\n", block_ident));
 
             // Disassemble the instructions
-            let code_slice = unsafe { std::slice::from_raw_parts(start_addr, code_size) };
+            out.push_str(&disasm_addr_range(global_cb, start_addr, code_size));
-            let insns = cs.disasm_all(code_slice, start_addr as u64).unwrap();
-
-            // For each instruction in this block
-            for insn in insns.as_ref() {
-                // Comments for this block
-                if let Some(comment_list) = global_cb.comments_at(insn.address() as usize) {
-                    for comment in comment_list {
-                        out.push_str(&format!("  \x1b[1m# {}\x1b[0m\n", comment));
-                    }
-                }
-                out.push_str(&format!("  {}\n", insn));
-            }
 
             // If this is not the last block
             if block_idx < block_list.len() - 1 {
@@ -147,6 +118,49 @@ pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u32, end_idx: u32) -> St
     return out;
 }
 
+
+#[cfg(feature = "disasm")]
+pub fn disasm_addr_range(cb: &CodeBlock, start_addr: *const u8, code_size: usize) -> String {
+    let mut out = String::from("");
+
+    // Initialize capstone
+    use capstone::prelude::*;
+
+    #[cfg(target_arch = "x86_64")]
+    let mut cs = Capstone::new()
+        .x86()
+        .mode(arch::x86::ArchMode::Mode64)
+        .syntax(arch::x86::ArchSyntax::Intel)
+        .build()
+        .unwrap();
+
+    #[cfg(target_arch = "aarch64")]
+    let mut cs = Capstone::new()
+        .arm64()
+        .mode(arch::arm64::ArchMode::Arm)
+        .detail(true)
+        .build()
+        .unwrap();
+    cs.set_skipdata(true);
+
+    // Disassemble the instructions
+    let code_slice = unsafe { std::slice::from_raw_parts(start_addr, code_size) };
+    let insns = cs.disasm_all(code_slice, start_addr as u64).unwrap();
+
+    // For each instruction in this block
+    for insn in insns.as_ref() {
+        // Comments for this block
+        if let Some(comment_list) = cb.comments_at(insn.address() as usize) {
+            for comment in comment_list {
+                write!(&mut out, "  \x1b[1m# {}\x1b[0m\n", comment).unwrap();
+            }
+        }
+        write!(&mut out, "  {}\n", insn).unwrap();
+    }
+
+    return out;
+}
+
 /// Primitive called in yjit.rb
 /// Produce a list of instructions compiled for an isew
 #[no_mangle]

yjit/src/options.rs

@@ -30,6 +30,9 @@ pub struct Options {
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
 
+    /// Dump all compiled instructions in inlined CodeBlock
+    pub dump_disasm: bool,
+
     /// Print when specific ISEQ items are compiled or invalidated
     pub dump_iseq_disasm: Option<String>,
 
@@ -53,6 +56,7 @@ pub static mut OPTIONS: Options = Options {
     gen_stats: false,
     gen_trace_exits: false,
     dump_insns: false,
+    dump_disasm: false,
     verify_ctx: false,
     global_constant_state: false,
     dump_iseq_disasm: None,
@@ -128,6 +132,7 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         ("stats", "") => unsafe { OPTIONS.gen_stats = true },
         ("trace-exits", "") => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true },
         ("dump-insns", "") => unsafe { OPTIONS.dump_insns = true },
+        ("dump-disasm", "") => unsafe { OPTIONS.dump_disasm = true },
         ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },
         ("global-constant-state", "") => unsafe { OPTIONS.global_constant_state = true },