user0/rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Directly invoke the main function in JIT mode

Browse source 
Fixes #1151
This commit is contained in:
bjorn3 2021-03-17 13:56:54 +01:00
parent 218b383f8f
commit e86b95480f
1 changed files with 52 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

74
src/driver/jit.rs
View file

@ -7,6 +7,7 @@ use std::os::raw::{c_char, c_int};
use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;
use rustc_session::config::EntryFnType;
use cranelift_jit::{JITBuilder, JITModule};
@ -32,16 +33,6 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
let mut jit_module = JITModule::new(jit_builder);
assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
let sig = Signature {
params: vec![
AbiParam::new(jit_module.target_config().pointer_type()),
AbiParam::new(jit_module.target_config().pointer_type()),
],
returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
};
let main_func_id = jit_module.declare_function("main", Linkage::Import, &sig).unwrap();
let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
let mono_items = cgus
.iter()
@ -86,24 +77,17 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
tcx.sess.fatal("Inline asm is not supported in JIT mode");
}
crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context);
crate::allocator::codegen(tcx, &mut jit_module, &mut unwind_context);
tcx.sess.abort_if_errors();
jit_module.finalize_definitions();
let _unwind_register_guard = unsafe { unwind_context.register_jit(&jit_module) };
let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);
println!(
"Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
);
let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
unsafe { ::std::mem::transmute(finalized_main) };
let args = ::std::env::var("CG_CLIF_JIT_ARGS").unwrap_or_else(|_| String::new());
let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
.chain(args.split(' '))
@ -118,12 +102,58 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
BACKEND_CONFIG.with(|tls_backend_config| {
assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
});
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());
let (main_def_id, entry_ty) = tcx.entry_fn(LOCAL_CRATE).unwrap();
let instance = Instance::mono(tcx, main_def_id.to_def_id()).polymorphize(tcx);
std::process::exit(ret);
match entry_ty {
EntryFnType::Main => {
// FIXME set program arguments somehow
let main_sig = Signature {
params: vec![],
returns: vec![],
call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
};
let main_func_id = jit_module
.declare_function(tcx.symbol_name(instance).name, Linkage::Import, &main_sig)
.unwrap();
let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);
CURRENT_MODULE.with(|current_module| {
assert!(current_module.borrow_mut().replace(jit_module).is_none())
});
let f: extern "C" fn() = unsafe { ::std::mem::transmute(finalized_main) };
f();
std::process::exit(0);
}
EntryFnType::Start => {
let start_sig = Signature {
params: vec![
AbiParam::new(jit_module.target_config().pointer_type()),
AbiParam::new(jit_module.target_config().pointer_type()),
],
returns: vec![AbiParam::new(
jit_module.target_config().pointer_type(), /*isize*/
)],
call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
};
let start_func_id = jit_module
.declare_function(tcx.symbol_name(instance).name, Linkage::Import, &start_sig)
.unwrap();
let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);
CURRENT_MODULE.with(|current_module| {
assert!(current_module.borrow_mut().replace(jit_module).is_none())
});
let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
unsafe { ::std::mem::transmute(finalized_start) };
let ret = f(args.len() as c_int, argv.as_ptr());
std::process::exit(ret);
}
}
}
#[no_mangle]
@ -220,7 +250,7 @@ fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
imported_symbols
}
pub(super) fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
let tcx = cx.tcx;
let pointer_type = cx.module.target_config().pointer_type();