]> git.lizzy.rs Git - rust.git/blobdiff - compiler/rustc_codegen_cranelift/src/driver/jit.rs
Auto merge of #84677 - CDirkx:pal, r=Mark-Simulacrum
[rust.git] / compiler / rustc_codegen_cranelift / src / driver / jit.rs
index dbe1ff083f0db24d12f9846c5cdf7147a4866bd1..53c93f6a9ddca82bbc39e246f253759ca373539a 100644 (file)
@@ -1,4 +1,4 @@
-//! The JIT driver uses [`cranelift_simplejit`] to JIT execute programs without writing any object
+//! The JIT driver uses [`cranelift_jit`] to JIT execute programs without writing any object
 //! files.
 
 use std::cell::RefCell;
 use crate::{prelude::*, BackendConfig};
 use crate::{CodegenCx, CodegenMode};
 
-thread_local! {
-    pub static BACKEND_CONFIG: RefCell<Option<BackendConfig>> = RefCell::new(None);
-    pub static CURRENT_MODULE: RefCell<Option<JITModule>> = RefCell::new(None);
+struct JitState {
+    backend_config: BackendConfig,
+    jit_module: JITModule,
 }
 
-pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
-    if !tcx.sess.opts.output_types.should_codegen() {
-        tcx.sess.fatal("JIT mode doesn't work with `cargo check`.");
-    }
+thread_local! {
+    static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
+}
 
+fn create_jit_module<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    backend_config: &BackendConfig,
+    hotswap: bool,
+) -> (JITModule, CodegenCx<'tcx>) {
     let imported_symbols = load_imported_symbols_for_jit(tcx);
 
-    let mut jit_builder =
-        JITBuilder::with_isa(crate::build_isa(tcx.sess), cranelift_module::default_libcall_names());
-    jit_builder.hotswap(matches!(backend_config.codegen_mode, CodegenMode::JitLazy));
+    let isa = crate::build_isa(tcx.sess, backend_config);
+    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
+    jit_builder.hotswap(hotswap);
     crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
     jit_builder.symbols(imported_symbols);
     let mut jit_module = JITModule::new(jit_builder);
-    assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
+
+    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);
+
+    crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
+    crate::main_shim::maybe_create_entry_wrapper(
+        tcx,
+        &mut jit_module,
+        &mut cx.unwind_context,
+        true,
+    );
+
+    (jit_module, cx)
+}
+
+pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
+    if !tcx.sess.opts.output_types.should_codegen() {
+        tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
+    }
+
+    if !tcx.sess.crate_types().contains(&rustc_session::config::CrateType::Executable) {
+        tcx.sess.fatal("can't jit non-executable crate");
+    }
+
+    let (mut jit_module, mut cx) = create_jit_module(
+        tcx,
+        &backend_config,
+        matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
+    );
 
     let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
     let mono_items = cgus
@@ -44,52 +75,45 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         .into_iter()
         .collect::<Vec<(_, (_, _))>>();
 
-    let mut cx = crate::CodegenCx::new(tcx, backend_config, &mut jit_module, false);
-
-    super::time(tcx, "codegen mono items", || {
-        super::predefine_mono_items(&mut cx, &mono_items);
+    super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
+        super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
         for (mono_item, _) in mono_items {
             match mono_item {
                 MonoItem::Fn(inst) => match backend_config.codegen_mode {
                     CodegenMode::Aot => unreachable!(),
                     CodegenMode::Jit => {
-                        cx.tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, inst));
+                        cx.tcx.sess.time("codegen fn", || {
+                            crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
+                        });
                     }
-                    CodegenMode::JitLazy => codegen_shim(&mut cx, inst),
+                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
                 },
                 MonoItem::Static(def_id) => {
-                    crate::constant::codegen_static(&mut cx.constants_cx, def_id);
+                    crate::constant::codegen_static(tcx, &mut jit_module, def_id);
                 }
                 MonoItem::GlobalAsm(item_id) => {
-                    let item = cx.tcx.hir().item(item_id);
+                    let item = tcx.hir().item(item_id);
                     tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                 }
             }
         }
     });
 
-    let (global_asm, _debug, mut unwind_context) =
-        tcx.sess.time("finalize CodegenCx", || cx.finalize());
-    jit_module.finalize_definitions();
-
-    if !global_asm.is_empty() {
+    if !cx.global_asm.is_empty() {
         tcx.sess.fatal("Inline asm is not supported in JIT mode");
     }
 
-    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) };
+    unsafe { cx.unwind_context.register_jit(&jit_module) };
 
     println!(
         "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
     );
 
-    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(backend_config.jit_args.iter().map(|arg| &**arg))
         .map(|arg| CString::new(arg).unwrap())
         .collect::<Vec<_>>();
     let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
@@ -98,61 +122,27 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
     // useful as some dynamic linkers use it as a marker to jump over.
     argv.push(std::ptr::null());
 
-    BACKEND_CONFIG.with(|tls_backend_config| {
-        assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
+    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("main", Linkage::Import, &start_sig).unwrap();
+    let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);
+
+    LAZY_JIT_STATE.with(|lazy_jit_state| {
+        let mut lazy_jit_state = lazy_jit_state.borrow_mut();
+        assert!(lazy_jit_state.is_none());
+        *lazy_jit_state = Some(JitState { backend_config, jit_module });
     });
 
-    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);
-
-    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);
-        }
-    }
+    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]
@@ -161,24 +151,23 @@ extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8
         // 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 backend_config =
-                BACKEND_CONFIG.with(|backend_config| backend_config.borrow().clone().unwrap());
+        LAZY_JIT_STATE.with(|lazy_jit_state| {
+            let mut lazy_jit_state = lazy_jit_state.borrow_mut();
+            let lazy_jit_state = lazy_jit_state.as_mut().unwrap();
+            let jit_module = &mut lazy_jit_state.jit_module;
+            let backend_config = lazy_jit_state.backend_config.clone();
 
-            let name = tcx.symbol_name(instance).name.to_string();
+            let name = tcx.symbol_name(instance).name;
             let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
-            let func_id = jit_module.declare_function(&name, Linkage::Export, &sig).unwrap();
+            let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
             jit_module.prepare_for_function_redefine(func_id).unwrap();
 
-            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module, false);
-            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, instance));
+            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module.isa(), false);
+            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));
 
-            let (global_asm, _debug_context, unwind_context) = cx.finalize();
-            assert!(global_asm.is_empty());
+            assert!(cx.global_asm.is_empty());
             jit_module.finalize_definitions();
-            std::mem::forget(unsafe { unwind_context.register_jit(&jit_module) });
+            unsafe { cx.unwind_context.register_jit(&jit_module) };
             jit_module.get_finalized_function(func_id)
         })
     })
@@ -248,35 +237,37 @@ fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
     imported_symbols
 }
 
-fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
+fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
     let tcx = cx.tcx;
 
-    let pointer_type = cx.module.target_config().pointer_type();
+    let pointer_type = 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 name = tcx.symbol_name(inst).name;
+    let sig = crate::abi::get_function_sig(tcx, module.isa().triple(), inst);
+    let func_id = module.declare_function(name, Linkage::Export, &sig).unwrap();
 
     let instance_ptr = Box::into_raw(Box::new(inst));
 
-    let jit_fn = cx
-        .module
+    let jit_fn = module
         .declare_function(
             "__clif_jit_fn",
             Linkage::Import,
             &Signature {
-                call_conv: cx.module.target_config().default_call_conv,
+                call_conv: 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());
+    cx.cached_context.clear();
+    let trampoline = &mut cx.cached_context.func;
+    trampoline.signature = sig.clone();
+
     let mut builder_ctx = FunctionBuilderContext::new();
-    let mut trampoline_builder = FunctionBuilder::new(&mut trampoline, &mut builder_ctx);
+    let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);
 
-    let jit_fn = cx.module.declare_func_in_func(jit_fn, trampoline_builder.func);
+    let jit_fn = 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();
@@ -291,10 +282,10 @@ fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
     let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
     trampoline_builder.ins().return_(&ret_vals);
 
-    cx.module
+    module
         .define_function(
             func_id,
-            &mut Context::for_function(trampoline),
+            &mut cx.cached_context,
             &mut NullTrapSink {},
             &mut NullStackMapSink {},
         )