Merge pull request #202 from sampsyo/brilift-jit

Brilift: Enable JIT mode

brilift/src/main.rs

@@ -1,5 +1,8 @@
+mod rt;
+
 use argh::FromArgs;
 use bril_rs as bril;
+use core::mem;
 use cranelift_codegen::entity::EntityRef;
 use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::InstBuilder;
@@ -57,6 +60,15 @@ impl RTFunc {
             Self::PrintEnd => "_bril_print_end",
         }
     }
+
+    fn rt_impl(&self) -> *const u8 {
+        match self {
+            RTFunc::PrintInt => rt::print_int as *const u8,
+            RTFunc::PrintBool => rt::print_int as *const u8,
+            RTFunc::PrintSep => rt::print_sep as *const u8,
+            RTFunc::PrintEnd => rt::print_end as *const u8,
+        }
+    }
 }
 
 /// Runtime functions used in the native `main` function, which dispatches to the proper Bril
@@ -243,11 +255,50 @@ impl Translator<ObjectModule> {
     }
 }
 
+/// Run the JITted code.
+// TODO Support `main` arguments somehow.
+unsafe fn run(main_ptr: *const u8) {
+    let func = mem::transmute::<_, fn() -> ()>(main_ptr);
+    func();
+}
+
 /// JIT compiler that totally does not work yet.
 impl Translator<JITModule> {
+    // `cranelift_jit` does not yet support PIC on AArch64:
+    // https://github.com/bytecodealliance/wasmtime/issues/2735
+    // The default initialization path for `JITBuilder` is hard-coded to use PIC, so we manually
+    // disable it here. Once this is fully supported in `cranelift_jit`, we can switch to the
+    // generic versin below unconditionally.
+    #[cfg(target_arch = "aarch64")]
+    fn jit_builder() -> JITBuilder {
+        let mut flag_builder = settings::builder();
+        flag_builder.set("use_colocated_libcalls", "false").unwrap();
+        flag_builder.set("is_pic", "false").unwrap(); // PIC unsupported on ARM.
+        let isa_builder = cranelift_native::builder().unwrap();
+        let isa = isa_builder
+            .finish(settings::Flags::new(flag_builder))
+            .unwrap();
+        JITBuilder::with_isa(isa, cranelift_module::default_libcall_names())
+    }
+
+    // The normal way to set up a JIT builder.
+    #[cfg(not(target_arch = "aarch64"))]
+    fn jit_builder() -> JITBuilder {
+        JITBuilder::new(cranelift_module::default_libcall_names()).unwrap()
+    }
+
     fn new() -> Self {
-        // Cranelift JIT scaffolding.
-        let builder = JITBuilder::new(cranelift_module::default_libcall_names()).unwrap();
+        // Set up the JIT.
+        let mut builder = Self::jit_builder();
+
+        // Provide runtime functions.
+        enum_map! {
+            rtfunc => {
+                let f: RTFunc = rtfunc;
+                builder.symbol(f.name(), f.rt_impl());
+            }
+        };
+
         let mut module = JITModule::new(builder);
 
         Self {
@@ -258,11 +309,11 @@ impl Translator<JITModule> {
         }
     }
 
-    fn compile(mut self) -> *const u8 {
+    // Dispose of the translator and obtain the entry-point code pointer.
+    fn get_main(mut self) -> *const u8 {
         self.module.clear_context(&mut self.context);
         self.module.finalize_definitions();
 
-        // TODO Compile all functions.
         let id = self.funcs["main"];
         self.module.get_finalized_function(id)
     }
@@ -748,7 +799,8 @@ fn main() {
     if args.jit {
         let mut trans = Translator::<JITModule>::new();
         trans.compile_prog(prog, args.dump_ir, false);
-        trans.compile();
+        let code = trans.get_main();
+        unsafe { run(code) };
     } else {
         let mut trans = Translator::<ObjectModule>::new(args.target, &args.opt_level);
         trans.compile_prog(prog, args.dump_ir, true);

brilift/src/rt.rs

@@ -0,0 +1,19 @@
+#[no_mangle]
+pub extern "C" fn print_int(i: i64) {
+    print!("{}", i);
+}
+
+#[no_mangle]
+pub extern "C" fn print_bool(b: bool) {
+    print!("{}", b);
+}
+
+#[no_mangle]
+pub extern "C" fn print_sep() {
+    print!(" ");
+}
+
+#[no_mangle]
+pub extern "C" fn print_end() {
+    println!();
+}

docs/tools/brilift.md

@@ -1,12 +1,12 @@
 Cranelift Compiler
 ==================
 
-Brilift is a compiler from Bril to native code using the [Cranelift][] code generator.
+Brilift is a ahead-of-time or just-in-time compiler from Bril to native code using the [Cranelift][] code generator.
 It supports [core Bril][core] only.
 
-Brilift is an ahead-of-time compiler.
-It emits `.o` files and also provides a simple run-time library.
+In AOT mode, Brilift emits `.o` files and also provides a simple run-time library.
 By linking these together, you get a complete native executable.
+In JIT mode, Brilift mimics an interpreter.
 
 [cranelift]: https://github.com/bytecodealliance/wasmtime/tree/main/cranelift
 [core]: ../lang/core.md
@@ -25,8 +25,8 @@ You can build it using [Cargo][]:
 [bril-rs]: rust.md
 [cargo]: https://doc.rust-lang.org/cargo/
 
-Compile Stuff
--------------
+Ahead-of-Time Compilation
+-------------------------
 
 Provide the `brilift` executable with a Bril JSON program:
 
@@ -46,18 +46,27 @@ Then, you will want to link `rt.o` and `bril.o` to produce an executable:
 
 If your Bril `@main` function takes arguments, those are now command-line arguments to the `myprog` executable.
 
+Just-in-Time Compilation
+------------------------
+
+Use the `-j` flag to compile and run the program immediately:
+
+    $ bril2json < something.bril | brilift -j
+
 Options
 -------
 
 Type `brilift --help` to see the full list of options:
 
-* `-o <FILE>`: Place the output object file in `<FILE>` instead of `bril.o` (the default).
-* `-t <TARGET>`: Specify the target triple, as interpreted by Cranelift. These triples resemble the [target triples][triple] that LLVM also understands, for example. For instance, `x86_64-unknown-darwin-macho` is the triple for macOS on Intel processors.
+* `-j`: JIT-compile the code and run it immediately, instead of AOT-compiling an object file (the default).
 * `-O [none|speed|speed_and_size]`: An [optimization level][opt_level], according to Cranelift. The default is `none`.
 * `-v`: Enable lots of logging from the Cranelift library.
 * `-d`: Dump the Cranelift IR text for debugging.
 
-There is also a `-j` option that tries to use a JIT to run the code immediately, instead of the default AOT mode, but that does not work at all yet.
+These options are only relevant in AOT mode:
+
+* `-o <FILE>`: Place the output object file in `<FILE>` instead of `bril.o` (the default).
+* `-t <TARGET>`: Specify the target triple, as interpreted by Cranelift. These triples resemble the [target triples][triple] that LLVM also understands, for example. For instance, `x86_64-unknown-darwin-macho` is the triple for macOS on Intel processors.
 
 [opt_level]: https://docs.rs/cranelift-codegen/0.84.0/cranelift_codegen/settings/struct.Flags.html#method.opt_level
 [triple]: https://clang.llvm.org/docs/CrossCompilation.html#target-triple