YJIT: Non-incremental code GC

k0kubun · k0kubun · commit 2629505fbf98 · 2022-10-17T15:21:56.000-07:00
diff --git a/yjit.c b/yjit.c
@@ -65,10 +65,7 @@ STATIC_ASSERT(pointer_tagging_scheme, USE_FLONUM);
 bool
 rb_yjit_mark_writable(void *mem_block, uint32_t mem_size)
 {
-    if (mprotect(mem_block, mem_size, PROT_READ | PROT_WRITE)) {
-        return false;
-    }
-    return true;
+    return mprotect(mem_block, mem_size, PROT_READ | PROT_WRITE) == 0;
 }
 
 void
@@ -85,6 +82,21 @@ rb_yjit_mark_executable(void *mem_block, uint32_t mem_size)
     }
 }
 
+// Free the specified memory block.
+bool
+rb_yjit_mark_unused(void *mem_block, uint32_t mem_size)
+{
+    // It seems like there's no way to free a memory block that works for all platforms.
+    // So this implementation is OS-dependent.
+#ifdef __APPLE__
+    // On macOS, mprotect PROT_NONE seems to reduce RSS immediately.
+    return mprotect(mem_block, mem_size, PROT_NONE) == 0;
+#else
+    // On Linux, however, you need to use madvise MADV_DONTNEED to do the same thing.
+    return madvise(mem_block, mem_size, MADV_DONTNEED) == 0;
+#endif
+}
+
 // `start` is inclusive and `end` is exclusive.
 void
 rb_yjit_icache_invalidate(void *start, void *end)
@@ -1018,6 +1030,19 @@ rb_yjit_invalidate_all_method_lookup_assumptions(void)
     // method caches, so we do nothing here for now.
 }
 
+static void
+update_in_use_flag(const rb_iseq_t *iseq)
+{
+    VALUE rb_yjit_update_payload_in_use(const rb_iseq_t *iseq, bool flag);
+    rb_yjit_update_payload_in_use(iseq, false);
+}
+
+void
+rb_yjit_scan_in_use_payloads(void)
+{
+    rb_yjit_for_each_iseq(update_in_use_flag);
+}
+
 // Primitives used by yjit.rb
 VALUE rb_yjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_trace_exit_locations_enabled_p(rb_execution_context_t *ec, VALUE self);
diff --git a/yjit.rb b/yjit.rb
@@ -218,7 +218,9 @@ def _print_stats
       $stderr.puts "invalidation_count:    " + ("%10d" % stats[:invalidation_count])
       $stderr.puts "constant_state_bumps:  " + ("%10d" % stats[:constant_state_bumps])
       $stderr.puts "inline_code_size:      " + ("%10d" % stats[:inline_code_size])
+      $stderr.puts "inline_freed_size:     " + ("%10d" % stats[:inline_freed_size])
       $stderr.puts "outlined_code_size:    " + ("%10d" % stats[:outlined_code_size])
+      $stderr.puts "outlined_freed_size:   " + ("%10d" % stats[:outlined_freed_size])
       $stderr.puts "num_gc_obj_refs:       " + ("%10d" % stats[:num_gc_obj_refs])
 
       $stderr.puts "total_exit_count:      " + ("%10d" % total_exits)
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
@@ -260,6 +260,7 @@ fn main() {
         .allowlist_function("rb_yjit_reserve_addr_space")
         .allowlist_function("rb_yjit_mark_writable")
         .allowlist_function("rb_yjit_mark_executable")
+        .allowlist_function("rb_yjit_mark_unused")
         .allowlist_function("rb_yjit_get_page_size")
         .allowlist_function("rb_leaf_invokebuiltin_iseq_p")
         .allowlist_function("rb_leaf_builtin_function")
diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
@@ -82,6 +82,9 @@ pub struct CodeBlock {
     // for example, when there is not enough space or when a jump
     // target is too far away.
     dropped_bytes: bool,
+
+    // The number of bytes that have been freed by code GC
+    pub freed_bytes: usize,
 }
 
 /// Set of CodeBlock label states. Used for recovering the previous state.
@@ -108,6 +111,7 @@ impl CodeBlock {
             asm_comments: BTreeMap::new(),
             outlined,
             dropped_bytes: false,
+            freed_bytes: 0,
         };
         cb.write_pos = cb.page_start();
         cb
@@ -124,6 +128,7 @@ impl CodeBlock {
         let next_page_idx = self.write_pos / self.page_size + 1;
         if !self.set_page(next_page_idx, &jmp_ptr) {
             self.set_write_ptr(old_write_ptr); // rollback if there are no more pages
+            self.code_gc(); // TODO: return true if there's a freed page once we implement freed page reuse
             return false;
         }
 
@@ -293,6 +298,11 @@ impl CodeBlock {
         self.mem_block.borrow().start_ptr().add_bytes(offset)
     }
 
+    pub fn addrs_to_pages(&self, start_addr: CodePtr, end_addr: CodePtr) -> Vec<usize> {
+        // TODO: move page management to codeblock
+        self.mem_block.borrow().addrs_to_pages(start_addr, end_addr)
+    }
+
     /// Get a (possibly dangling) direct pointer to the current write position
     pub fn get_write_ptr(&self) -> CodePtr {
         self.get_ptr(self.write_pos)
@@ -431,6 +441,44 @@ impl CodeBlock {
         self.mem_block.borrow_mut().mark_all_executable();
     }
 
+    /// Free unused code pages
+    fn code_gc(&mut self) {
+        //println!("code_gc");
+        let iseq_payloads = CodegenGlobals::get_iseq_payloads();
+
+        // Check which pages are still in use
+        let num_pages = self.mem_block.borrow().num_pages();
+        let mut alive_pages = vec![false; num_pages]; // TODO: consider using a bitmap
+        for iseq_payload in iseq_payloads.iter() {
+            let iseq_payload = unsafe { iseq_payload.as_ref() }.unwrap();
+            for page in &iseq_payload.pages {
+                alive_pages[*page] = true;
+            }
+        }
+
+        // Ask VirtualMem to free unused pages
+        for (page, alive) in alive_pages.iter().enumerate() {
+            if !alive {
+                // TODO: reduce the number of syscalls by merging them for consecutive pages
+                let freed_bytes = self.mem_block.borrow_mut().free_page(page);
+                //println!("free_page: {}", page);
+                self.freed_bytes += freed_bytes;
+            }
+        }
+
+        let mut cb = CodegenGlobals::get_inline_cb();
+        cb.dropped_bytes = true;
+        let mut ocb = CodegenGlobals::get_outlined_cb().unwrap();
+        cb.dropped_bytes = true;
+
+        // TODO: Implement start-over; allow reusing freed pages
+        // 1. scan cfp stack of all Fibers and list up which JIT payloads are still in use
+        // 2. invalidate existing in-use code
+        // 3. list up pages used by invalidation code
+        // 4. make vector of freed pages + free existing pages
+        // 5. move write_pos back to first freed page
+    }
+
     pub fn inline(&self) -> bool {
         !self.outlined
     }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -17,6 +17,7 @@ use std::cell::RefCell;
 use std::cell::RefMut;
 use std::cmp;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::ffi::CStr;
 use std::mem::{self, size_of};
 use std::os::raw::c_uint;
@@ -653,6 +654,10 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, iseq: IseqPtr, insn_idx: u32) -> O
     if (cb.has_dropped_bytes()) {
         None
     } else {
+        let mut iseq_payload = unsafe { get_or_create_iseq_payload(iseq) };
+        for page in cb.addrs_to_pages(code_ptr, cb.get_write_ptr()) {
+            iseq_payload.pages.insert(page);
+        }
         Some(code_ptr)
     }
 }
@@ -6519,6 +6524,9 @@ pub struct CodegenGlobals {
 
     // Methods for generating code for hardcoded (usually C) methods
     method_codegen_table: HashMap<usize, MethodGenFn>,
+
+    /// List of the payload of ISEQs that are JITed and not GCed yet
+    iseq_payloads: HashSet<*const IseqPayload>,
 }
 
 /// For implementing global code invalidation. A position in the inline
@@ -6600,6 +6608,7 @@ impl CodegenGlobals {
             global_inval_patches: Vec::new(),
             inline_frozen_bytes: 0,
             method_codegen_table: HashMap::new(),
+            iseq_payloads: HashSet::new(),
         };
 
         // Register the method codegen functions
@@ -6735,6 +6744,10 @@ impl CodegenGlobals {
             Some(&mgf) => Some(mgf), // Deref
         }
     }
+
+    pub fn get_iseq_payloads() -> &'static mut HashSet<*const IseqPayload> {
+        &mut CodegenGlobals::get_instance().iseq_payloads
+    }
 }
 
 #[cfg(test)]
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
@@ -11,6 +11,7 @@ use crate::utils::*;
 use crate::disasm::*;
 use core::ffi::c_void;
 use std::cell::*;
+use std::collections::HashSet;
 use std::hash::{Hash, Hasher};
 use std::mem;
 use std::rc::{Rc};
@@ -321,7 +322,7 @@ struct Branch {
 
     // Positions where the generated code starts and ends
     start_addr: Option<CodePtr>,
-    end_addr: Option<CodePtr>,
+    end_addr: Option<CodePtr>, // exclusive
 
     // Context right after the branch instruction
     #[allow(unused)] // set but not read at the moment
@@ -475,7 +476,9 @@ impl Eq for BlockRef {}
 /// when calling into YJIT
 #[derive(Default)]
 pub struct IseqPayload {
+    pub pages: HashSet<usize>,
     version_map: VersionMap,
+    pub in_use: bool,
 }
 
 impl IseqPayload {
@@ -498,7 +501,7 @@ pub fn get_iseq_payload(iseq: IseqPtr) -> Option<&'static mut IseqPayload> {
 }
 
 /// Get the payload object associated with an iseq. Create one if none exists.
-fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
+pub fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
     type VoidPtr = *mut c_void;
 
     let payload_non_null = unsafe {
@@ -513,6 +516,7 @@ fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
             // We allocate in those cases anyways.
             let new_payload = Box::into_raw(Box::new(IseqPayload::default()));
             rb_iseq_set_yjit_payload(iseq, new_payload as VoidPtr);
+            CodegenGlobals::get_iseq_payloads().insert(new_payload);
 
             new_payload
         } else {
@@ -539,17 +543,20 @@ pub extern "C" fn rb_yjit_iseq_free(payload: *mut c_void) {
         }
     };
 
-    use crate::invariants;
+    // Allow code GC to free this ISEQ's JIT code
+    CodegenGlobals::get_iseq_payloads().remove(&(payload as *const IseqPayload));
 
     // Take ownership of the payload with Box::from_raw().
     // It drops right before this function returns.
     // SAFETY: We got the pointer from Box::into_raw().
     let payload = unsafe { Box::from_raw(payload) };
+    //println!("pages:{:?}, versions:{:?}", payload.pages, payload.version_map.len());
 
     // Increment the freed iseq count
     incr_counter!(freed_iseq_count);
 
     // Remove all blocks in the payload from global invariants table.
+    use crate::invariants;
     for versions in &payload.version_map {
         for block in versions {
             invariants::block_assumptions_free(&block);
@@ -844,6 +851,11 @@ fn add_block_version(blockref: &BlockRef, cb: &CodeBlock) {
     }
 
     incr_counter!(compiled_block_count);
+
+    let mut iseq_payload = unsafe { get_iseq_payload(block.blockid.iseq) }.unwrap();
+    for page in cb.addrs_to_pages(block.start_addr.unwrap(), block.end_addr.unwrap()) {
+        iseq_payload.pages.insert(page);
+    }
 }
 
 /// Remove a block version from the version map of its parent ISEQ
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
@@ -1271,6 +1271,9 @@ extern "C" {
 extern "C" {
     pub fn rb_yjit_mark_executable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32);
 }
+extern "C" {
+    pub fn rb_yjit_mark_unused(mem_block: *mut ::std::os::raw::c_void, mem_size: u32) -> bool;
+}
 extern "C" {
     pub fn rb_yjit_icache_invalidate(
         start: *mut ::std::os::raw::c_void,
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
@@ -364,10 +364,20 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
         let value = VALUE::fixnum_from_usize(cb.get_write_pos());
         rb_hash_aset(hash, key, value);
 
+        // GCed inline code size
+        let key = rust_str_to_sym("inline_freed_size");
+        let value = VALUE::fixnum_from_usize(cb.freed_bytes);
+        rb_hash_aset(hash, key, value);
+
         // Outlined code size
         let key = rust_str_to_sym("outlined_code_size");
         let value = VALUE::fixnum_from_usize(ocb.unwrap().get_write_pos());
         rb_hash_aset(hash, key, value);
+
+        // GCed outlined code size
+        let key = rust_str_to_sym("outlined_freed_size");
+        let value = VALUE::fixnum_from_usize(ocb.unwrap().freed_bytes);
+        rb_hash_aset(hash, key, value);
     }
 
     // If we're not generating stats, the hash is done
diff --git a/yjit/src/virtualmem.rs b/yjit/src/virtualmem.rs
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
