RJIT: Introduce --rjit-exec-mem-size

lib/ruby_vm/rjit/code_block.rb

@@ -18,9 +18,9 @@ module RubyVM::RJIT
       start_addr = write_addr
 
       # Write machine code
-      C.rjit_mark_writable
+      C.mprotect_write(@mem_block, @mem_size)
       @write_pos += asm.assemble(start_addr)
-      C.rjit_mark_executable
+      C.mprotect_exec(@mem_block, @mem_size)
 
       end_addr = write_addr
 

lib/ruby_vm/rjit/compiler.rb

@@ -39,9 +39,10 @@ module RubyVM::RJIT
       INSNS.fetch(C.rb_vm_insn_decode(encoded))
     end
 
-    # @param mem_block [Integer] JIT buffer address
+    # @param mem_size [Integer] JIT buffer size
-    # @param mem_size  [Integer] JIT buffer size
+    def initialize
-    def initialize(mem_block, mem_size)
+      mem_size = C.rjit_opts.exec_mem_size * 1024 * 1024
+      mem_block = C.mmap(mem_size)
       @cb = CodeBlock.new(mem_block: mem_block, mem_size: mem_size / 2)
       @ocb = CodeBlock.new(mem_block: mem_block + mem_size / 2, mem_size: mem_size / 2, outlined: true)
       @exit_compiler = ExitCompiler.new

rjit.c

@@ -74,9 +74,6 @@ bool rb_rjit_call_p = false;
 // A flag to communicate that rb_rjit_call_p should be disabled while it's temporarily false.
 static bool rjit_cancel_p = false;
 
-// JIT buffer
-uint8_t *rb_rjit_mem_block = NULL;
-
 // `rb_ec_ractor_hooks(ec)->events` is moved to this variable during compilation.
 rb_event_flag_t rb_rjit_global_events = 0;
 
@@ -98,6 +95,8 @@ static VALUE rb_mRJITHooks = 0;
 
 // A default threshold used to add iseq to JIT.
 #define DEFAULT_CALL_THRESHOLD 30
+// Size of executable memory block in MiB.
+#define DEFAULT_EXEC_MEM_SIZE 64
 
 #define opt_match_noarg(s, l, name) \
     opt_match(s, l, name) && (*(s) ? (rb_warn("argument to --rjit-" name " is ignored"), 1) : 1)
@@ -117,6 +116,9 @@ rb_rjit_setup_options(const char *s, struct rjit_options *rjit_opt)
     else if (opt_match_arg(s, l, "call-threshold")) {
         rjit_opt->call_threshold = atoi(s + 1);
     }
+    else if (opt_match_arg(s, l, "exec-mem-size")) {
+        rjit_opt->exec_mem_size = atoi(s + 1);
+    }
     // --rjit=pause is an undocumented feature for experiments
     else if (opt_match_noarg(s, l, "pause")) {
         rjit_opt->pause = true;
@@ -135,6 +137,7 @@ const struct ruby_opt_message rb_rjit_option_messages[] = {
 #if RJIT_STATS
     M("--rjit-stats",              "", "Enable collecting RJIT statistics"),
 #endif
+    M("--rjit-exec-mem-size=num",  "", "Size of executable memory block in MiB (default: " STRINGIZE(DEFAULT_EXEC_MEM_SIZE) ")"),
     M("--rjit-call-threshold=num", "", "Number of calls to trigger JIT (default: " STRINGIZE(DEFAULT_CALL_THRESHOLD) ")"),
 #ifdef HAVE_LIBCAPSTONE
     M("--rjit-dump-disasm",        "", "Dump all JIT code"),
@@ -143,105 +146,6 @@ const struct ruby_opt_message rb_rjit_option_messages[] = {
 };
 #undef M
 
-#if defined(MAP_FIXED_NOREPLACE) && defined(_SC_PAGESIZE)
-// Align the current write position to a multiple of bytes
-static uint8_t *
-align_ptr(uint8_t *ptr, uint32_t multiple)
-{
-    // Compute the pointer modulo the given alignment boundary
-    uint32_t rem = ((uint32_t)(uintptr_t)ptr) % multiple;
-
-    // If the pointer is already aligned, stop
-    if (rem == 0)
-        return ptr;
-
-    // Pad the pointer by the necessary amount to align it
-    uint32_t pad = multiple - rem;
-
-    return ptr + pad;
-}
-#endif
-
-// Address space reservation. Memory pages are mapped on an as needed basis.
-// See the Rust mm module for details.
-static uint8_t *
-rjit_reserve_addr_space(uint32_t mem_size)
-{
-#ifndef _WIN32
-    uint8_t *mem_block;
-
-    // On Linux
-    #if defined(MAP_FIXED_NOREPLACE) && defined(_SC_PAGESIZE)
-        uint32_t const page_size = (uint32_t)sysconf(_SC_PAGESIZE);
-        uint8_t *const cfunc_sample_addr = (void *)&rjit_reserve_addr_space;
-        uint8_t *const probe_region_end = cfunc_sample_addr + INT32_MAX;
-        // Align the requested address to page size
-        uint8_t *req_addr = align_ptr(cfunc_sample_addr, page_size);
-
-        // Probe for addresses close to this function using MAP_FIXED_NOREPLACE
-        // to improve odds of being in range for 32-bit relative call instructions.
-        do {
-            mem_block = mmap(
-                req_addr,
-                mem_size,
-                PROT_NONE,
-                MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE,
-                -1,
-                0
-            );
-
-            // If we succeeded, stop
-            if (mem_block != MAP_FAILED) {
-                break;
-            }
-
-            // +4MB
-            req_addr += 4 * 1024 * 1024;
-        } while (req_addr < probe_region_end);
-
-    // On MacOS and other platforms
-    #else
-        // Try to map a chunk of memory as executable
-        mem_block = mmap(
-            (void *)rjit_reserve_addr_space,
-            mem_size,
-            PROT_NONE,
-            MAP_PRIVATE | MAP_ANONYMOUS,
-            -1,
-            0
-        );
-    #endif
-
-    // Fallback
-    if (mem_block == MAP_FAILED) {
-        // Try again without the address hint (e.g., valgrind)
-        mem_block = mmap(
-            NULL,
-            mem_size,
-            PROT_NONE,
-            MAP_PRIVATE | MAP_ANONYMOUS,
-            -1,
-            0
-        );
-    }
-
-    // Check that the memory mapping was successful
-    if (mem_block == MAP_FAILED) {
-        perror("ruby: yjit: mmap:");
-        if(errno == ENOMEM) {
-            // No crash report if it's only insufficient memory
-            exit(EXIT_FAILURE);
-        }
-        rb_bug("mmap failed");
-    }
-
-    return mem_block;
-#else
-    // Windows not supported for now
-    return NULL;
-#endif
-}
-
 #if RJIT_STATS
 struct rb_rjit_runtime_counters rb_rjit_counters = { 0 };
 
@@ -469,9 +373,17 @@ void
 rb_rjit_init(const struct rjit_options *opts)
 {
     VM_ASSERT(rb_rjit_enabled);
-    rb_rjit_opts = *opts;
 
-    rb_rjit_mem_block = rjit_reserve_addr_space(RJIT_CODE_SIZE);
+    // Normalize options
+    rb_rjit_opts = *opts;
+    if (rb_rjit_opts.exec_mem_size == 0)
+        rb_rjit_opts.exec_mem_size = DEFAULT_EXEC_MEM_SIZE;
+    if (rb_rjit_opts.call_threshold == 0)
+        rb_rjit_opts.call_threshold = DEFAULT_CALL_THRESHOLD;
+#ifndef HAVE_LIBCAPSTONE
+    if (rb_rjit_opts.dump_disasm)
+        rb_warn("libcapstone has not been linked. Ignoring --rjit-dump-disasm.");
+#endif
 
     // RJIT doesn't support miniruby, but it might reach here by RJIT_FORCE_ENABLE.
     rb_mRJIT = rb_const_get(rb_cRubyVM, rb_intern("RJIT"));
@@ -482,22 +394,14 @@ rb_rjit_init(const struct rjit_options *opts)
     }
     rb_mRJITC = rb_const_get(rb_mRJIT, rb_intern("C"));
     VALUE rb_cRJITCompiler = rb_const_get(rb_mRJIT, rb_intern("Compiler"));
-    rb_RJITCompiler = rb_funcall(rb_cRJITCompiler, rb_intern("new"), 2,
+    rb_RJITCompiler = rb_funcall(rb_cRJITCompiler, rb_intern("new"), 0);
-                                 SIZET2NUM((size_t)rb_rjit_mem_block), UINT2NUM(RJIT_CODE_SIZE));
     rb_cRJITIseqPtr = rb_funcall(rb_mRJITC, rb_intern("rb_iseq_t"), 0);
     rb_cRJITCfpPtr = rb_funcall(rb_mRJITC, rb_intern("rb_control_frame_t"), 0);
     rb_mRJITHooks = rb_const_get(rb_mRJIT, rb_intern("Hooks"));
 
+    // Enable RJIT and stats from here
     rb_rjit_call_p = !rb_rjit_opts.pause;
     rjit_stats_p = rb_rjit_opts.stats;
-
-    // Normalize options
-    if (rb_rjit_opts.call_threshold == 0)
-        rb_rjit_opts.call_threshold = DEFAULT_CALL_THRESHOLD;
-#ifndef HAVE_LIBCAPSTONE
-    if (rb_rjit_opts.dump_disasm)
-        rb_warn("libcapstone has not been linked. Ignoring --rjit-dump-disasm.");
-#endif
 }
 
 //

rjit.h

@@ -54,8 +54,10 @@ struct rjit_options {
     char* debug_flags;
     // If true, all ISeqs are synchronously compiled. For testing.
     bool wait;
-    // Number of calls to trigger JIT compilation. For testing.
+    // Number of calls to trigger JIT compilation.
     unsigned int call_threshold;
+    // Size of executable memory block in MiB
+    unsigned int exec_mem_size;
     // Collect RJIT statistics
     bool stats;
     // Force printing info about RJIT work of level VERBOSE or

rjit_c.c

@@ -33,24 +33,125 @@
 
 #include <errno.h>
 
-static bool
+#if defined(MAP_FIXED_NOREPLACE) && defined(_SC_PAGESIZE)
-rjit_mark_writable(void *mem_block, uint32_t mem_size)
+// Align the current write position to a multiple of bytes
+static uint8_t *
+align_ptr(uint8_t *ptr, uint32_t multiple)
 {
-    return mprotect(mem_block, mem_size, PROT_READ | PROT_WRITE) == 0;
+    // Compute the pointer modulo the given alignment boundary
+    uint32_t rem = ((uint32_t)(uintptr_t)ptr) % multiple;
+
+    // If the pointer is already aligned, stop
+    if (rem == 0)
+        return ptr;
+
+    // Pad the pointer by the necessary amount to align it
+    uint32_t pad = multiple - rem;
+
+    return ptr + pad;
+}
+#endif
+
+// Address space reservation. Memory pages are mapped on an as needed basis.
+// See the Rust mm module for details.
+static uint8_t *
+rjit_reserve_addr_space(uint32_t mem_size)
+{
+#ifndef _WIN32
+    uint8_t *mem_block;
+
+    // On Linux
+    #if defined(MAP_FIXED_NOREPLACE) && defined(_SC_PAGESIZE)
+        uint32_t const page_size = (uint32_t)sysconf(_SC_PAGESIZE);
+        uint8_t *const cfunc_sample_addr = (void *)&rjit_reserve_addr_space;
+        uint8_t *const probe_region_end = cfunc_sample_addr + INT32_MAX;
+        // Align the requested address to page size
+        uint8_t *req_addr = align_ptr(cfunc_sample_addr, page_size);
+
+        // Probe for addresses close to this function using MAP_FIXED_NOREPLACE
+        // to improve odds of being in range for 32-bit relative call instructions.
+        do {
+            mem_block = mmap(
+                req_addr,
+                mem_size,
+                PROT_NONE,
+                MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE,
+                -1,
+                0
+            );
+
+            // If we succeeded, stop
+            if (mem_block != MAP_FAILED) {
+                break;
+            }
+
+            // +4MB
+            req_addr += 4 * 1024 * 1024;
+        } while (req_addr < probe_region_end);
+
+    // On MacOS and other platforms
+    #else
+        // Try to map a chunk of memory as executable
+        mem_block = mmap(
+            (void *)rjit_reserve_addr_space,
+            mem_size,
+            PROT_NONE,
+            MAP_PRIVATE | MAP_ANONYMOUS,
+            -1,
+            0
+        );
+    #endif
+
+    // Fallback
+    if (mem_block == MAP_FAILED) {
+        // Try again without the address hint (e.g., valgrind)
+        mem_block = mmap(
+            NULL,
+            mem_size,
+            PROT_NONE,
+            MAP_PRIVATE | MAP_ANONYMOUS,
+            -1,
+            0
+        );
+    }
+
+    // Check that the memory mapping was successful
+    if (mem_block == MAP_FAILED) {
+        perror("ruby: yjit: mmap:");
+        if(errno == ENOMEM) {
+            // No crash report if it's only insufficient memory
+            exit(EXIT_FAILURE);
+        }
+        rb_bug("mmap failed");
+    }
+
+    return mem_block;
+#else
+    // Windows not supported for now
+    return NULL;
+#endif
 }
 
-static void
+static VALUE
-rjit_mark_executable(void *mem_block, uint32_t mem_size)
+mprotect_write(rb_execution_context_t *ec, VALUE self, VALUE rb_mem_block, VALUE rb_mem_size)
 {
-    // Do not call mprotect when mem_size is zero. Some platforms may return
+    void *mem_block = (void *)NUM2SIZET(rb_mem_block);
-    // an error for it. https://github.com/Shopify/ruby/issues/450
+    uint32_t mem_size = NUM2UINT(rb_mem_size);
-    if (mem_size == 0) {
+    return RBOOL(mprotect(mem_block, mem_size, PROT_READ | PROT_WRITE) == 0);
-        return;
+}
-    }
+
+static VALUE
+mprotect_exec(rb_execution_context_t *ec, VALUE self, VALUE rb_mem_block, VALUE rb_mem_size)
+{
+    void *mem_block = (void *)NUM2SIZET(rb_mem_block);
+    uint32_t mem_size = NUM2UINT(rb_mem_size);
+    if (mem_size == 0) return Qfalse; // Some platforms return an error for mem_size 0.
+
     if (mprotect(mem_block, mem_size, PROT_READ | PROT_EXEC)) {
         rb_bug("Couldn't make JIT page (%p, %lu bytes) executable, errno: %s\n",
             mem_block, (unsigned long)mem_size, strerror(errno));
     }
+    return Qtrue;
 }
 
 static VALUE

rjit_c.h

@@ -102,9 +102,6 @@ struct compile_status {
 // New stuff from here
 //
 
-// TODO: Make it configurable
-#define RJIT_CODE_SIZE 64 * 1024 * 1024
-
 extern uint8_t *rb_rjit_mem_block;
 
 #define RJIT_RUNTIME_COUNTERS(...) struct rb_rjit_runtime_counters { size_t __VA_ARGS__; };

rjit_c.rb

@@ -5,20 +5,16 @@ module RubyVM::RJIT # :nodoc: all
   # This `class << C` section is for calling C functions. For importing variables
   # or macros as is, please consider using tool/rjit/bindgen.rb instead.
   class << C = Object.new
-    def rjit_mark_writable
+    def mmap(mem_size)
-      Primitive.cstmt! %{
+      Primitive.cexpr! 'SIZET2NUM((size_t)rjit_reserve_addr_space(NUM2UINT(mem_size)))'
-        extern bool rjit_mark_writable(void *mem_block, uint32_t mem_size);
-        rjit_mark_writable(rb_rjit_mem_block, RJIT_CODE_SIZE);
-        return Qnil;
-      }
     end
 
-    def rjit_mark_executable
+    def mprotect_write(mem_block, mem_size)
-      Primitive.cstmt! %{
+      Primitive.mprotect_write(mem_block, mem_size)
-        extern void rjit_mark_executable(void *mem_block, uint32_t mem_size);
+    end
-        rjit_mark_executable(rb_rjit_mem_block, RJIT_CODE_SIZE);
+
-        return Qnil;
+    def mprotect_exec(mem_block, mem_size)
-      }
+      Primitive.mprotect_exec(mem_block, mem_size)
     end
 
     def rjit_insn_exits
@@ -1623,6 +1619,7 @@ module RubyVM::RJIT # :nodoc: all
       debug_flags: [CType::Pointer.new { CType::Immediate.parse("char") }, Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), debug_flags)")],
       wait: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), wait)")],
       call_threshold: [CType::Immediate.parse("unsigned int"), Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), call_threshold)")],
+      exec_mem_size: [CType::Immediate.parse("unsigned int"), Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), exec_mem_size)")],
       stats: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), stats)")],
       verbose: [CType::Immediate.parse("int"), Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), verbose)")],
       max_cache_size: [CType::Immediate.parse("int"), Primitive.cexpr!("OFFSETOF((*((struct rjit_options *)NULL)), max_cache_size)")],