From 3f6a3b5ebece5f6d3b1010b68a5fd86756593098 Mon Sep 17 00:00:00 2001 From: bjorn3 Date: Fri, 25 Dec 2020 11:41:48 +0100 Subject: [PATCH] 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. --- example/std_example.rs | 2 + scripts/tests.sh | 6 ++ src/constant.rs | 3 +- src/driver/jit.rs | 129 ++++++++++++++++++++++++++++++++++++++--- src/driver/mod.rs | 4 +- src/lib.rs | 2 + src/vtable.rs | 3 +- 7 files changed, 138 insertions(+), 11 deletions(-) diff --git a/example/std_example.rs b/example/std_example.rs index b38e25328a4..015bbdfed46 100644 --- a/example/std_example.rs +++ b/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!"); }); diff --git a/scripts/tests.sh b/scripts/tests.sh index 60ed7aac8aa..a61774f479e 100755 --- a/scripts/tests.sh +++ b/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 diff --git a/src/constant.rs b/src/constant.rs index 544b020b711..bfa2076d3d9 100644 --- a/src/constant.rs +++ b/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); } diff --git a/src/driver/jit.rs b/src/driver/jit.rs index cb386ea83b9..9a42c675cc1 100644 --- a/src/driver/jit.rs +++ b/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> = 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(); +} diff --git a/src/driver/mod.rs b/src/driver/mod.rs index 49bb79c101a..9f4ea9a3865 100644 --- a/src/driver/mod.rs +++ b/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 diff --git a/src/lib.rs b/src/lib.rs index 91e1efb71b9..6e4f3bf2898 100644 --- a/src/lib.rs +++ b/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)), } } diff --git a/src/vtable.rs b/src/vtable.rs index 238abc0d8bd..8f15586a9dc 100644 --- a/src/vtable.rs +++ b/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 }