Add -Zjit-mode flag and use it in cg_clif

This flag tells the codegen backend to immediately execute the code it
builds. This is exactly what miri needs as it is unable to produce real
executables. Adding a builtin flag would allow cargo to add native
support for this without miri needing to do a bunch of hacks in its
cargo wrapper. Furthermore cg_clif also has a jit mode that benefits
from it.

Author: bjorn3

compiler/rustc_codegen_cranelift/docs/usage.md

@@ -38,7 +38,7 @@ $ $cg_clif_dir/dist/cargo-clif jit
 or
 
 ```bash
-$ $cg_clif_dir/dist/rustc-clif -Cllvm-args=jit-mode -Cprefer-dynamic my_crate.rs
+$ $cg_clif_dir/dist/rustc-clif -Zjit-mode -Cprefer-dynamic my_crate.rs
 ```
 
 ## Shell
@@ -47,7 +47,7 @@ These are a few functions that allow you to easily run rust code from the shell
 
 ```bash
 function jit_naked() {
-    echo "$@" | $cg_clif_dir/dist/rustc-clif - -Zunstable-options -Cllvm-args=jit-mode -Cprefer-dynamic
+    echo "$@" | $cg_clif_dir/dist/rustc-clif - -Zunstable-options -Zjit-mode -Cprefer-dynamic
 }
 
 function jit() {

compiler/rustc_codegen_cranelift/scripts/cargo-clif.rs

@@ -51,11 +51,7 @@ fn main() {
             args.remove(0);
             IntoIterator::into_iter(["rustc".to_string()])
                 .chain(args)
-                .chain([
-                    "--".to_string(),
-                    "-Zunstable-options".to_string(),
-                    "-Cllvm-args=jit-mode".to_string(),
-                ])
+                .chain(["--".to_string(), "-Zjit-mode".to_string()])
                 .collect()
         }
         _ => args,

compiler/rustc_codegen_cranelift/scripts/filter_profile.rs

@@ -4,7 +4,7 @@
 pushd $(dirname "$0")/../
 RUSTC="$(pwd)/dist/rustc-clif"
 popd
-PROFILE=$1 OUTPUT=$2 exec $RUSTC -Zunstable-options -Cllvm-args=jit-mode -Cprefer-dynamic $0
+PROFILE=$1 OUTPUT=$2 exec $RUSTC -Zjit-mode -Cprefer-dynamic $0
 #*/
 
 //! This program filters away uninteresting samples and trims uninteresting frames for stackcollapse

compiler/rustc_codegen_cranelift/src/config.rs

@@ -3,6 +3,7 @@
 
 use std::ffi::CString;
 use std::os::raw::{c_char, c_int};
+use std::process::ExitCode;
 
 use cranelift_jit::{JITBuilder, JITModule};
 use rustc_codegen_ssa::CrateInfo;
@@ -18,12 +19,12 @@ use crate::unwind_module::UnwindModule;
 
 fn create_jit_module(
     tcx: TyCtxt<'_>,
-    crate_info: &CrateInfo,
+    crate_info: CrateInfo,
 ) -> (UnwindModule<JITModule>, Option<DebugContext>) {
     let isa = crate::build_isa(tcx.sess, true);
     let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
     crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
-    jit_builder.symbol_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info.clone()));
+    jit_builder.symbol_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info));
     let mut jit_module = UnwindModule::new(JITModule::new(jit_builder), false);
 
     let cx = DebugContext::new(tcx, jit_module.isa(), false, "dummy_cgu_name");
@@ -33,13 +34,10 @@ fn create_jit_module(
     (jit_module, cx)
 }
 
-pub(crate) fn run_jit(tcx: TyCtxt<'_>, crate_info: &CrateInfo, jit_args: Vec<String>) -> ! {
-    if !tcx.crate_types().contains(&rustc_session::config::CrateType::Executable) {
-        tcx.dcx().fatal("can't jit non-executable crate");
-    }
-
+pub(crate) fn run_jit(tcx: TyCtxt<'_>, target_cpu: String, jit_args: Vec<String>) -> ExitCode {
     let output_filenames = tcx.output_filenames(());
     let should_write_ir = crate::pretty_clif::should_write_ir(tcx.sess);
+    let crate_info = CrateInfo::new(tcx, target_cpu);
     let (mut jit_module, mut debug_context) = create_jit_module(tcx, crate_info);
     let mut cached_context = Context::new();
 
@@ -114,7 +112,7 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, crate_info: &CrateInfo, jit_args: Vec<Str
     argv.push(std::ptr::null());
 
     let ret = f(args.len() as c_int, argv.as_ptr());
-    std::process::exit(ret);
+    ExitCode::from(ret as u8)
 }
 
 fn codegen_and_compile_fn<'tcx>(

compiler/rustc_codegen_cranelift/src/driver/jit.rs

@@ -33,8 +33,8 @@ extern crate rustc_target;
 extern crate rustc_driver;
 
 use std::any::Any;
-use std::cell::OnceCell;
 use std::env;
+use std::process::ExitCode;
 use std::sync::Arc;
 
 use cranelift_codegen::isa::TargetIsa;
@@ -48,7 +48,6 @@ use rustc_session::config::OutputFilenames;
 use rustc_span::{Symbol, sym};
 use rustc_target::spec::{Abi, Arch, Env, Os};
 
-pub use crate::config::*;
 use crate::prelude::*;
 
 mod abi;
@@ -61,7 +60,6 @@ mod codegen_i128;
 mod common;
 mod compiler_builtins;
 mod concurrency_limiter;
-mod config;
 mod constant;
 mod debuginfo;
 mod discriminant;
@@ -119,9 +117,7 @@ impl<F: Fn() -> String> Drop for PrintOnPanic<F> {
     }
 }
 
-pub struct CraneliftCodegenBackend {
-    pub config: OnceCell<BackendConfig>,
-}
+pub struct CraneliftCodegenBackend;
 
 impl CodegenBackend for CraneliftCodegenBackend {
     fn name(&self) -> &'static str {
@@ -141,15 +137,6 @@ impl CodegenBackend for CraneliftCodegenBackend {
             sess.dcx()
                 .fatal("`-Cinstrument-coverage` is LLVM specific and not supported by Cranelift");
         }
-
-        let config = self.config.get_or_init(|| {
-            BackendConfig::from_opts(&sess.opts.cg.llvm_args)
-                .unwrap_or_else(|err| sess.dcx().fatal(err))
-        });
-
-        if config.jit_mode && !sess.opts.output_types.should_codegen() {
-            sess.dcx().fatal("JIT mode doesn't work with `cargo check`");
-        }
     }
 
     fn target_config(&self, sess: &Session) -> TargetConfig {
@@ -211,16 +198,17 @@ impl CodegenBackend for CraneliftCodegenBackend {
 
     fn codegen_crate(&self, tcx: TyCtxt<'_>, _crate_info: &CrateInfo) -> Box<dyn Any> {
         info!("codegen crate {}", tcx.crate_name(LOCAL_CRATE));
-        let config = self.config.get().unwrap();
-        if config.jit_mode {
-            #[cfg(feature = "jit")]
-            driver::jit::run_jit(tcx, _crate_info, config.jit_args.clone());
 
-            #[cfg(not(feature = "jit"))]
-            tcx.dcx().fatal("jit support was disabled when compiling rustc_codegen_cranelift");
-        } else {
-            driver::aot::run_aot(tcx)
+        for opt in &tcx.sess.opts.cg.llvm_args {
+            if opt.starts_with("-import-instr-limit") {
+                // Silently ignore -import-instr-limit. It is set by rust's build system even when
+                // testing cg_clif.
+                continue;
+            }
+            tcx.sess.dcx().fatal(format!("Unknown option `{}`", opt));
         }
+
+        driver::aot::run_aot(tcx)
     }
 
     fn join_codegen(
@@ -231,6 +219,29 @@ impl CodegenBackend for CraneliftCodegenBackend {
     ) -> (CompiledModules, FxIndexMap<WorkProductId, WorkProduct>) {
         ongoing_codegen.downcast::<driver::aot::OngoingCodegen>().unwrap().join(sess, outputs)
     }
+
+    fn jit_crate<'tcx>(&self, tcx: TyCtxt<'tcx>, args: Vec<String>) -> ExitCode {
+        info!("jit crate {}", tcx.crate_name(LOCAL_CRATE));
+
+        for opt in &tcx.sess.opts.cg.llvm_args {
+            if opt.starts_with("-import-instr-limit") {
+                // Silently ignore -import-instr-limit. It is set by rust's build system even when
+                // testing cg_clif.
+                continue;
+            }
+            tcx.sess.dcx().fatal(format!("Unknown option `{}`", opt));
+        }
+
+        #[cfg(feature = "jit")]
+        #[allow(unreachable_code)]
+        return driver::jit::run_jit(tcx, self.target_cpu(&tcx.sess), args);
+
+        #[cfg(not(feature = "jit"))]
+        {
+            let _ = args;
+            tcx.dcx().fatal("jit support was disabled when compiling rustc_codegen_cranelift");
+        }
+    }
 }
 
 /// Determine if the Cranelift ir verifier should run.
@@ -360,5 +371,5 @@ fn build_isa(sess: &Session, jit: bool) -> Arc<dyn TargetIsa + 'static> {
 /// This is the entrypoint for a hot plugged rustc_codegen_cranelift
 #[unsafe(no_mangle)]
 pub fn __rustc_codegen_backend() -> Box<dyn CodegenBackend> {
-    Box::new(CraneliftCodegenBackend { config: OnceCell::new() })
+    Box::new(CraneliftCodegenBackend)
 }

compiler/rustc_codegen_cranelift/src/lib.rs

@@ -1,5 +1,6 @@
 use std::any::Any;
 use std::hash::Hash;
+use std::process::ExitCode;
 
 use rustc_ast::expand::allocator::AllocatorMethod;
 use rustc_data_structures::fx::FxIndexMap;
@@ -138,6 +139,12 @@ pub trait CodegenBackend {
             self.name(),
         );
     }
+
+    /// Used in place of [`codegen_crate`](Self::codegen_crate) when `-Zjit-mode` is passed.
+    fn jit_crate<'tcx>(&self, tcx: TyCtxt<'tcx>, args: Vec<String>) -> ExitCode {
+        let _ = args;
+        tcx.sess.dcx().fatal("-Zjit-mode not supported by the active codegen backend")
+    }
 }
 
 pub trait ExtraBackendMethods:

compiler/rustc_codegen_ssa/src/traits/backend.rs

@@ -180,7 +180,15 @@ pub fn run_compiler(at_args: &[String], callbacks: &mut (dyn Callbacks + Send))
     // the compiler with @empty_file as argv[0] and no more arguments.
     let at_args = at_args.get(1..).unwrap_or_default();
 
-    let args = args::arg_expand_all(&default_early_dcx, at_args);
+    let mut args = args::arg_expand_all(&default_early_dcx, at_args);
+
+    let (args, jit_args) = if let Some(idx) = args.iter().position(|arg| arg == "--") {
+        let mut jit_args = args.split_off(idx);
+        jit_args.remove(0);
+        (args, jit_args)
+    } else {
+        (args, vec![])
+    };
 
     let (matches, help_only) = match handle_options(&default_early_dcx, &args) {
         HandledOptions::None => return,
@@ -201,6 +209,10 @@ pub fn run_compiler(at_args: &[String], callbacks: &mut (dyn Callbacks + Send))
     let has_input = input.is_some();
     let (odir, ofile) = make_output(&matches);
 
+    if !jit_args.is_empty() && !sopts.unstable_opts.jit_mode {
+        default_early_dcx.early_fatal("passing arguments after -- requires -Zjit-mode");
+    }
+
     drop(default_early_dcx);
 
     let mut config = interface::Config {
@@ -338,15 +350,15 @@ pub fn run_compiler(at_args: &[String], callbacks: &mut (dyn Callbacks + Send))
                 }
             }
 
-            Some(Linker::codegen_and_build_linker(tcx, &*compiler.codegen_backend))
+            Some(Linker::codegen_and_build_linker(tcx, &*compiler.codegen_backend, jit_args))
         });
 
         // Linking is done outside the `compiler.enter()` so that the
         // `GlobalCtxt` within `Queries` can be freed as early as possible.
         if let Some(linker) = linker {
             linker.link(sess, codegen_backend);
         }
-    })
+    });
 }
 
 fn dump_feature_usage_metrics(tcxt: TyCtxt<'_>, metrics_dir: &Path) {

compiler/rustc_driver_impl/src/lib.rs

@@ -1,5 +1,6 @@
 // tidy-alphabetical-start
 #![feature(decl_macro)]
+#![feature(exitcode_exit_method)]
 #![feature(file_buffered)]
 #![feature(iter_intersperse)]
 #![feature(try_blocks)]

compiler/rustc_interface/src/lib.rs

@@ -2,6 +2,7 @@ use std::any::Any;
 use std::ffi::{OsStr, OsString};
 use std::io::{self, BufWriter, Write};
 use std::path::{Path, PathBuf};
+use std::process::ExitCode;
 use std::sync::{Arc, LazyLock, OnceLock};
 use std::{env, fs, iter};
 
@@ -1189,14 +1190,8 @@ fn analysis(tcx: TyCtxt<'_>, (): ()) {
     }
 }
 
-/// Runs the codegen backend, after which the AST and analysis can
-/// be discarded.
-pub(crate) fn start_codegen<'tcx>(
-    codegen_backend: &dyn CodegenBackend,
-    tcx: TyCtxt<'tcx>,
-) -> (Box<dyn Any>, CrateInfo, EncodedMetadata) {
-    tcx.sess.timings.start_section(tcx.sess.dcx(), TimingSection::Codegen);
-
+/// A couple of checks that need to run before we run codegen.
+fn pre_codegen_checks(tcx: TyCtxt<'_>) {
     // Hook for tests.
     if let Some((def_id, _)) = tcx.entry_fn(())
         && find_attr!(tcx, def_id, RustcDelayedBugFromInsideQuery)
@@ -1216,6 +1211,17 @@ pub(crate) fn start_codegen<'tcx>(
     if let Some(guar) = tcx.sess.dcx().has_errors_or_delayed_bugs() {
         guar.raise_fatal();
     }
+}
+
+/// Runs the codegen backend, after which the AST and analysis can
+/// be discarded.
+pub(crate) fn start_codegen<'tcx>(
+    codegen_backend: &dyn CodegenBackend,
+    tcx: TyCtxt<'tcx>,
+) -> (Box<dyn Any>, CrateInfo, EncodedMetadata) {
+    tcx.sess.timings.start_section(tcx.sess.dcx(), TimingSection::Codegen);
+
+    pre_codegen_checks(tcx);
 
     info!("Pre-codegen\n{:?}", tcx.debug_stats());
 
@@ -1246,6 +1252,16 @@ pub(crate) fn start_codegen<'tcx>(
     (codegen, crate_info, metadata)
 }
 
+pub fn jit_crate<'tcx>(
+    codegen_backend: &dyn CodegenBackend,
+    tcx: TyCtxt<'tcx>,
+    args: Vec<String>,
+) -> ExitCode {
+    pre_codegen_checks(tcx);
+
+    tcx.sess.time("jit_crate", move || codegen_backend.jit_crate(tcx, args))
+}
+
 /// Compute and validate the crate name.
 pub fn get_crate_name(sess: &Session, krate_attrs: &[ast::Attribute]) -> Symbol {
     // We validate *all* occurrences of `#![crate_name]`, pick the first find and