ruby · maximecb · Jun 9, 2022 · May 26, 2022 · maximecb · Jun 8, 2022
diff --git a/test/ruby/test_yjit_exit_locations.rb 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
@@ -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
@@ -18,11 +18,111 @@ def self.stats_enabled?
     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
@@ -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
@@ -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
@@ -967,12 +967,27 @@ extern "C" {
 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;
 }