Implement lazy compilation in JIT mode

Lazy compilation has the potential to significantly improve the startup
time of a program. While functions have to be codegened when called, it
is expected that a significant amount of all code is only required when
an error occurs or only when the program is used in certain ways.

The basic approach is to first codegen a shim for each function. This
shim calls the `__cg_clif_jit` function of cg_clif with a pointer to the
`Instance` corresponding to the function for which it is a shim.
`__cg_clif_jit` function then codegens this function and uses the hot
code swapping support of SimpleJIT to redirect future calls to the
function to the real version. Finally it calls the newly codegened
function.

Author: bjorn3

example/std_example.rs

@@ -15,6 +15,8 @@ fn main() {
     let stderr = ::std::io::stderr();
     let mut stderr = stderr.lock();
 
+    // FIXME support lazy jit when multi threading
+    #[cfg(not(lazy_jit))]
     std::thread::spawn(move || {
         println!("Hello from another thread!");
     });

scripts/tests.sh

@@ -16,6 +16,9 @@ function no_sysroot_tests() {
     if [[ "$JIT_SUPPORTED" = "1" ]]; then
         echo "[JIT] mini_core_hello_world"
         CG_CLIF_JIT_ARGS="abc bcd" $MY_RUSTC -Cllvm-args=mode=jit -Cprefer-dynamic example/mini_core_hello_world.rs --cfg jit --target "$HOST_TRIPLE"
+
+        echo "[JIT-lazy] mini_core_hello_world"
+        CG_CLIF_JIT_ARGS="abc bcd" $MY_RUSTC -Cllvm-args=mode=jit-lazy -Cprefer-dynamic example/mini_core_hello_world.rs --cfg jit --target "$HOST_TRIPLE"
     else
         echo "[JIT] mini_core_hello_world (skipped)"
     fi
@@ -38,6 +41,9 @@ function base_sysroot_tests() {
     if [[ "$JIT_SUPPORTED" = "1" ]]; then
         echo "[JIT] std_example"
         $MY_RUSTC -Cllvm-args=mode=jit -Cprefer-dynamic example/std_example.rs --target "$HOST_TRIPLE"
+
+        echo "[JIT-lazy] std_example"
+        $MY_RUSTC -Cllvm-args=mode=jit-lazy -Cprefer-dynamic example/std_example.rs --cfg lazy_jit --target "$HOST_TRIPLE"
     else
         echo "[JIT] std_example (skipped)"
     fi

src/constant.rs

@@ -447,7 +447,8 @@ fn define_all_allocs(tcx: TyCtxt<'_>, module: &mut impl Module, cx: &mut Constan
             data_ctx.write_data_addr(offset.bytes() as u32, global_value, addend as i64);
         }
 
-        module.define_data(data_id, &data_ctx).unwrap();
+        // FIXME don't duplicate definitions in lazy jit mode
+        let _ = module.define_data(data_id, &data_ctx);
         cx.done.insert(data_id);
     }
 

src/driver/jit.rs

@@ -1,6 +1,7 @@
 //! The JIT driver uses [`cranelift_simplejit`] to JIT execute programs without writing any object
 //! files.
 
+use std::cell::RefCell;
 use std::ffi::CString;
 use std::os::raw::{c_char, c_int};
 
@@ -10,8 +11,13 @@ use rustc_middle::mir::mono::MonoItem;
 use cranelift_jit::{JITBuilder, JITModule};
 
 use crate::prelude::*;
+use crate::{CodegenCx, CodegenMode};
 
-pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
+thread_local! {
+    pub static CURRENT_MODULE: RefCell<Option<JITModule>> = RefCell::new(None);
+}
+
+pub(super) fn run_jit(tcx: TyCtxt<'_>, codegen_mode: CodegenMode) -> ! {
     if !tcx.sess.opts.output_types.should_codegen() {
         tcx.sess.fatal("JIT mode doesn't work with `cargo check`.");
     }
@@ -40,6 +46,7 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
         crate::build_isa(tcx.sess),
         cranelift_module::default_libcall_names(),
     );
+    jit_builder.hotswap(matches!(codegen_mode, CodegenMode::JitLazy));
     jit_builder.symbols(imported_symbols);
     let mut jit_module = JITModule::new(jit_builder);
     assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
@@ -74,13 +81,17 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
         for (mono_item, (linkage, visibility)) in mono_items {
             let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
             match mono_item {
-                MonoItem::Fn(inst) => {
-                    cx.tcx.sess.time("codegen fn", || {
-                        crate::base::codegen_fn(&mut cx, inst, linkage)
-                    });
-                }
+                MonoItem::Fn(inst) => match codegen_mode {
+                    CodegenMode::Aot => unreachable!(),
+                    CodegenMode::Jit => {
+                        cx.tcx.sess.time("codegen fn", || {
+                            crate::base::codegen_fn(&mut cx, inst, linkage)
+                        });
+                    }
+                    CodegenMode::JitLazy => codegen_shim(&mut cx, inst),
+                },
                 MonoItem::Static(def_id) => {
-                    crate::constant::codegen_static(&mut cx.constants_cx, def_id)
+                    crate::constant::codegen_static(&mut cx.constants_cx, def_id);
                 }
                 MonoItem::GlobalAsm(hir_id) => {
                     let item = cx.tcx.hir().expect_item(hir_id);
@@ -126,11 +137,50 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
     // useful as some dynamic linkers use it as a marker to jump over.
     argv.push(std::ptr::null());
 
+    CURRENT_MODULE
+        .with(|current_module| assert!(current_module.borrow_mut().replace(jit_module).is_none()));
+
     let ret = f(args.len() as c_int, argv.as_ptr());
 
     std::process::exit(ret);
 }
 
+#[no_mangle]
+extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
+    rustc_middle::ty::tls::with(|tcx| {
+        // lift is used to ensure the correct lifetime for instance.
+        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();
+
+        CURRENT_MODULE.with(|jit_module| {
+            let mut jit_module = jit_module.borrow_mut();
+            let jit_module = jit_module.as_mut().unwrap();
+            let mut cx = crate::CodegenCx::new(tcx, jit_module, false, false);
+
+            let (name, sig) = crate::abi::get_function_name_and_sig(
+                tcx,
+                cx.module.isa().triple(),
+                instance,
+                true,
+            );
+            let func_id = cx
+                .module
+                .declare_function(&name, Linkage::Export, &sig)
+                .unwrap();
+            cx.module.prepare_for_function_redefine(func_id).unwrap();
+
+            tcx.sess.time("codegen fn", || {
+                crate::base::codegen_fn(&mut cx, instance, Linkage::Export)
+            });
+
+            let (jit_module, global_asm, _debug_context, unwind_context) = cx.finalize();
+            assert!(global_asm.is_empty());
+            jit_module.finalize_definitions();
+            std::mem::forget(unsafe { unwind_context.register_jit(&jit_module) });
+            jit_module.get_finalized_function(func_id)
+        })
+    })
+}
+
 fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
     use rustc_middle::middle::dependency_format::Linkage;
 
@@ -190,3 +240,68 @@ fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
 
     imported_symbols
 }
+
+pub(super) fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx, impl Module>, inst: Instance<'tcx>) {
+    let tcx = cx.tcx;
+
+    let pointer_type = cx.module.target_config().pointer_type();
+
+    let (name, sig) =
+        crate::abi::get_function_name_and_sig(tcx, cx.module.isa().triple(), inst, true);
+    let func_id = cx
+        .module
+        .declare_function(&name, Linkage::Export, &sig)
+        .unwrap();
+
+    let instance_ptr = Box::into_raw(Box::new(inst));
+
+    let jit_fn = cx
+        .module
+        .declare_function(
+            "__clif_jit_fn",
+            Linkage::Import,
+            &Signature {
+                call_conv: cx.module.target_config().default_call_conv,
+                params: vec![AbiParam::new(pointer_type)],
+                returns: vec![AbiParam::new(pointer_type)],
+            },
+        )
+        .unwrap();
+
+    let mut trampoline = Function::with_name_signature(ExternalName::default(), sig.clone());
+    let mut builder_ctx = FunctionBuilderContext::new();
+    let mut trampoline_builder = FunctionBuilder::new(&mut trampoline, &mut builder_ctx);
+
+    let jit_fn = cx
+        .module
+        .declare_func_in_func(jit_fn, trampoline_builder.func);
+    let sig_ref = trampoline_builder.func.import_signature(sig);
+
+    let entry_block = trampoline_builder.create_block();
+    trampoline_builder.append_block_params_for_function_params(entry_block);
+    let fn_args = trampoline_builder
+        .func
+        .dfg
+        .block_params(entry_block)
+        .to_vec();
+
+    trampoline_builder.switch_to_block(entry_block);
+    let instance_ptr = trampoline_builder
+        .ins()
+        .iconst(pointer_type, instance_ptr as u64 as i64);
+    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
+    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
+    let call_inst = trampoline_builder
+        .ins()
+        .call_indirect(sig_ref, jitted_fn, &fn_args);
+    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
+    trampoline_builder.ins().return_(&ret_vals);
+
+    cx.module
+        .define_function(
+            func_id,
+            &mut Context::for_function(trampoline),
+            &mut cranelift_codegen::binemit::NullTrapSink {},
+        )
+        .unwrap();
+}

src/driver/mod.rs

@@ -23,7 +23,7 @@ pub(crate) fn codegen_crate(
 
     match config.codegen_mode {
         CodegenMode::Aot => aot::run_aot(tcx, metadata, need_metadata_module),
-        CodegenMode::Jit => {
+        CodegenMode::Jit | CodegenMode::JitLazy => {
             let is_executable = tcx
                 .sess
                 .crate_types()
@@ -33,7 +33,7 @@ pub(crate) fn codegen_crate(
             }
 
             #[cfg(feature = "jit")]
-            let _: ! = jit::run_jit(tcx);
+            let _: ! = jit::run_jit(tcx, config.codegen_mode);
 
             #[cfg(not(feature = "jit"))]
             tcx.sess

src/lib.rs

@@ -177,6 +177,7 @@ impl<'tcx, M: Module> CodegenCx<'tcx, M> {
 pub enum CodegenMode {
     Aot,
     Jit,
+    JitLazy,
 }
 
 impl Default for CodegenMode {
@@ -192,6 +193,7 @@ impl FromStr for CodegenMode {
         match s {
             "aot" => Ok(CodegenMode::Aot),
             "jit" => Ok(CodegenMode::Jit),
+            "jit-lazy" => Ok(CodegenMode::JitLazy),
             _ => Err(format!("Unknown codegen mode `{}`", s)),
         }
     }

src/vtable.rs

@@ -158,7 +158,8 @@ fn build_vtable<'tcx>(
         )
         .unwrap();
 
-    fx.cx.module.define_data(data_id, &data_ctx).unwrap();
+    // FIXME don't duplicate definitions in lazy jit mode
+    let _ = fx.cx.module.define_data(data_id, &data_ctx);
 
     data_id
 }