BYOJ: Build your Own JIT (RubyConfTH 2023)

1. BYOJ: Build your own JIT @Faraaz98
2. What is a compiler? Ahead-of-Time compiler Just-in-Time compiler C, C++, Rust, etc Ruby, Python, Java, etc Compile, Then Execute Compile While Executing
3. The many JITs of Ruby MJIT YJIT First ever JIT for Ruby. Introduced in Ruby 2.6 To be discontinued RJIT New first-class JIT released with Ruby 3.0 Default JIT written in Ruby, shipping with Ruby 3.3 Experimental only
4. https://browse.arxiv.org/abs/1411.0352
5. The many JITs of Ruby MJIT YJIT First ever JIT for Ruby. Introduced in Ruby 2.6 To be discontinued RJIT New first-class JIT released with Ruby 3.0 Default JIT written in Ruby, shipping with Ruby 3.3 Experimental only
6. What is a machine?
7. Physical Machine Virtual Machine CPU Virtual CPU Core i9, Apple M2 Ruby interpreter (YARV)
8. YARV (Yet Another Ruby VM) A Stack machine: Stores/retrieves all your objects onto/from the stack 1 + 2 SP push 1 push 2 add PC PC: Program counter SP: Stack pointer
9. YARV (Yet Another Ruby VM) A Stack machine: Stores/retrieves all your objects onto/from the stack 1 SP 1 + 2 push 1 push 2 add PC
10. YARV (Yet Another Ruby VM) A Stack machine: Stores/retrieves all your objects onto/from the stack 1 2 SP 1 + 2 push 1 push 2 add PC
11. YARV (Yet Another Ruby VM) A Stack machine: Stores/retrieves all your objects onto/from the stack SP 1 + 2 push 1 push 2 add PC 2, 1
12. YARV (Yet Another Ruby VM) A Stack machine: Stores/retrieves all your objects onto/from the stack 3 SP 1 + 2 push 1 push 2 add PC
13. CPU A Register machine: Stores/retrieves all your values onto/from registers 1 + 2 mov r1, 1 mov r2, 2 add r1, r2 PC Register Value r1 1 r2 r3
14. CPU A Register machine: Stores/retrieves all your values onto/from registers 1 + 2 mov r1, 1 mov r2, 2 add r1, r2 PC Register Value r1 1 r2 2 r3
15. CPU A Register machine: Stores/retrieves all your values onto/from registers 1 + 2 mov r1, 1 mov r2, 2 add r1, r2 PC Register Value r1 3 r2 2 r3
16. Compiling nil 1. Put nil onto the stack 2. return nil SP def foo nil end
17. Compiling nil 1. Put nil onto the stack 2. return nil nil SP def foo nil end
18. Compiling nil 1. Put nil onto the stack 2. return nil nil SP WHERE DOES THE nil GO?? def foo nil end
19. Instruction sequence $ ruby --dump=insns foo.rb Looking under the hood == disasm: #<ISeq:<main>foo.rb:1 (1,0)-(3,3)> 0000 definemethod :foo, foo 0003 putobject :foo 0005 leave == disasm: #<ISeq:foo@foo.rb:1 (1,0)-(3,3)> 0000 putnil 0001 leave
20. nil SP 0000 putnil 0001 leave
21. nil SP WHERE DO WE leave FROM?? 0000 putnil 0001 leave
22. nil SP WHERE DO WE leave FROM?? WHERE DO WE leave TO?? 0000 putnil 0001 leave
23. Stack traces • Shows the exact path from which you errored out • Traces it all the way to the top level: <main> def foo raise StandardError end def bar foo end bar frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>'
24. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now
25. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now 2. Where in the program you are in right now
26. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return
27. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return
28. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return
29. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return
30. A data structure that knows the following: frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>' 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return 1. Which method you are in right now 2. Where in the program you are in right now 3. What value you need to return Needs to be stacked: 1. So Every method has its own copy of context 2. Can be unrolled when a complete backtrace needs to be shown
31. Control frames rb_control_frame_t [TOP] rb_control_frame_t [EVAL] rb_control_frame_t [METHOD:bar] CFP rb_control_frame_t [METHOD:foo] Represent the path taken by your program pc sp self ... rb_control_frame_t Program counter Stack pointer frames.rb:2:in `foo': StandardError from frames.rb:6:in `bar' from frames.rb:9:in `<main>'
32. Increment CFP to leave == disasm: #<ISeq:<main>@main.rb:1 (1,0)-(1,3)> (catch: FALSE) 0000 putnil 0001 leave Increment Control Frame Pointer
33. Building your own JIT (This time for real) Using RJIT
34. Using RJIT # Replace RJIT with JIT::Compiler RubyVM::RJIT::Compiler.prepend(Module.new { def compile(iseq, _) @compiler ||= JIT::Compiler.new @compiler.compile(iseq) end }) # Enable JIT compilation (paused by --rjit=pause) RubyVM::RJIT.resume
35. Building our own JIT # Replace RJIT with JIT::Compiler RubyVM::RJIT::Compiler.prepend(Module.new { def compile(iseq, _) @compiler ||= JIT::Compiler.new @compiler.compile(iseq) end }) # Enable JIT compilation (paused by --rjit=pause) RubyVM::RJIT.resume We write the JIT::Compiler class
36. The JIT::Compiler class module JIT class Compiler # Utilities to call C functions and interact with the Ruby VM. C = RubyVM::RJIT::C # Metadata for each YARV instruction. INSNS = RubyVM::RJIT::INSNS # Compile a method def compile(iseq) # Write machine code to this assembler. asm = Assembler.new # Iterate over each YARV instruction. insn_index = 0 while insn_index < iseq.body.iseq_size # Compile our instructions here end end end end
37. The JIT::Compiler class module JIT class Compiler # Utilities to call C functions and interact with the Ruby VM. C = RubyVM::RJIT::C # Metadata for each YARV instruction. INSNS = RubyVM::RJIT::INSNS # Compile a method def compile(iseq) # Write machine code to this assembler. asm = Assembler.new # Iterate over each YARV instruction. insn_index = 0 while insn_index < iseq.body.iseq_size # Compile our instructions here end end end end # Utilities to call C functions and interact with the Ruby VM. C = RubyVM::RJIT::C # Metadata for each YARV instruction. INSNS = RubyVM::RJIT::INSNS
38. The JIT::Compiler class module JIT class Compiler # Utilities to call C functions and interact with the Ruby VM. C = RubyVM::RJIT::C # Metadata for each YARV instruction. INSNS = RubyVM::RJIT::INSNS # Compile a method def compile(iseq) # Write machine code to this assembler. asm = Assembler.new # Iterate over each YARV instruction. insn_index = 0 while insn_index < iseq.body.iseq_size # Compile our instructions here end end end end def compile(iseq) # Write machine code to this assembler. asm = Assembler.new # Iterate over each YARV instruction. insn_index = 0 while insn_index < iseq.body.iseq_size # Compile our instructions here end end
39. Compiling putnil 0000 putnil 0001 leave def foo nil end
40. Compiling putnil # Iterate over each YARV instruction. insn_index = 0 while insn_index < iseq.body.iseq_size insn = INSNS.fetch( C.rb_vm_insn_decode(iseq.body.iseq_encoded[insn_index]) ) case insn.name in :putnil # ... end insn_index += insn.len end
41. Compiling putnil STACK = [:r8, :r9] stack_size = 0 while insn_index < iseq.body.iseq_size # ... case insn.name # ... in :putnil asm.mov(STACK[stack_size], C.to_value(nil)) stack_size += 1 end end :r8 :r9 SP nil r8 register
42. Compiling leave 0000 putnil 0001 leave def foo nil end
43. Compiling leave EC = :rdi CFP = :rsi while insn_index < iseq.body.iseq_size case insn.name # ... in :leave asm.add(CFP, C.rb_control_frame_t.size) asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP) end end
44. Compiling leave asm.add( CFP, C.rb_control_frame_t.size ) rb_control_frame_t [TOP] rb_control_frame_t [EVAL] rb_control_frame_t [METHOD:bar] CFP 0000 0007 0008 0015 0016 0024
45. Compiling leave rb_control_frame_t [TOP] rb_control_frame_t [EVAL] rb_control_frame_t [METHOD:bar] CFP 0000 0007 0008 0015 0016 0024 “Pop a control frame from the stack” asm.add( CFP, C.rb_control_frame_t.size )
46. Compiling leave asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP) rb_control_frame_t [TOP] rb_control_frame_t [EVAL] rb_control_frame_t [METHOD:bar] CFP pc sp self ... rb_control_frame_t
47. Compiling leave rb_control_frame_t [TOP] rb_control_frame_t [EVAL] rb_control_frame_t [METHOD:bar] CFP “Set execution context to current control frame” pc sp self ... rb_control_frame_t asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP)
48. Compiling leave EC = :rdi CFP = :rsi while insn_index < iseq.body.iseq_size case insn.name # ... in :leave asm.add(CFP, C.rb_control_frame_t.size) asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP) asm.mov(:rax, STACK[stack_size - 1]) asm.ret end end
49. Compiling leave EC = :rdi CFP = :rsi while insn_index < iseq.body.iseq_size case insn.name # ... in :leave asm.add(CFP, C.rb_control_frame_t.size) asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP) asm.mov(:rax, STACK[stack_size - 1]) asm.ret end end All return values go into :rax
50. Compiling 2 + 3 def five 2 + 3 end 0000 putobject 2 0002 putobject 3 0004 opt_plus 0006 leave
51. Compiling 2 + 3 0000 putobject 2 0002 putobject 3 0004 opt_plus 0006 leave PC SP
52. Compiling 2 + 3 2 PC SP 0000 putobject 2 0002 putobject 3 0004 opt_plus 0006 leave
53. Compiling 2 + 3 2 3 PC SP 0000 putobject 2 0002 putobject 3 0004 opt_plus 0006 leave
54. Compiling 2 + 3 5 PC SP 0000 putobject 2 0002 putobject 3 0004 opt_plus 0006 leave
55. Compiling putobject while insn_index < iseq.body.iseq_size case insn.name # ... in :putobject operand = iseq.body.iseq_encoded[insn_index + 1] asm.mov(STACK[stack_size], operand) stack_size += 1 end end
56. Compiling opt_plus while insn_index < iseq.body.iseq_size case insn.name # ... in :opt_plus recv = STACK[stack_size - 2] obj = STACK[stack_size - 1] asm.add(recv, obj) stack_size -= 1 end end 2 3 SP :r8 :r9
57. Compiling opt_plus while insn_index < iseq.body.iseq_size case insn.name # ... in :opt_plus recv = STACK[stack_size - 2] obj = STACK[stack_size - 1] asm.add(recv, obj) stack_size -= 1 end end 5 3 SP :r8 :r9
58. and that’s it!
59. Benchmarks (Higher is better) rubybench.github.io
60. One More Thing
61. def sum 2 + 3 end puts sum ➜ ruby plus.rb 420
62. Bonus - Monkey patching at the machine level while insn_index < iseq.body.iseq_size case insn.name # ... in :opt_plus recv = STACK[stack_size - 2] obj = STACK[stack_size - 1] # asm.add(recv, obj) asm.mov(recv, C.to_value(420)) stack_size -= 1 end end
63. Credits Takashi Kokubun @k0kubun Github: k0kubun/ruby-jit-challenge (github.com/k0kubun/ruby-jit-challenge)
64. Thank you!