From 0710bec01e405af596fa3e8e29ba38bb8176e39d Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Tue, 23 Feb 2021 15:22:20 -0500
Subject: [PATCH] Implement --ujit-call-threshold

---
 mjit.h       |  2 +-
 ruby.c       |  7 ++++---
 ujit.h       | 19 +++++++++----------
 ujit_iface.c | 20 +++++++++++++++++---
 4 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/mjit.h b/mjit.h
index 0264962da8..83f5cb919a 100644
--- a/mjit.h
+++ b/mjit.h
@@ -150,7 +150,7 @@ mjit_exec(rb_execution_context_t *ec)
     }
 
 #ifndef MJIT_HEADER
-    if (rb_ujit_enabled_p() && !mjit_call_p && body->total_calls == UJIT_CALL_THRESHOLD)  {
+    if (rb_ujit_enabled_p() && !mjit_call_p && body->total_calls == rb_ujit_call_threshold())  {
         rb_ujit_compile_iseq(iseq);
         return Qundef;
     }
diff --git a/ruby.c b/ruby.c
index b5c705002a..78d400278f 100644
--- a/ruby.c
+++ b/ruby.c
@@ -1033,12 +1033,13 @@ set_option_encoding_once(const char *type, VALUE *name, const char *e, long elen
 static void
 setup_ujit_options(const char *s, struct rb_ujit_options *ujit_opt)
 {
-    *ujit_opt = (struct rb_ujit_options) { 0 };
-
     if (*s != '-') return;
     const size_t l = strlen(++s);
 
-    if (opt_match_noarg(s, l, "stats")) {
+    if (opt_match_arg(s, l, "call-threshold")) {
+        ujit_opt->call_threshold = atoi(s + 1);
+    }
+    else if (opt_match_noarg(s, l, "stats")) {
         ujit_opt->gen_stats = true;
     }
     else {
diff --git a/ujit.h b/ujit.h
index c2e03e5e26..6957e9178f 100644
--- a/ujit.h
+++ b/ujit.h
@@ -32,22 +32,21 @@ typedef struct rb_iseq_struct rb_iseq_t;
 #endif
 
 struct rb_ujit_options {
+    bool ujit_enabled;
+
+    // Number of method calls after which to start generating code
+    // Threshold==1 means compile on first execution
+    unsigned call_threshold;
+
+    // Capture and print out stats
     bool gen_stats;
 };
 
 RUBY_SYMBOL_EXPORT_BEGIN
-RUBY_EXTERN bool rb_ujit_enabled;
+bool rb_ujit_enabled_p(void);
+unsigned rb_ujit_call_threshold(void);
 RUBY_SYMBOL_EXPORT_END
 
-static inline
-bool rb_ujit_enabled_p(void)
-{
-    return rb_ujit_enabled;
-}
-
-// Threshold==1 means compile on first execution
-#define UJIT_CALL_THRESHOLD (2u)
-
 void rb_ujit_collect_vm_usage_insn(int insn);
 void rb_ujit_method_lookup_change(VALUE cme_or_cc);
 void rb_ujit_compile_iseq(const rb_iseq_t *iseq);
diff --git a/ujit_iface.c b/ujit_iface.c
index b9368b2fe1..9893c73467 100644
--- a/ujit_iface.c
+++ b/ujit_iface.c
@@ -24,12 +24,12 @@ VALUE cUjitBlock;
 VALUE cUjitDisasm;
 VALUE cUjitDisasmInsn;
 
-bool rb_ujit_enabled;
-
+#if RUBY_DEBUG
 static int64_t vm_insns_count = 0;
 int64_t rb_ujit_exec_insns_count = 0;
 static int64_t exit_op_count[VM_INSTRUCTION_SIZE] = { 0 };
 int64_t rb_compiled_iseq_count = 0;
+#endif
 
 // Machine code blocks (executable memory)
 extern codeblock_t *cb;
@@ -627,6 +627,16 @@ rb_ujit_iseq_free(const struct rb_iseq_constant_body *body)
     rb_darray_free(body->ujit_blocks);
 }
 
+bool rb_ujit_enabled_p(void)
+{
+    return rb_ujit_opts.ujit_enabled;
+}
+
+unsigned rb_ujit_call_threshold(void)
+{
+    return rb_ujit_opts.call_threshold;
+}
+
 void
 rb_ujit_init(struct rb_ujit_options *options)
 {
@@ -636,8 +646,12 @@ rb_ujit_init(struct rb_ujit_options *options)
     }
 
     rb_ujit_opts = *options;
+    rb_ujit_opts.ujit_enabled = true;
 
-    rb_ujit_enabled = true;
+    // Normalize options
+    if (rb_ujit_opts.call_threshold < 1) {
+        rb_ujit_opts.call_threshold = 10;
+    }
 
     ujit_init_core();
     ujit_init_codegen();