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 }