RJIT: Implement --rjit-trace-exits

2 files changed, 95 insertions, 7 deletions

diff --git a/lib/ruby_vm/rjit/exit_compiler.rb b/lib/ruby_vm/rjit/exit_compiler.rb
index c082cc3660..b7beb22177 100644
--- a/lib/ruby_vm/rjit/exit_compiler.rb
+++ b/lib/ruby_vm/rjit/exit_compiler.rb
@@ -6,12 +6,12 @@ module RubyVM::RJIT
     # @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 @@ module RubyVM::RJIT
     # @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 @@ module RubyVM::RJIT
 
     # @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]
diff --git a/lib/ruby_vm/rjit/stats.rb b/lib/ruby_vm/rjit/stats.rb
index 7cef634991..2fde44bc8e 100644
--- a/lib/ruby_vm/rjit/stats.rb
+++ b/lib/ruby_vm/rjit/stats.rb
@@ -30,6 +30,7 @@ module RubyVM::RJIT
   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 @@ module RubyVM::RJIT
       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