]> git.lizzy.rs Git - rust.git/blobdiff - src/driver/jit.rs
Use an ItemId inside mir::GlobalAsm.
[rust.git] / src / driver / jit.rs
index fd3a628922d0bc5ffbae730dc97febdd81c0d6d3..f784d8d27cc7423de24c787030cf680637214ec4 100644 (file)
@@ -1,15 +1,26 @@
 //! 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};
 
 use rustc_codegen_ssa::CrateInfo;
+use rustc_middle::mir::mono::MonoItem;
+
+use cranelift_jit::{JITBuilder, JITModule};
 
 use crate::prelude::*;
+use crate::{CodegenCx, CodegenMode};
+
+thread_local! {
+    pub static CURRENT_MODULE: RefCell<Option<JITModule>> = RefCell::new(None);
+}
 
-pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
-    use cranelift_simplejit::{SimpleJITBackend, SimpleJITBuilder};
+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`.");
+    }
 
     #[cfg(unix)]
     unsafe {
@@ -31,12 +42,13 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
 
     let imported_symbols = load_imported_symbols_for_jit(tcx);
 
-    let mut jit_builder = SimpleJITBuilder::with_isa(
-        crate::build_isa(tcx.sess, false),
+    let mut jit_builder = JITBuilder::with_isa(
+        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: Module<SimpleJITBackend> = Module::new(jit_builder);
+    let mut jit_module = JITModule::new(jit_builder);
     assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
 
     let sig = Signature {
@@ -53,10 +65,6 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
         .declare_function("main", Linkage::Import, &sig)
         .unwrap();
 
-    if !tcx.sess.opts.output_types.should_codegen() {
-        tcx.sess.fatal("JIT mode doesn't work with `cargo check`.");
-    }
-
     let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
     let mono_items = cgus
         .iter()
@@ -66,35 +74,60 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
         .into_iter()
         .collect::<Vec<(_, (_, _))>>();
 
-    let mut cx = crate::CodegenCx::new(tcx, jit_module, false);
+    let mut cx = crate::CodegenCx::new(tcx, jit_module, false, false);
+
+    super::time(tcx, "codegen mono items", || {
+        super::predefine_mono_items(&mut cx, &mono_items);
+        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) => 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);
+                }
+                MonoItem::GlobalAsm(item_id) => {
+                    let item = cx.tcx.hir().item(item_id);
+                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
+                }
+            }
+        }
+    });
 
     let (mut jit_module, global_asm, _debug, mut unwind_context) =
-        super::time(tcx, "codegen mono items", || {
-            super::codegen_mono_items(&mut cx, mono_items);
-            tcx.sess.time("finalize CodegenCx", || cx.finalize())
-        });
+        tcx.sess.time("finalize CodegenCx", || cx.finalize());
+    jit_module.finalize_definitions();
+
     if !global_asm.is_empty() {
-        tcx.sess.fatal("Global asm is not supported in JIT mode");
+        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::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context, true);
     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(&mut jit_module) };
-
-    tcx.sess.abort_if_errors();
+    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 the CG_CLIF_JIT env var is set");
+    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(" "))
+        .chain(args.split(' '))
         .map(|arg| CString::new(arg).unwrap())
         .collect::<Vec<_>>();
     let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
@@ -103,12 +136,46 @@ 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());
 
-    jit_module.finish();
     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 = tcx.symbol_name(instance).name.to_string();
+            let sig = crate::abi::get_function_sig(tcx, cx.module.isa().triple(), instance);
+            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;
 
@@ -141,18 +208,18 @@ fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
 
     let mut imported_symbols = Vec::new();
     for path in dylib_paths {
-        use object::Object;
+        use object::{Object, ObjectSymbol};
         let lib = libloading::Library::new(&path).unwrap();
         let obj = std::fs::read(path).unwrap();
         let obj = object::File::parse(&obj).unwrap();
-        imported_symbols.extend(obj.dynamic_symbols().filter_map(|(_idx, symbol)| {
+        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
             let name = symbol.name().unwrap().to_string();
             if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                 return None;
             }
             let dlsym_name = if cfg!(target_os = "macos") {
                 // On macOS `dlsym` expects the name without leading `_`.
-                assert!(name.starts_with("_"), "{:?}", name);
+                assert!(name.starts_with('_'), "{:?}", name);
                 &name[1..]
             } else {
                 &name
@@ -168,3 +235,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 = tcx.symbol_name(inst).name.to_string();
+    let sig = crate::abi::get_function_sig(tcx, cx.module.isa().triple(), inst);
+    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();
+}