Expose methods for inspecting Micro JIT code blocks

[tenderlove]

This commit adds a module UJIT. The module allows you to insert the
initial Micro JIT instruction in to an arbitrary iseq like this:

def foo(x)
  if x < 1
    "less than one"
  else
    "something else"
  end
end

iseq = RubyVM::InstructionSequence.of(method(:foo))

UJIT.insert(iseq) # Add initial jump

After the initial jump is added, we can make Micro JIT do some work:

100.times { foo(0) }

The UJIT module also exposes a method for finding all compiled blocks
for a given iseq, like this:

blocks = UJIT.blocks_for(iseq)

We can sort the blocks by address and use the Crabstone gem (which is a
wrapper around capstone) to disassemble the generated code.

Here is the full code example:

def foo(x)
  if x < 1
    "less than one"
  else
    "something else"
  end
end

iseq = RubyVM::InstructionSequence.of(method(:foo))

UJIT.insert(iseq) # Add initial jump

100.times { foo(0) }

blocks = UJIT.blocks_for(iseq)

 # brew install capstone
 # gem install crabstone
require "crabstone"

cs = Crabstone::Disassembler.new(Crabstone::ARCH_X86, Crabstone::MODE_64)

puts iseq.disasm

blocks.sort_by(&:address).reverse.each do |block|
  puts "== ISEQ RANGE: #{block.iseq_start_index} -> #{block.iseq_end_index} ".ljust(80, "=")
  cs.disasm(block.code, 0).each do |i|
    printf(
      "\t0x%<address>x:\t%<instruction>s\t%<details>s\n",
      address: i.address,
      instruction: i.mnemonic,
      details: i.op_str
    )
  end
end

Here is the output:

$ ./ruby test.rb
== disasm: #<ISeq:[email protected]:1 (1,0)-(7,3)> (catch: FALSE)
local table (size: 1, argc: 1 [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])
[ 1] x@0<Arg>
0000 getlocal_WC_0                          x@0                       (   2)[LiCa]
0002 putobject_INT2FIX_1_
0003 opt_lt                                 <calldata!mid:<, argc:1, ARGS_SIMPLE>
0005 branchunless                           10
0007 putstring                              "less than one"           (   3)[Li]
0009 leave                                                            (   7)[Re]
0010 putstring                              "something else"          (   5)[Li]
0012 leave                                                            (   7)[Re]
== ISEQ RANGE: 7 -> 7 ==========================================================
	0x0:	movabs	rax, 0x7fcd014cd518
	0xa:	mov	qword ptr [rdi], rax
	0xd:	mov	r8, rax
	0x10:	mov	r9, rax
	0x13:	mov	r11, r12
	0x16:	jmp	qword ptr [rax]
== ISEQ RANGE: 0 -> 7 ==========================================================
	0x0:	mov	rax, qword ptr [rdi + 0x20]
	0x4:	mov	rax, qword ptr [rax - 0x18]
	0x8:	mov	qword ptr [rdx], rax
	0xb:	mov	qword ptr [rdx + 8], 3
	0x13:	movabs	rax, 0x7fcd0180ac00
	0x1d:	test	byte ptr [rax + 0x3e6], 1
	0x24:	jne	0x3ffe0da
	0x2a:	test	byte ptr [rdx], 1
	0x2d:	je	0x3ffe0da
	0x33:	test	byte ptr [rdx + 8], 1
	0x37:	je	0x3ffe0da
	0x3d:	mov	rax, qword ptr [rdx]
	0x40:	cmp	rax, qword ptr [rdx + 8]
	0x44:	movabs	rax, 0
	0x4e:	movabs	rcx, 0x14
	0x58:	cmovl	rax, rcx
	0x5c:	mov	qword ptr [rdx], rax
	0x5f:	test	qword ptr [rdx], -9
	0x66:	je	0x3ffe111
	0x6c:	jmp	0xffffffffffffffa3

[XrXr]

I should fix the CI builds...

[maximecb]

I like what you've done. It seems like a good idea to have a ujit module with an API for this and various things. I think that will serve us well in the future when we want to write tests and add more introspection.

I have 3 small requests and a question:

• Can we name the module lowercase ujit instead of UJIT ?
• Can we name .insert just install_entry or attach_entry?
• I'd rather have the code in ujit_ruby.c be in ujit_iface.c. The goal of the ujit_iface.c file was to hold things that relate to our interface to the outside.

Question: you wrote code to produce a dissassembly with crabstone. How can we package this code so that it's easy to use and people don't have to look for it? Is there any way we can add method written in Ruby to the ujit module that does the disasm of a method or iseq? This can be a separate pull request if you prefer.

Ideally if we could do something like ujit.disasm(method or iseq) that would be super convenient.

[tenderlove]

• Can we name the module lowercase ujit instead of UJIT ?

We can name it Ujit. The first letter of a Ruby constant must be uppercase (which is why I just did UJIT). Given the constraints of constant names in Ruby, what would you prefer?

• Can we name .insert just install_entry or attach_entry?

Yep, I'll change that

• I'd rather have the code in ujit_ruby.c be in ujit_iface.c. The goal of the ujit_iface.c file was to hold things that relate to our interface to the outside.

Sure, that's fine. I can move it in that file.

Question: you wrote code to produce a dissassembly with crabstone. How can we package this code so that it's easy to use and people don't have to look for it? Is there any way we can add method written in Ruby to the ujit module that does the disasm of a method or iseq? This can be a separate pull request if you prefer.

IMO we should package that as an external gem (maybe a micro jit gem?). I don't think it's a good idea to add a dependency to a disassembler to Ruby. For now, I think we could just check in a script for us to use while debugging, but eventually I think we need a gem or library that exists outside Ruby.

[maximecb]

IMO we should package that as an external gem (maybe a micro jit gem?). I don't think it's a good idea to add a dependency to a disassembler to Ruby. For now, I think we could just check in a script for us to use while debugging, but eventually I think we need a gem or library that exists outside Ruby.

I agree that we shouldn't add a dependency to Ruby. Though we do need some kind of way to package the code in a convenient way.

In Python it's easy to lazily import a module, or let the import fail and catch the exception. We could have ujit.disasm() lazily import the disassembler lib right when the method is called, and just print an error message (please install this lib) and return nil if not found. That way we can kind of have the best of both worlds. People can just install that lib manually if they need disasm?

Alternatively, some custom script we can easily require would be OK too. I'll let you pick the solution you prefer.

[maximecb]

We can name it Ujit. The first letter of a Ruby constant must be uppercase (which is why I just did UJIT). Given the constraints of constant names in Ruby, what would you prefer?

Hawkward 🦅 Hmmm. I guess UJIT is OK.

[tenderlove]

I added a script under misc as a helper for disassembly. Here is an example of how to use the script:

require "misc/ujit_disasm"

def foo(x)
  if x < 1
    "less than one"
  else
    "something else"
  end
end

iseq = RubyVM::InstructionSequence.of(method(:foo))

puts UJIT.disasm(iseq) # returns nil because nothing has been generated

100.times { foo 1 }

puts UJIT.disasm(iseq) # returns stuff because the JIT executed

Then in the shell:

$ ./ruby -I. test.rb

== disasm: #<ISeq:[email protected]:3 (3,0)-(9,3)> (catch: FALSE)
local table (size: 1, argc: 1 [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])
[ 1] x@0<Arg>
0000 getlocal_WC_0                          x@0                       (   4)[LiCa]
0002 putobject_INT2FIX_1_
0003 opt_lt                                 <calldata!mid:<, argc:1, ARGS_SIMPLE>
0005 branchunless                           10
0007 putstring                              "less than one"           (   5)[Li]
0009 leave                                                            (   9)[Re]
0010 putstring                              "something else"          (   7)[Li]
0012 leave                                                            (   9)[Re]
== ISEQ RANGE: 10 -> 10 ========================================================
	0x0:	movabs	rax, 0x7fc38f5d17b0
	0xa:	mov	qword ptr [rdi], rax
	0xd:	mov	r8, rax
	0x10:	mov	r9, rax
	0x13:	mov	r11, r12
	0x16:	jmp	qword ptr [rax]
== ISEQ RANGE: 0 -> 7 ==========================================================
	0x0:	mov	rax, qword ptr [rdi + 0x20]
	0x4:	mov	rax, qword ptr [rax - 0x18]
	0x8:	mov	qword ptr [rdx], rax
	0xb:	mov	qword ptr [rdx + 8], 3
	0x13:	movabs	rax, 0x7fc390009400
	0x1d:	test	byte ptr [rax + 0x3e6], 1
	0x24:	jne	0x3ffca7a
	0x2a:	test	byte ptr [rdx], 1
	0x2d:	je	0x3ffca7a
	0x33:	test	byte ptr [rdx + 8], 1
	0x37:	je	0x3ffca7a
	0x3d:	mov	rax, qword ptr [rdx]
	0x40:	cmp	rax, qword ptr [rdx + 8]
	0x44:	movabs	rax, 0
	0x4e:	movabs	rcx, 0x14
	0x58:	cmovl	rax, rcx
	0x5c:	mov	qword ptr [rdx], rax
	0x5f:	test	qword ptr [rdx], -9
	0x66:	jne	0x3ffcad4

[maximecb]

Nice, very well done!