RJIT: Implement --rjit-trace-exits

lib/ruby_vm/rjit/exit_compiler.rb

@@ -6,12 +6,12 @@ def initialize = freeze
     # @param pc [Integer]
     # @param asm [RubyVM::RJIT::Assembler]
     def compile_entry_exit(pc, ctx, asm, cause:)
-      # Increment per-insn exit counter
-      incr_insn_exit(pc, asm)
-
       # Fix pc/sp offsets for the interpreter
       save_pc_and_sp(pc, ctx, asm, reset_sp_offset: false)
 
+      # Increment per-insn exit counter
+      count_insn_exit(pc, asm)
+
       # Restore callee-saved registers
       asm.comment("#{cause}: entry exit")
       asm.pop(SP)
@@ -62,12 +62,12 @@ def compile_full_cfunc_return(asm)
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
     def compile_side_exit(pc, ctx, asm)
-      # Increment per-insn exit counter
-      incr_insn_exit(pc, asm)
-
       # Fix pc/sp offsets for the interpreter
       save_pc_and_sp(pc, ctx.dup, asm) # dup to avoid sp_offset update
 
+      # Increment per-insn exit counter
+      count_insn_exit(pc, asm)
+
       # Restore callee-saved registers
       asm.comment("exit to interpreter on #{pc_to_insn(pc).name}")
       asm.pop(SP)
@@ -105,13 +105,18 @@ def pc_to_insn(pc)
 
     # @param pc [Integer]
     # @param asm [RubyVM::RJIT::Assembler]
-    def incr_insn_exit(pc, asm)
+    def count_insn_exit(pc, asm)
       if C.rjit_opts.stats
         insn = Compiler.decode_insn(C.VALUE.new(pc).*)
         asm.comment("increment insn exit: #{insn.name}")
         asm.mov(:rax, (C.rjit_insn_exits + insn.bin).to_i)
         asm.add([:rax], 1) # TODO: lock
       end
+      if C.rjit_opts.trace_exits
+        asm.comment('rjit_record_exit_stack')
+        asm.mov(C_ARGS[0], pc)
+        asm.call(C.rjit_record_exit_stack)
+      end
     end
 
     # @param jit [RubyVM::RJIT::JITState]

lib/ruby_vm/rjit/stats.rb

@@ -30,6 +30,7 @@ def self.runtime_stats
   class << self
     private
 
+    # --yjit-stats at_exit
     def print_stats
       stats = runtime_stats
       $stderr.puts("***RJIT: Printing RJIT statistics on exit***")
@@ -98,5 +99,87 @@ def format_number(pad, number)
       with_commas = d_groups.map(&:join).join(',').reverse
       [with_commas, decimal].compact.join('.').rjust(pad, ' ')
     end
+
+    # --yjit-trace-exits at_exit
+    def dump_trace_exits
+      filename = "#{Dir.pwd}/rjit_exit_locations.dump"
+      File.binwrite(filename, Marshal.dump(exit_traces))
+      $stderr.puts("RJIT exit locations dumped to:\n#{filename}")
+    end
+
+    # Convert rb_rjit_raw_samples and rb_rjit_line_samples into a StackProf format.
+    def exit_traces
+      results = C.rjit_exit_traces
+      raw_samples = results[:raw].dup
+      line_samples = results[:lines].dup
+      frames = results[:frames].dup
+      samples_count = 0
+
+      # 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, lines: {} }
+        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.
+      #
+      # Raw Samples:
+      # [ length, frame1, frame2, frameN, ..., instruction, count
+      #
+      # Line Samples
+      # [ length, line_1, line_2, line_n, ..., dummy value, count
+      i = 0
+      while i < raw_samples.length
+        stack_length = raw_samples[i] + 1
+        i += 1 # consume the stack length
+
+        prev_frame_id = nil
+        stack_length.times do |idx|
+          idx += i
+          frame_id = raw_samples[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] += 1
+
+          frame_info[:lines][line_samples[idx]] ||= [0, 0]
+          frame_info[:lines][line_samples[idx]][0] += 1
+
+          prev_frame_id = frame_id
+        end
+
+        i += stack_length # consume the stack
+
+        top_frame_id = prev_frame_id
+        top_frame_line = 1
+
+        sample_count = raw_samples[i]
+
+        frames[top_frame_id][:samples] += sample_count
+        frames[top_frame_id][:lines] ||= {}
+        frames[top_frame_id][:lines][top_frame_line] ||= [0, 0]
+        frames[top_frame_id][:lines][top_frame_line][1] += sample_count
+
+        samples_count += sample_count
+        i += 1
+      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
   end
 end

rjit.c

@@ -67,7 +67,10 @@ struct rjit_options rb_rjit_opts;
 
 // true if RJIT is enabled.
 bool rb_rjit_enabled = false;
+// true if --rjit-stats (used before rb_rjit_opts is set)
 bool rb_rjit_stats_enabled = false;
+// true if --rjit-trace-exits (used before rb_rjit_opts is set)
+bool rb_rjit_trace_exits_enabled = false;
 // true if JIT-ed code should be called. When `ruby_vm_event_enabled_global_flags & ISEQ_TRACE_EVENTS`
 // and `rb_rjit_call_p == false`, any JIT-ed code execution is cancelled as soon as possible.
 bool rb_rjit_call_p = false;
@@ -93,6 +96,11 @@ static VALUE rb_cRJITCfpPtr = 0;
 // RubyVM::RJIT::Hooks
 static VALUE rb_mRJITHooks = 0;
 
+// Frames for --rjit-trace-exits
+VALUE rb_rjit_raw_samples = 0;
+// Line numbers for --rjit-trace-exits
+VALUE rb_rjit_line_samples = 0;
+
 // A default threshold used to add iseq to JIT.
 #define DEFAULT_CALL_THRESHOLD 30
 // Size of executable memory block in MiB.
@@ -113,6 +121,9 @@ rb_rjit_setup_options(const char *s, struct rjit_options *rjit_opt)
     else if (opt_match_noarg(s, l, "stats")) {
         rjit_opt->stats = true;
     }
+    else if (opt_match_noarg(s, l, "trace-exits")) {
+        rjit_opt->trace_exits = true;
+    }
     else if (opt_match_arg(s, l, "call-threshold")) {
         rjit_opt->call_threshold = atoi(s + 1);
     }
@@ -136,6 +147,7 @@ rb_rjit_setup_options(const char *s, struct rjit_options *rjit_opt)
 const struct ruby_opt_message rb_rjit_option_messages[] = {
 #if RJIT_STATS
     M("--rjit-stats",              "", "Enable collecting RJIT statistics"),
+    M("--rjit-trace-exits",        "", "Trace side exit locations"),
 #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) ")"),
@@ -314,6 +326,8 @@ rb_rjit_mark(void)
     rb_gc_mark(rb_cRJITIseqPtr);
     rb_gc_mark(rb_cRJITCfpPtr);
     rb_gc_mark(rb_mRJITHooks);
+    rb_gc_mark(rb_rjit_raw_samples);
+    rb_gc_mark(rb_rjit_line_samples);
 
     RUBY_MARK_LEAVE("rjit");
 }
@@ -398,6 +412,10 @@ rb_rjit_init(const struct rjit_options *opts)
     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"));
+    if (rb_rjit_opts.trace_exits) {
+        rb_rjit_raw_samples = rb_ary_new();
+        rb_rjit_line_samples = rb_ary_new();
+    }
 
     // Enable RJIT and stats from here
     rb_rjit_call_p = !rb_rjit_opts.pause;
@@ -408,13 +426,20 @@ rb_rjit_init(const struct rjit_options *opts)
 // Primitive for rjit.rb
 //
 
-// Same as `RubyVM::RJIT::C.enabled?`, but this is used before rjit_init.
+// Same as `rb_rjit_opts.stats`, but this is used before rb_rjit_opts is set.
 static VALUE
 rjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self)
 {
     return RBOOL(rb_rjit_stats_enabled);
 }
 
+// Same as `rb_rjit_opts.trace_exits`, but this is used before rb_rjit_opts is set.
+static VALUE
+rjit_trace_exits_enabled_p(rb_execution_context_t *ec, VALUE self)
+{
+    return RBOOL(rb_rjit_trace_exits_enabled);
+}
+
 // Disable anything that could impact stats. It ends up disabling JIT calls as well.
 static VALUE
 rjit_stop_stats(rb_execution_context_t *ec, VALUE self)

rjit.h

@@ -32,6 +32,8 @@ struct rjit_options {
     unsigned int exec_mem_size;
     // Collect RJIT statistics
     bool stats;
+    // Trace side exit locations
+    bool trace_exits;
     // Enable disasm of all JIT code
     bool dump_disasm;
     // [experimental] Do not start RJIT until RJIT.resume is called.
@@ -69,6 +71,7 @@ extern void rb_rjit_collect_vm_usage_insn(int insn);
 
 extern bool rb_rjit_enabled;
 extern bool rb_rjit_stats_enabled;
+extern bool rb_rjit_trace_exits_enabled;
 
 # else // USE_RJIT
 
@@ -88,6 +91,7 @@ static inline void rb_rjit_tracing_invalidate_all(rb_event_flag_t new_iseq_event
 #define rb_rjit_enabled false
 #define rb_rjit_call_p false
 #define rb_rjit_stats_enabled false
+#define rb_rjit_trace_exits_enabled false
 
 #define rb_rjit_call_threshold() UINT_MAX
 

rjit.rb

@@ -18,6 +18,12 @@ def self.resume
       print_stats
     end
   end
+  if Primitive.rjit_trace_exits_enabled_p
+    at_exit do
+      Primitive.rjit_stop_stats
+      dump_trace_exits
+    end
+  end
 end
 
 if RubyVM::RJIT.enabled?