What is the status of JIT compilation for OCaml today?

Back in the day there were various attempts to write a JIT compiler for OCaml. Amongst other things this promises a much faster top level.

There was the ocaml-jit project but that was x86 and I have since moved to x64 and then Aarch64 (M1). Then, 14 years ago, there was the ocamljit2 project that targeted OCaml v3.12 and touted huge benefits.

What is the status of JIT compilation for OCaml today?

3 Likes

If I remember correctly, ocamljit2 was not maintained.

Currently there is an in-progress effort to revive ocamlnat and make it easier for native applications to programmatically compile code fragments and link them into the currently running program: Fast native toplevel using JIT · Pull Request #15 · ocaml/RFCs · GitHub .

5 Likes

That’s great, thanks!

I was actually banking on being able to use natdynlink in a future project. Good to know it isn’t working!

natdynlink is working - it’s more that ocamlnat (the native toplevel for OCaml) has been receiving some considerable attention.

1 Like

Just a small thing I’d like to add: I extended @dra27’s ocaml-jit to support dynlinking directly from OCaml source without a toplevel: Allow dynlink/eval usage by copy · Pull Request #5 · NathanReb/ocaml-jit · GitHub The patch is pretty small, mostly moving usages of the toplevel api out of the library and exposing functions to add symbols that can be used by the evaluated code. An eval example looks like this: ocaml-jit/eval.ml at 57ef5c906681562186565ff28eeb49eb8e05527e · NathanReb/ocaml-jit · GitHub

I think it’s a much lighter alternative to ocaml_plugin, which embeds the OCaml compiler in a tar archive. Feedback is welcome.

5 Likes

Brilliant. That’s actually much closer to what I was hoping for. Thanks!

Hi,

I am wondering what would be the efforts to JIT an evaluation loop with some constant arguments?
I would expect it to be simpler and have less dependencies than full OCaml compilation code. Maybe it already exists off-the-shelf in the compiler.

Here is a basic example:

type op = A | B

let eval_a x = x + 1
let eval_b x = x + x

let rec eval x = function
 | [] -> x
 | A :: l -> eval (eval_a x) l
 | B :: l -> eval (eval_b x) l

let const_eval x = eval x [ B; A; B; A ]

With JIT, I would like to achieve what would have been the result of the compilation of the following:

let const_eval x =
  let x1 = eval_b x in
  let x2 = eval_a x1 in
  let x3 = eval_b x2 in
  eval_a x3

I think it would mainly require three primitives – lambda, apply and return – in order to create custom jit functions:

type 'a t
val lambda : ('a t -> 'b t) -> ('a -> 'b) t
val apply : ('a -> 'b) -> 'a t -> 'b t
val return : 'a t -> 'a
let comp_eval l : int -> int =
  let rec unroll l x = match l with
    | [] -> x
    | A :: l -> unroll l (apply eval_a x)
    | B :: l -> unroll l (apply eval_b x) in
  return (lambda (unroll l)) 
let const_eval = comp_eval [ B; A; B; A ]

Regards,

1 Like

You should have a look at MetaOCaml which is basically an implementation of this idea – staged programming, when you first explicitly construct program fragments and then you “run” them.

6 Likes

This is what I was looking for, thank you @gasche.

Would be awesome if it also exposes an API through standalone library, but I will definitely start playing with 4.11.1+BER.

1 Like

It would be awesome to hear about progress and plans regarding BER MetaOCaml.

1 Like

I was just thinking the same thing! I started Googling MetaOCaml last night because I’m in need of some staged metaprogramming and, last I tried it, MetaOCaml was freaking awesome.

Crazy idea: what if there was a jit intrinsic that took a closure and gave you a JIT compiled version with the body of the function partially specialized over the closure’s environment? Wouldn’t that offer more benefit with less complexity?

Hi @gasche,

I tried MetaOCaml (the latest version 4.14.1+BER), mainly for comparison purpose, but I am facing a hard limitation: native code generator does not support CSP (cross stage persistency) for non trivial values (according to BER MetaOCaml). It is said that it may be implemented in the future if there is a demand for it, but (maybe I am blind) I do not see any information on the way to contact the authors. Do you know who they are and how to contact them?

Also, since the byte code version should be able to handle complex CSP, I wanted to try byte code compilation, but I ran into the issue that compiler-libs has only the cmx files installed in the switch while I need (I guess) the cmo. Do you know how I can install the bytecode version of this library? Is there a flag to pass to opam?

Regards,
Frédéric

1 Like

You don’t need the .cmo, as the information is contained in the corresponding cma archive. The following works on my machine:

$ cat test.ml
let () =
  "foo"
  |> Pparse.preprocess
  |> print_endline

$ ocamlfind ocamlopt -package compiler-libs.common -linkpkg -o test test.ml

$ ./test
foo

I don’t know the details of what MetaOCaml expects in terms of native-compilation dependencies (maybe compiler-libs.optcomp?) and how to integrate it in the build process, but above is a minimal usage of linking a native program to compiler-libs.

1 Like

Thank you. You are right, I have to look for .cma files.
At the end, I solved the compilation issue by adding manually the following dependencies in the dune file:

compiler-libs.common
compiler-libs.bytecomp
compiler-libs.toplevel

That does not mean it works since I got a runtime error (Env.error in file "typing/env.ml") but it is another problem. Quand ça veut pas, ça veut pas! I think I should try with a much simpler code ¯\_(ツ)_/¯

I know that Oleg Kiselyov (whose website you are linking to) has co-authored most of papers using MetaOCaml for a number of years, although I don’t know whether he maintains it. His contact details are in his papers.

Is it specifically the runtime aspect of JIT that you are interested in, or rather staged programming? If the latter, modular macros may be of interest. They have many similarities with MetaOCaml (and they have (a form of) CSP).

1 Like

Hi @otini,

I am more interested in the runtime aspect of JIT.
In fact, last year, I have written a small JIT library, mainly inspired by ocaml-jit that provides the three primitives lambda, apply and return I mentioned above.
I am planning to release it by the end of the year.

I developed it mainly because it brought me fun to do it by myself. I made (what I believe to be) some optimizations and now, I am looking for related works to compare against to see if they were just a premature optimization or if they are game changer.

1 Like

A solution that I’m guessing is in between modular macros and MetaOCaml on the macros-jit spectrum is ppx_stage.

Please ping me on release, please. I’m interested does it scales to more complicated languages.

Hi @Kakadu, for sure, I will keep you informed of progress.

For now, it was designed to be really simple (PoC prototype) and is restricted on purpose to const (value from the current OCaml runtime) and apply (delayed application).
This was making the front end (the library) and the back end (x86 assembler and linker) very small (as I said, it can be viewed as a very specialized fork of ocaml-jit).

It would be possible to add some primitives too for basic type manipulation (bool, int, etc.). I also have some plans to add some form of if-then-else, but it will not be for quite a while.