Add ability to trace exit locations in yjit (#5970)

When running with `--yjit-stats` turned on, yjit can inform the user
what the most common exits are. While this is useful information it
doesn't tell you the source location of the code that exited or what the
code that exited looks like. This change intends to fix that.

To use the feature, run yjit with the `--yjit-trace-exits` option,
which will record the backtrace for every exit that occurs. This functionality
requires the stats feature to be turned on. Calling `--yjit-trace-exits`
will automatically set the `--yjit-stats` option.

Users must call `RubyVM::YJIT.dump_exit_locations(filename)` which will
Marshal dump the contents of `RubyVM::YJIT.exit_locations` into a file
based on the passed filename.

*Example usage:*

Given the following script, we write to a file called
`concat_array.dump` the results of `RubyVM::YJIT.exit_locations`.

```ruby
def concat_array
  ["t", "r", *x = "u", "e"].join
end

1000.times do
  concat_array
end

RubyVM::YJIT.dump_exit_locations("concat_array.dump")
```

When we run the file with this branch and the appropriate flags the
stacktrace will be recorded. Note Stackprof needs to be installed or you
need to point to the library directly.

```
./ruby --yjit --yjit-call-threshold=1 --yjit-trace-exits -I/Users/eileencodes/open_source/stackprof/lib test.rb
```

We can then read the dump file with Stackprof:

```
./ruby -I/Users/eileencodes/open_source/stackprof/lib/ /Users/eileencodes/open_source/stackprof/bin/stackprof --text concat_array.dump
```

Results will look similar to the following:

```
==================================
  Mode: ()
  Samples: 1817 (0.00% miss rate)
  GC: 0 (0.00%)
==================================
     TOTAL    (pct)     SAMPLES    (pct)     FRAME
      1001  (55.1%)        1001  (55.1%)     concatarray
       335  (18.4%)         335  (18.4%)     invokeblock
       178   (9.8%)         178   (9.8%)     send
       140   (7.7%)         140   (7.7%)     opt_getinlinecache
       ...etc...
```

Simply inspecting the `concatarray` method will give `SOURCE
UNAVAILABLE` because the source is insns.def.

```
./ruby -I/Users/eileencodes/open_source/stackprof/lib/ /Users/eileencodes/open_source/stackprof/bin/stackprof --text concat_array.dump --method concatarray
```

Result:

```
concatarray (nonexistent.def:1)
  samples:  1001 self (55.1%)  /   1001 total (55.1%)
  callers:
    1000  (   99.9%)  Object#concat_array
       1  (    0.1%)  Gem.suffixes
  callees (0 total):
  code:
        SOURCE UNAVAILABLE
```

However if we go deeper to the callee we can see the exact
source of the `concatarray` exit.

```
./ruby -I/Users/eileencodes/open_source/stackprof/lib/ /Users/eileencodes/open_source/stackprof/bin/stackprof --text concat_array.dump --method Object#concat_array
```

```
Object#concat_array (/Users/eileencodes/open_source/rust_ruby/test.rb:1)
  samples:     0 self (0.0%)  /   1000 total (55.0%)
  callers:
    1000  (  100.0%)  block in <main>
  callees (1000 total):
    1000  (  100.0%)  concatarray
  code:
                                  |     1  | def concat_array
 1000   (55.0%)                   |     2  |   ["t", "r", *x = "u", "e"].join
                                  |     3  | end
```

The `--walk` option is recommended for this feature as it make it
easier to traverse the tree of exits.

*Goals of this feature:*

This feature is meant to give more information when working on YJIT.
The idea is that if we know what code is exiting we can decide what
areas to prioritize when fixing exits. In some cases this means adding
prioritizing avoiding certain exits in yjit. In more complex cases it
might mean changing the Ruby code to be more performant when run with
yjit. Ultimately the more information we have about what code is exiting
AND why, the better we can make yjit.

*Known limitations:*

* Due to tracing exits, running this on large codebases like Rails
can be quite slow.
* On complex methods it can still be difficult to pinpoint the exact cause of
an exit.
* Stackprof is a requirement to to view the backtrace information from
the dump file.

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>

10 files changed, 557 insertions, 0 deletions

diff --git a/test/ruby/test_yjit_exit_locations.rb b/test/ruby/test_yjit_exit_locations.rb
new file mode 100644
index 0000000000..d708bed5e9
--- /dev/null
+++ b/test/ruby/test_yjit_exit_locations.rb
@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+#
+# This set of tests can be run with:
+# make test-all TESTS='test/ruby/test_yjit_exit_locations.rb' RUN_OPTS="--yjit-call-threshold=1"
+
+require 'test/unit'
+require 'envutil'
+require 'tmpdir'
+require_relative '../lib/jit_support'
+
+return unless defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled? && RubyVM::YJIT.trace_exit_locations_enabled?
+
+# Tests for YJIT with assertions on tracing exits
+# insipired by the MJIT tests in test/ruby/test_yjit.rb
+class TestYJITExitLocations < Test::Unit::TestCase
+  def test_yjit_trace_exits_and_v_no_error
+    _stdout, stderr, _status = EnvUtil.invoke_ruby(%w(-v --yjit-trace-exits), '', true, true)
+    refute_includes(stderr, "NoMethodError")
+  end
+
+  def test_trace_exits_setclassvariable
+    script = 'class Foo; def self.foo; @@foo = 1; end; end; Foo.foo'
+    assert_exit_locations(script)
+  end
+
+  def test_trace_exits_putobject
+    assert_exit_locations('true')
+    assert_exit_locations('123')
+    assert_exit_locations(':foo')
+  end
+
+  def test_trace_exits_opt_not
+    assert_exit_locations('!false')
+    assert_exit_locations('!nil')
+    assert_exit_locations('!true')
+    assert_exit_locations('![]')
+  end
+
+  private
+
+  def assert_exit_locations(test_script)
+    write_results = <<~RUBY
+      IO.open(3).write Marshal.dump(RubyVM::YJIT.exit_locations)
+    RUBY
+
+    script = <<~RUBY
+      _test_proc = -> {
+        #{test_script}
+      }
+      result = _test_proc.call
+      #{write_results}
+    RUBY
+
+    exit_locations = eval_with_jit(script)
+
+    assert exit_locations.key?(:raw)
+    assert exit_locations.key?(:frames)
+    assert exit_locations.key?(:lines)
+    assert exit_locations.key?(:samples)
+    assert exit_locations.key?(:missed_samples)
+    assert exit_locations.key?(:gc_samples)
+
+    assert_equal 0, exit_locations[:missed_samples]
+    assert_equal 0, exit_locations[:gc_samples]
+
+    assert_not_empty exit_locations[:raw]
+    assert_not_empty exit_locations[:frames]
+    assert_not_empty exit_locations[:lines]
+
+    exit_locations[:frames].each do |frame_id, frame|
+      assert frame.key?(:name)
+      assert frame.key?(:file)
+      assert frame.key?(:samples)
+      assert frame.key?(:total_samples)
+      assert frame.key?(:edges)
+    end
+  end
+
+  def eval_with_jit(script)
+    args = [
+      "--disable-gems",
+      "--yjit-call-threshold=1",
+      "--yjit-trace-exits"
+    ]
+    args << "-e" << script_shell_encode(script)
+    stats_r, stats_w = IO.pipe
+    out, err, status = EnvUtil.invoke_ruby(args,
+                                           '', true, true, timeout: 1000, ios: { 3 => stats_w }
+                                          )
+    stats_w.close
+    stats = stats_r.read
+    stats = Marshal.load(stats) if !stats.empty?
+    stats_r.close
+    stats
+  end
+
+  def script_shell_encode(s)
+    # We can't pass utf-8-encoded characters directly in a shell arg. But we can use Ruby \u constants.
+    s.chars.map { |c| c.ascii_only? ? c : "\\u%x" % c.codepoints[0] }.join
+  end
+end
diff --git a/yjit.c b/yjit.c
index b97fcdb62c..b4e36e0b68 100644
--- a/yjit.c
+++ b/yjit.c
@@ -25,6 +25,7 @@
 #include "probes.h"
 #include "probes_helper.h"
 #include "iseq.h"
+#include "ruby/debug.h"
 
 // For mmapp(), sysconf()
 #ifndef _WIN32
@@ -83,6 +84,94 @@ rb_yjit_mark_executable(void *mem_block, uint32_t mem_size)
     }
 }
 
+# define PTR2NUM(x)   (rb_int2inum((intptr_t)(void *)(x)))
+
+// For a given raw_sample (frame), set the hash with the caller's
+// name, file, and line number. Return the  hash with collected frame_info.
+static void
+rb_yjit_add_frame(VALUE hash, VALUE frame)
+{
+    VALUE frame_id = PTR2NUM(frame);
+
+    if (RTEST(rb_hash_aref(hash, frame_id))) {
+        return;
+    } else {
+        VALUE frame_info = rb_hash_new();
+        // Full label for the frame
+        VALUE name = rb_profile_frame_full_label(frame);
+        // Absolute path of the frame from rb_iseq_realpath
+        VALUE file = rb_profile_frame_absolute_path(frame);
+        // Line number of the frame
+        VALUE line = rb_profile_frame_first_lineno(frame);
+
+        // If absolute path isn't available use the rb_iseq_path
+        if (NIL_P(file)) {
+            file = rb_profile_frame_path(frame);
+        }
+
+        rb_hash_aset(frame_info, ID2SYM(rb_intern("name")), name);
+        rb_hash_aset(frame_info, ID2SYM(rb_intern("file")), file);
+
+        if (line != INT2FIX(0)) {
+            rb_hash_aset(frame_info, ID2SYM(rb_intern("line")), line);
+        }
+
+       rb_hash_aset(hash, frame_id, frame_info);
+    }
+}
+
+// Parses the YjitExtiLocations raw_samples and line_samples collected by
+// rb_yjit_record_exit_stack and turns them into 3 hashes (raw, lines, and frames) to
+// be used by RubyVM::YJIT.exit_locations. yjit_raw_samples represents the raw frames information
+// (without name, file, and line), and yjit_line_samples represents the line information
+// of the iseq caller.
+VALUE
+rb_yjit_exit_locations_dict(VALUE *yjit_raw_samples, int *yjit_line_samples, int samples_len)
+{
+    VALUE result = rb_hash_new();
+    VALUE raw_samples = rb_ary_new_capa(samples_len);
+    VALUE line_samples = rb_ary_new_capa(samples_len);
+    VALUE frames = rb_hash_new();
+    int idx = 0;
+
+    // While the index is less than samples_len, parse yjit_raw_samples and
+    // yjit_line_samples, then add casted values to raw_samples and line_samples array.
+    while (idx < samples_len) {
+        int num = (int)yjit_raw_samples[idx];
+        int line_num = (int)yjit_line_samples[idx];
+        idx++;
+
+        rb_ary_push(raw_samples, SIZET2NUM(num));
+        rb_ary_push(line_samples, INT2NUM(line_num));
+
+        // Loop through the length of samples_len and add data to the
+        // frames hash. Also push the current value onto the raw_samples
+        // and line_samples arrary respectively.
+        for (int o = 0; o < num; o++) {
+            rb_yjit_add_frame(frames, yjit_raw_samples[idx]);
+            rb_ary_push(raw_samples, SIZET2NUM(yjit_raw_samples[idx]));
+            rb_ary_push(line_samples, INT2NUM(yjit_line_samples[idx]));
+            idx++;
+        }
+
+        rb_ary_push(raw_samples, SIZET2NUM(yjit_raw_samples[idx]));
+        rb_ary_push(line_samples, INT2NUM(yjit_line_samples[idx]));
+        idx++;
+
+        rb_ary_push(raw_samples, SIZET2NUM(yjit_raw_samples[idx]));
+        rb_ary_push(line_samples, INT2NUM(yjit_line_samples[idx]));
+        idx++;
+    }
+
+    // Set add the raw_samples, line_samples, and frames to the results
+    // hash.
+    rb_hash_aset(result, ID2SYM(rb_intern("raw")), raw_samples);
+    rb_hash_aset(result, ID2SYM(rb_intern("lines")), line_samples);
+    rb_hash_aset(result, ID2SYM(rb_intern("frames")), frames);
+
+    return result;
+}
+
 uint32_t
 rb_yjit_get_page_size(void)
 {
@@ -860,12 +949,14 @@ rb_yjit_invalidate_all_method_lookup_assumptions(void)
 
 // 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);
 VALUE rb_yjit_get_stats(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_reset_stats_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_disasm_iseq(rb_execution_context_t *ec, VALUE self, VALUE iseq);
 VALUE rb_yjit_insns_compiled(rb_execution_context_t *ec, VALUE self, VALUE iseq);
 VALUE rb_yjit_simulate_oom_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_get_stats(rb_execution_context_t *ec, VALUE self);
+VALUE rb_yjit_get_exit_locations(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"
diff --git a/yjit.rb b/yjit.rb
index 7c5311d79a..cbaea7d003 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -18,11 +18,111 @@ module RubyVM::YJIT
     Primitive.rb_yjit_stats_enabled_p
   end
 
+  # Check if rb_yjit_trace_exit_locations_enabled_p is enabled.
+  def self.trace_exit_locations_enabled?
+    Primitive.rb_yjit_trace_exit_locations_enabled_p
+  end
+
   # Discard statistics collected for --yjit-stats.
   def self.reset_stats!
     Primitive.rb_yjit_reset_stats_bang
   end
 
+  # If --yjit-trace-exits is enabled parse the hashes from
+  # Primitive.rb_yjit_get_exit_locations into a format readable
+  # by Stackprof. This will allow us to find the exact location of a
+  # side exit in YJIT based on the instruction that is exiting.
+  def self.exit_locations
+    return unless trace_exit_locations_enabled?
+
+    results = Primitive.rb_yjit_get_exit_locations
+    raw_samples = results[:raw].dup
+    line_samples = results[:lines].dup
+    frames = results[:frames].dup
+    samples_count = 0
+
+    frames.each do |frame_id, frame|
+      frame[:samples] = 0
+      frame[:edges] = {}
+    end
+
+    # Loop through the instructions and set the frame hash with the data.
+    # We use nonexistent.def for the file name, otherwise insns.def will be displayed
+    # and that information isn't useful in this context.
+    RubyVM::INSTRUCTION_NAMES.each_with_index do |name, frame_id|
+      frame_hash = { samples: 0, total_samples: 0, edges: {}, name: name, file: "nonexistent.def", line: nil }
+      results[:frames][frame_id] = frame_hash
+      frames[frame_id] = frame_hash
+    end
+
+    # Loop through the raw_samples and build the hashes for StackProf.
+    # The loop is based off an example in the StackProf documentation and therefore
+    # this functionality can only work with that library.
+    while raw_samples.length > 0
+      stack_trace = raw_samples.shift(raw_samples.shift + 1)
+      lines = line_samples.shift(line_samples.shift + 1)
+      prev_frame_id = nil
+
+      stack_trace.each_with_index do |frame_id, idx|
+        if prev_frame_id
+          prev_frame = frames[prev_frame_id]
+          prev_frame[:edges][frame_id] ||= 0
+          prev_frame[:edges][frame_id] += 1
+        end
+
+        frame_info = frames[frame_id]
+        frame_info[:total_samples] ||= 0
+        frame_info[:total_samples] += 1
+
+        frame_info[:lines] ||= {}
+        frame_info[:lines][lines[idx]] ||= [0, 0]
+        frame_info[:lines][lines[idx]][0] += 1
+
+        prev_frame_id = frame_id
+      end
+
+      top_frame_id = stack_trace.last
+      top_frame_line = 1
+
+      frames[top_frame_id][:samples] += 1
+      frames[top_frame_id][:lines] ||= {}
+      frames[top_frame_id][:lines][top_frame_line] ||= [0, 0]
+      frames[top_frame_id][:lines][top_frame_line][1] += 1
+
+      samples_count += raw_samples.shift
+      line_samples.shift
+    end
+
+    results[:samples] = samples_count
+    # Set missed_samples and gc_samples to 0 as their values
+    # don't matter to us in this context.
+    results[:missed_samples] = 0
+    results[:gc_samples] = 0
+    results
+  end
+
+  # Marshal dumps exit locations to the given filename.
+  #
+  # Usage:
+  #
+  # In a script call:
+  #
+  #   RubyVM::YJIT.dump_exit_locations("my_file.dump")
+  #
+  # Then run the file with the following options:
+  #
+  #   ruby --yjit --yjit-stats --yjit-trace-exits test.rb
+  #
+  # Once the code is done running, use Stackprof to read the dump file.
+  # See Stackprof documentation for options.
+  def self.dump_exit_locations(filename)
+    unless trace_exit_locations_enabled?
+      raise ArgumentError, "--yjit-trace-exits must be enabled to use dump_exit_locations."
+    end
+
+    File.write(filename, Marshal.dump(RubyVM::YJIT.exit_locations))
+  end
+
   # Return a hash for statistics generated for the --yjit-stats command line option.
   # Return nil when option is not passed or unavailable.
   def self.runtime_stats
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index c3a18c9c6f..212013d70c 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -259,6 +259,7 @@ fn main() {
         .allowlist_function("rb_yjit_obj_written")
         .allowlist_function("rb_yjit_str_simple_append")
         .allowlist_function("rb_ENCODING_GET")
+        .allowlist_function("rb_yjit_exit_locations_dict")
 
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")
@@ -293,6 +294,9 @@ fn main() {
         .allowlist_function("rb_class_allocate_instance")
         .allowlist_function("rb_obj_info")
 
+        // From include/ruby/debug.h
+        .allowlist_function("rb_profile_frames")
+
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 4203196e34..4d7e7344ad 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -424,6 +424,13 @@ fn gen_exit(exit_pc: *mut VALUE, ctx: &Context, cb: &mut CodeBlock) -> CodePtr {
     if get_option!(gen_stats) {
         mov(cb, RDI, const_ptr_opnd(exit_pc as *const u8));
         call_ptr(cb, RSI, rb_yjit_count_side_exit_op as *const u8);
+
+        // If --yjit-trace-exits option is enabled, record the exit stack
+        // while recording the side exits.
+        if get_option!(gen_trace_exits) {
+            mov(cb, C_ARG_REGS[0], const_ptr_opnd(exit_pc as *const u8));
+            call_ptr(cb, REG0, rb_yjit_record_exit_stack as *const u8);
+        }
     }
 
     pop(cb, REG_SP);
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 758a8de2fd..b5dd356aef 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -968,12 +968,27 @@ extern "C" {
     pub fn rb_vm_barrier();
 }
 extern "C" {
+    pub fn rb_profile_frames(
+        start: ::std::os::raw::c_int,
+        limit: ::std::os::raw::c_int,
+        buff: *mut VALUE,
+        lines: *mut ::std::os::raw::c_int,
+    ) -> ::std::os::raw::c_int;
+}
+extern "C" {
     pub fn rb_yjit_mark_writable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32);
 }
 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_exit_locations_dict(
+        yjit_raw_samples: *mut VALUE,
+        yjit_line_samples: *mut ::std::os::raw::c_int,
+        samples_len: ::std::os::raw::c_int,
+    ) -> VALUE;
+}
+extern "C" {
     pub fn rb_yjit_get_page_size() -> u32;
 }
 extern "C" {
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index e457f2ea80..6329c70f87 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -360,6 +360,11 @@ pub extern "C" fn rb_yjit_constant_state_changed(id: ID) {
 /// See `struct yjijt_root_struct` in C.
 #[no_mangle]
 pub extern "C" fn rb_yjit_root_mark() {
+    // Call rb_gc_mark on exit location's raw_samples to
+    // wrap frames in a GC allocated object. This needs to be called
+    // at the same time as root mark.
+    YjitExitLocations::gc_mark_raw_samples();
+
     // Comment from C YJIT:
     //
     // Why not let the GC move the cme keys in this table?
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 07c501a108..704c709bae 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -24,6 +24,9 @@ pub struct Options {
     // Capture and print out stats
     pub gen_stats: bool,
 
+    // Trace locations of exits
+    pub gen_trace_exits: bool,
+
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
 
@@ -45,6 +48,7 @@ pub static mut OPTIONS: Options = Options {
     no_type_prop: false,
     max_versions: 4,
     gen_stats: false,
+    gen_trace_exits: false,
     dump_insns: false,
     verify_ctx: false,
     global_constant_state: false,
@@ -104,6 +108,7 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         ("greedy-versioning", "") => unsafe { OPTIONS.greedy_versioning = true },
         ("no-type-prop", "") => unsafe { OPTIONS.no_type_prop = true },
         ("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 },
         ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },
         ("global-constant-state", "") => unsafe { OPTIONS.global_constant_state = true },
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 3b7a3f31d9..e129cc2811 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -12,6 +12,105 @@ use crate::yjit::yjit_enabled_p;
 const VM_INSTRUCTION_SIZE_USIZE:usize = VM_INSTRUCTION_SIZE as usize;
 static mut EXIT_OP_COUNT: [u64; VM_INSTRUCTION_SIZE_USIZE] = [0; VM_INSTRUCTION_SIZE_USIZE];
 
+/// Global state needed for collecting backtraces of exits
+pub struct YjitExitLocations {
+    /// Vec to hold raw_samples which represent the control frames
+    /// of method entries.
+    raw_samples: Vec<VALUE>,
+    /// Vec to hold line_samples which represent line numbers of
+    /// the iseq caller.
+    line_samples: Vec<i32>
+}
+
+/// Private singleton instance of yjit exit locations
+static mut YJIT_EXIT_LOCATIONS: Option<YjitExitLocations> = None;
+
+impl YjitExitLocations {
+    /// Initialize the yjit exit locations
+    pub fn init() {
+        // Return if the stats feature is disabled
+        if !cfg!(feature = "stats") {
+            return;
+        }
+
+        // Return if --yjit-trace-exits isn't enabled
+        if !get_option!(gen_trace_exits) {
+            return;
+        }
+
+        let yjit_exit_locations = YjitExitLocations {
+            raw_samples: Vec::new(),
+            line_samples: Vec::new()
+        };
+
+        // Initialize the yjit exit locations instance
+        unsafe {
+            YJIT_EXIT_LOCATIONS = Some(yjit_exit_locations);
+        }
+    }
+
+    /// Get a mutable reference to the yjit exit locations globals instance
+    pub fn get_instance() -> &'static mut YjitExitLocations {
+        unsafe { YJIT_EXIT_LOCATIONS.as_mut().unwrap() }
+    }
+
+    /// Get a mutable reference to the yjit raw samples Vec
+    pub fn get_raw_samples() -> &'static mut Vec<VALUE> {
+        &mut YjitExitLocations::get_instance().raw_samples
+    }
+
+    /// Get a mutable reference to yjit the line samples Vec.
+    pub fn get_line_samples() -> &'static mut Vec<i32> {
+        &mut YjitExitLocations::get_instance().line_samples
+    }
+
+    /// Mark the data stored in YjitExitLocations::get_raw_samples that needs to be used by
+    /// rb_yjit_add_frame. YjitExitLocations::get_raw_samples are an array of
+    /// VALUE pointers, exit instruction, and number of times we've seen this stack row
+    /// as collected by rb_yjit_record_exit_stack.
+    ///
+    /// These need to have rb_gc_mark called so they can be used by rb_yjit_add_frame.
+    pub fn gc_mark_raw_samples() {
+        // Return if YJIT is not enabled
+        if !yjit_enabled_p() {
+            return;
+        }
+
+        // Return if the stats feature is disabled
+        if !cfg!(feature = "stats") {
+            return;
+        }
+
+        // Return if --yjit-trace-exits isn't enabled
+        if !get_option!(gen_trace_exits) {
+            return;
+        }
+
+        let mut idx: size_t = 0;
+        let yjit_raw_samples = YjitExitLocations::get_raw_samples();
+
+        while idx < yjit_raw_samples.len() as size_t {
+            let num = yjit_raw_samples[idx as usize];
+            let mut i = 0;
+            idx += 1;
+
+            // Mark the yjit_raw_samples at the given index. These represent
+            // the data that needs to be GC'd which are the current frames.
+            while i < i32::from(num) {
+                unsafe { rb_gc_mark(yjit_raw_samples[idx as usize]); }
+                i += 1;
+                idx += 1;
+            }
+
+            // Increase index for exit instruction.
+            idx += 1;
+            // Increase index for bookeeping value (number of times we've seen this
+            // row in a stack).
+            idx += 1;
+        }
+    }
+}
+
 // Macro to declare the stat counters
 macro_rules! make_counters {
     ($($counter_name:ident,)+) => {
@@ -168,6 +267,57 @@ pub extern "C" fn rb_yjit_get_stats(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     with_vm_lock(src_loc!(), || rb_yjit_gen_stats_dict())
 }
 
+/// Primitive called in yjit.rb
+///
+/// Check if trace_exits generation is enabled. Requires the stats feature
+/// to be enabled.
+#[no_mangle]
+pub extern "C" fn rb_yjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    #[cfg(feature = "stats")]
+    if get_option!(gen_trace_exits) {
+        return Qtrue;
+    }
+
+    return Qfalse;
+}
+
+/// Call the C function to parse the raw_samples and line_samples
+/// into raw, lines, and frames hash for RubyVM::YJIT.exit_locations.
+#[no_mangle]
+pub extern "C" fn rb_yjit_get_exit_locations(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    // Return if YJIT is not enabled
+    if !yjit_enabled_p() {
+        return Qnil;
+    }
+
+    // Return if the stats feature is disabled
+    if !cfg!(feature = "stats") {
+        return Qnil;
+    }
+
+    // Return if --yjit-trace-exits isn't enabled
+    if !get_option!(gen_trace_exits) {
+        return Qnil;
+    }
+
+    // If the stats feature is enabled, pass yjit_raw_samples and yjit_line_samples
+    // to the C function called rb_yjit_exit_locations_dict for parsing.
+    let yjit_raw_samples = YjitExitLocations::get_raw_samples();
+    let yjit_line_samples = YjitExitLocations::get_line_samples();
+
+    // Assert that the two Vec's are the same length. If they aren't
+    // equal something went wrong.
+    assert_eq!(yjit_raw_samples.len(), yjit_line_samples.len());
+
+    // yjit_raw_samples and yjit_line_samples are the same length so
+    // pass only one of the lengths in the C function.
+    let samples_len = yjit_raw_samples.len() as i32;
+
+    unsafe {
+        rb_yjit_exit_locations_dict(yjit_raw_samples.as_mut_ptr(), yjit_line_samples.as_mut_ptr(), samples_len)
+    }
+}
+
 /// Export all YJIT statistics as a Ruby hash.
 fn rb_yjit_gen_stats_dict() -> VALUE {
     // If YJIT is not enabled, return Qnil
@@ -231,6 +381,83 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
     hash
 }
 
+/// Record the backtrace when a YJIT exit occurs. This functionality requires
+/// that the stats feature is enabled as well as the --yjit-trace-exits option.
+///
+/// This function will fill two Vec's in YjitExitLocations to record the raw samples
+/// and line samples. Their length should be the same, however the data stored in
+/// them is different.
+#[no_mangle]
+pub extern "C" fn rb_yjit_record_exit_stack(exit_pc: *const VALUE)
+{
+    // Return if YJIT is not enabled
+    if !yjit_enabled_p() {
+        return;
+    }
+
+    // Return if the stats feature is disabled
+    if !cfg!(feature = "stats") {
+        return;
+    }
+
+    // Return if --yjit-trace-exits isn't enabled
+    if !get_option!(gen_trace_exits) {
+        return;
+    }
+
+    // rb_vm_insn_addr2opcode won't work in cargo test --all-features
+    // because it's a C function. Without insn call, this function is useless
+    // so wrap the whole thing in a not test check.
+    if cfg!(not(test)) {
+        // Get the opcode from the encoded insn handler at this PC
+        let insn = unsafe { rb_vm_insn_addr2opcode((*exit_pc).as_ptr()) };
+
+        // Use the same buffer size as Stackprof.
+        const BUFF_LEN: usize = 2048;
+
+        // Create 2 array buffers to be used to collect frames and lines.
+        let mut frames_buffer = [VALUE(0 as usize); BUFF_LEN];
+        let mut lines_buffer = [0; BUFF_LEN];
+
+        // Records call frame and line information for each method entry into two
+        // temporary buffers. Returns the number of times we added to the buffer (ie
+        // the length of the stack).
+        //
+        // Call frame info is stored in the frames_buffer, line number information
+        // in the lines_buffer. The first argument is the start point and the second
+        // argument is the buffer limit, set at 2048.
+        let num = unsafe { rb_profile_frames(0, BUFF_LEN as i32, frames_buffer.as_mut_ptr(), lines_buffer.as_mut_ptr()) };
+
+        let mut i = num - 1;
+        let yjit_raw_samples = YjitExitLocations::get_raw_samples();
+        let yjit_line_samples = YjitExitLocations::get_line_samples();
+
+        yjit_raw_samples.push(VALUE(num as usize));
+        yjit_line_samples.push(num);
+
+        while i >= 0 {
+            let frame = frames_buffer[i as usize];
+            let line = lines_buffer[i as usize];
+
+            yjit_raw_samples.push(frame);
+            yjit_line_samples.push(line);
+
+            i -= 1;
+        }
+
+        // Push the insn value into the yjit_raw_samples Vec.
+        yjit_raw_samples.push(VALUE(insn as usize));
+
+        // Push the current line onto the yjit_line_samples Vec. This
+        // points to the line in insns.def.
+        let line = yjit_line_samples.len() - 1;
+        yjit_line_samples.push(line as i32);
+
+        yjit_raw_samples.push(VALUE(1 as usize));
+        yjit_line_samples.push(1);
+    }
+}
+
 /// Primitive called in yjit.rb. Zero out all the counters.
 #[no_mangle]
 pub extern "C" fn rb_yjit_reset_stats_bang(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 06e428f34a..192e9753d9 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -3,6 +3,7 @@ use crate::core::*;
 use crate::cruby::*;
 use crate::invariants::*;
 use crate::options::*;
+use crate::stats::YjitExitLocations;
 
 use std::os::raw;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -54,6 +55,7 @@ pub extern "C" fn rb_yjit_init_rust() {
     let result = std::panic::catch_unwind(|| {
         Invariants::init();
         CodegenGlobals::init();
+        YjitExitLocations::init();
 
         // YJIT enabled and initialized successfully
         YJIT_ENABLED.store(true, Ordering::Release);