Remove support for the lazy jit mode

I might re-implement it in the future, but would probably do so by
replacing cranelift-jit. cranelift-jit's api doesn't quite work well for
lazy jitting. And even so it adds complexity to cranelift-jit and breaks
cranelift-jit outside of x86_64.

Author: bjorn3

build_system/tests.rs

@@ -384,20 +384,6 @@ impl<'a> TestRunner<'a> {
                         jit_cmd.env("CG_CLIF_JIT_ARGS", args);
                     }
                     spawn_and_wait(jit_cmd);
-
-                    eprintln!("[JIT-lazy] {testname}");
-                    let mut jit_cmd = self.rustc_command([
-                        "-Zunstable-options",
-                        "-Cllvm-args=mode=jit-lazy",
-                        "-Cprefer-dynamic",
-                        source,
-                        "--cfg",
-                        "jit",
-                    ]);
-                    if !args.is_empty() {
-                        jit_cmd.env("CG_CLIF_JIT_ARGS", args);
-                    }
-                    spawn_and_wait(jit_cmd);
                 }
             }
         }

docs/usage.md

@@ -41,20 +41,13 @@ or
 $ $cg_clif_dir/dist/rustc-clif -Cllvm-args=jit-mode -Cprefer-dynamic my_crate.rs
 ```
 
-There is also an experimental lazy jit mode. In this mode functions are only compiled once they are
-first called.
-
-```bash
-$ $cg_clif_dir/dist/cargo-clif lazy-jit
-```
-
 ## Shell
 
 These are a few functions that allow you to easily run rust code from the shell using cg_clif as jit.
 
 ```bash
 function jit_naked() {
-    echo "$@" | $cg_clif_dir/dist/rustc-clif - -Zunstable-options -Cllvm-args=jit-mode -Cllvm-args=jit-lazy -Cprefer-dynamic
+    echo "$@" | $cg_clif_dir/dist/rustc-clif - -Zunstable-options -Cllvm-args=jit-mode-Cprefer-dynamic
 }
 
 function jit() {

scripts/cargo-clif.rs

@@ -54,19 +54,6 @@ fn main() {
                 ])
                 .collect()
         }
-        Some("lazy-jit") => {
-            rustflags.push("-Cprefer-dynamic".to_owned());
-            args.remove(0);
-            IntoIterator::into_iter(["rustc".to_string()])
-                .chain(args)
-                .chain([
-                    "--".to_string(),
-                    "-Zunstable-options".to_string(),
-                    "-Cllvm-args=jit-mode".to_string(),
-                    "-Cllvm-args=jit-lazy".to_string(),
-                ])
-                .collect()
-        }
         _ => args,
     };
 

src/config.rs

@@ -6,11 +6,6 @@ pub struct BackendConfig {
     /// Defaults to AOT compilation. Can be set using `-Cllvm-args=jit-mode`.
     pub jit_mode: bool,
 
-    /// When JIT executing should the lazy JIT mode be used.
-    ///
-    /// Defaults to false. Can be set using `-Cllvm-args=jit-lazy`.
-    pub lazy_jit: bool,
-
     /// When JIT mode is enable pass these arguments to the program.
     ///
     /// Defaults to the value of `CG_CLIF_JIT_ARGS`.
@@ -22,7 +17,6 @@ impl BackendConfig {
     pub fn from_opts(opts: &[String]) -> Result<Self, String> {
         let mut config = BackendConfig {
             jit_mode: false,
-            lazy_jit: false,
             jit_args: match std::env::var("CG_CLIF_JIT_ARGS") {
                 Ok(args) => args.split(' ').map(|arg| arg.to_string()).collect(),
                 Err(std::env::VarError::NotPresent) => vec![],
@@ -40,7 +34,6 @@ impl BackendConfig {
             }
             match &**opt {
                 "jit-mode" => config.jit_mode = true,
-                "jit-lazy" => config.lazy_jit = true,
                 _ => return Err(format!("Unknown option `{}`", opt)),
             }
         }

src/driver/jit.rs

@@ -1,12 +1,9 @@
 //! The JIT driver uses [`cranelift_jit`] 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 std::sync::{Mutex, OnceLock, mpsc};
 
-use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
 use cranelift_jit::{JITBuilder, JITModule};
 use rustc_codegen_ssa::CrateInfo;
 use rustc_middle::mir::mono::MonoItem;
@@ -18,58 +15,13 @@ use crate::debuginfo::TypeDebugContext;
 use crate::prelude::*;
 use crate::unwind_module::UnwindModule;
 
-struct JitState {
-    jit_module: UnwindModule<JITModule>,
-}
-
-thread_local! {
-    static LAZY_JIT_STATE: RefCell<Option<JitState>> = const { RefCell::new(None) };
-}
-
-/// The Sender owned by the rustc thread
-static GLOBAL_MESSAGE_SENDER: OnceLock<Mutex<mpsc::Sender<UnsafeMessage>>> = OnceLock::new();
-
-/// A message that is sent from the jitted runtime to the rustc thread.
-/// Senders are responsible for upholding `Send` semantics.
-enum UnsafeMessage {
-    /// Request that the specified `Instance` be lazily jitted.
-    ///
-    /// Nothing accessible through `instance_ptr` may be moved or mutated by the sender after
-    /// this message is sent.
-    JitFn {
-        instance_ptr: *const Instance<'static>,
-        trampoline_ptr: *const u8,
-        tx: mpsc::Sender<*const u8>,
-    },
-}
-unsafe impl Send for UnsafeMessage {}
-
-impl UnsafeMessage {
-    /// Send the message.
-    fn send(self) {
-        thread_local! {
-            /// The Sender owned by the local thread
-            static LOCAL_MESSAGE_SENDER: mpsc::Sender<UnsafeMessage> =
-                GLOBAL_MESSAGE_SENDER
-                    .get().unwrap()
-                    .lock().unwrap()
-                    .clone();
-        }
-        LOCAL_MESSAGE_SENDER.with(|sender| {
-            sender.send(self).expect("rustc thread hung up before lazy JIT request was sent")
-        })
-    }
-}
-
-fn create_jit_module(tcx: TyCtxt<'_>, hotswap: bool) -> (UnwindModule<JITModule>, CodegenCx) {
+fn create_jit_module(tcx: TyCtxt<'_>) -> (UnwindModule<JITModule>, CodegenCx) {
     let crate_info = CrateInfo::new(tcx, "dummy_target_cpu".to_string());
 
     let isa = crate::build_isa(tcx.sess);
     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.symbol_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info));
-    jit_builder.symbol("__clif_jit_fn", clif_jit_fn as *const u8);
     let mut jit_module = UnwindModule::new(JITModule::new(jit_builder), false);
 
     let cx = crate::CodegenCx::new(tcx, jit_module.isa(), false, sym::dummy_cgu_name);
@@ -79,7 +31,7 @@ fn create_jit_module(tcx: TyCtxt<'_>, hotswap: bool) -> (UnwindModule<JITModule>
     (jit_module, cx)
 }
 
-pub(crate) fn run_jit(tcx: TyCtxt<'_>, jit_lazy: bool, jit_args: Vec<String>) -> ! {
+pub(crate) fn run_jit(tcx: TyCtxt<'_>, jit_args: Vec<String>) -> ! {
     if !tcx.sess.opts.output_types.should_codegen() {
         tcx.dcx().fatal("JIT mode doesn't work with `cargo check`");
     }
@@ -88,7 +40,7 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, jit_lazy: bool, jit_args: Vec<String>) ->
         tcx.dcx().fatal("can't jit non-executable crate");
     }
 
-    let (mut jit_module, mut cx) = create_jit_module(tcx, jit_lazy);
+    let (mut jit_module, mut cx) = create_jit_module(tcx);
     let mut cached_context = Context::new();
 
     let cgus = tcx.collect_and_partition_mono_items(()).codegen_units;
@@ -105,17 +57,13 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, jit_lazy: bool, jit_args: Vec<String>) ->
         for (mono_item, _) in mono_items {
             match mono_item {
                 MonoItem::Fn(inst) => {
-                    if jit_lazy {
-                        codegen_shim(tcx, &mut cached_context, &mut jit_module, inst)
-                    } else {
-                        codegen_and_compile_fn(
-                            tcx,
-                            &mut cx,
-                            &mut cached_context,
-                            &mut jit_module,
-                            inst,
-                        );
-                    }
+                    codegen_and_compile_fn(
+                        tcx,
+                        &mut cx,
+                        &mut cached_context,
+                        &mut jit_module,
+                        inst,
+                    );
                 }
                 MonoItem::Static(def_id) => {
                     crate::constant::codegen_static(tcx, &mut jit_module, def_id);
@@ -158,41 +106,17 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, jit_lazy: bool, jit_args: Vec<String>) ->
     let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
     let finalized_start: *const u8 = jit_module.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 { jit_module });
-    });
-
     let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
         unsafe { ::std::mem::transmute(finalized_start) };
 
-    let (tx, rx) = mpsc::channel();
-    GLOBAL_MESSAGE_SENDER.set(Mutex::new(tx)).unwrap();
-
-    // Spawn the jitted runtime in a new thread so that this rustc thread can handle messages
-    // (eg to lazily JIT further functions as required)
-    std::thread::spawn(move || {
-        let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
+    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
 
-        // Push a null pointer as a terminating argument. This is required by POSIX and
-        // useful as some dynamic linkers use it as a marker to jump over.
-        argv.push(std::ptr::null());
+    // Push a null pointer as a terminating argument. This is required by POSIX and
+    // useful as some dynamic linkers use it as a marker to jump over.
+    argv.push(std::ptr::null());
 
-        let ret = f(args.len() as c_int, argv.as_ptr());
-        std::process::exit(ret);
-    });
-
-    // Handle messages
-    loop {
-        match rx.recv().unwrap() {
-            // lazy JIT compilation request - compile requested instance and return pointer to result
-            UnsafeMessage::JitFn { instance_ptr, trampoline_ptr, tx } => {
-                tx.send(jit_fn(instance_ptr, trampoline_ptr))
-                    .expect("jitted runtime hung up before response to lazy JIT request was sent");
-            }
-        }
-    }
+    let ret = f(args.len() as c_int, argv.as_ptr());
+    std::process::exit(ret);
 }
 
 pub(crate) fn codegen_and_compile_fn<'tcx>(
@@ -224,58 +148,6 @@ pub(crate) fn codegen_and_compile_fn<'tcx>(
     });
 }
 
-extern "C" fn clif_jit_fn(
-    instance_ptr: *const Instance<'static>,
-    trampoline_ptr: *const u8,
-) -> *const u8 {
-    // send the JIT request to the rustc thread, with a channel for the response
-    let (tx, rx) = mpsc::channel();
-    UnsafeMessage::JitFn { instance_ptr, trampoline_ptr, tx }.send();
-
-    // block on JIT compilation result
-    rx.recv().expect("rustc thread hung up before responding to sent lazy JIT request")
-}
-
-fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) -> *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();
-
-        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 name = tcx.symbol_name(instance).name;
-            let sig = crate::abi::get_function_sig(
-                tcx,
-                jit_module.target_config().default_call_conv,
-                instance,
-            );
-            let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
-
-            let current_ptr = jit_module.module.read_got_entry(func_id);
-
-            // If the function's GOT entry has already been updated to point at something other
-            // than the shim trampoline, don't re-jit but just return the new pointer instead.
-            // This does not need synchronization as this code is executed only by a sole rustc
-            // thread.
-            if current_ptr != trampoline_ptr {
-                return current_ptr;
-            }
-
-            jit_module.module.prepare_for_function_redefine(func_id).unwrap();
-
-            let mut cx = crate::CodegenCx::new(tcx, jit_module.isa(), false, sym::dummy_cgu_name);
-            codegen_and_compile_fn(tcx, &mut cx, &mut Context::new(), jit_module, instance);
-
-            assert!(cx.global_asm.is_empty());
-            jit_module.finalize_definitions();
-            jit_module.module.get_finalized_function(func_id)
-        })
-    })
-}
-
 fn dep_symbol_lookup_fn(
     sess: &Session,
     crate_info: CrateInfo,
@@ -323,57 +195,3 @@ fn dep_symbol_lookup_fn(
         None
     })
 }
-
-fn codegen_shim<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    cached_context: &mut Context,
-    module: &mut UnwindModule<JITModule>,
-    inst: Instance<'tcx>,
-) {
-    let pointer_type = module.target_config().pointer_type();
-
-    let name = tcx.symbol_name(inst).name;
-    let sig = crate::abi::get_function_sig(tcx, module.target_config().default_call_conv, inst);
-    let func_id = module.declare_function(name, Linkage::Export, &sig).unwrap();
-
-    let instance_ptr = Box::into_raw(Box::new(inst));
-
-    let jit_fn = module
-        .declare_function(
-            "__clif_jit_fn",
-            Linkage::Import,
-            &Signature {
-                call_conv: module.target_config().default_call_conv,
-                params: vec![AbiParam::new(pointer_type), AbiParam::new(pointer_type)],
-                returns: vec![AbiParam::new(pointer_type)],
-            },
-        )
-        .unwrap();
-
-    let context = cached_context;
-    context.clear();
-    let trampoline = &mut context.func;
-    trampoline.signature = sig.clone();
-
-    let mut builder_ctx = FunctionBuilderContext::new();
-    let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);
-
-    let trampoline_fn = module.declare_func_in_func(func_id, 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();
-    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 trampoline_ptr = trampoline_builder.ins().func_addr(pointer_type, trampoline_fn);
-    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr, trampoline_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);
-
-    module.define_function(func_id, context).unwrap();
-}

src/lib.rs

@@ -215,7 +215,7 @@ impl CodegenBackend for CraneliftCodegenBackend {
         });
         if config.jit_mode {
             #[cfg(feature = "jit")]
-            driver::jit::run_jit(tcx, config.lazy_jit, config.jit_args);
+            driver::jit::run_jit(tcx, config.jit_args);
 
             #[cfg(not(feature = "jit"))]
             tcx.dcx().fatal("jit support was disabled when compiling rustc_codegen_cranelift");