JIT mode integration for LTO

Author: bjorn3

build_system/tests.rs

@@ -437,6 +437,7 @@ impl<'a> TestRunner<'a> {
                     }
                     spawn_and_wait(jit_cmd);
 
+                    /*
                     eprintln!("[JIT-lazy] {testname}");
                     let mut jit_cmd = self.rustc_command([
                         "-Zunstable-options",
@@ -450,6 +451,7 @@ impl<'a> TestRunner<'a> {
                         jit_cmd.env("CG_CLIF_JIT_ARGS", args);
                     }
                     spawn_and_wait(jit_cmd);
+                    */
                 }
             }
         }

example/mini_core.rs

@@ -735,6 +735,7 @@ struct PanicLocation {
     column: u32,
 }
 
+/*
 #[no_mangle]
 #[cfg(not(all(windows, target_env = "gnu")))]
 pub fn get_tls() -> u8 {
@@ -743,3 +744,4 @@ pub fn get_tls() -> u8 {
 
     A
 }
+*/

example/mini_core_hello_world.rs

@@ -333,6 +333,7 @@ fn main() {
     #[cfg(all(not(jit), not(all(windows, target_env = "gnu"))))]
     test_tls();
 
+    /*
     #[cfg(all(
         not(jit),
         not(no_unstable_features),
@@ -343,6 +344,7 @@ fn main() {
         global_asm_test();
         naked_test();
     }
+    */
 
     // Both statics have a reference that points to the same anonymous allocation.
     static REF1: &u8 = &42;
@@ -367,6 +369,7 @@ fn stack_val_align() {
     assert_eq!(&a as *const Foo as usize % 8192, 0);
 }
 
+/*
 #[cfg(all(
     not(jit),
     not(no_unstable_features),
@@ -396,6 +399,7 @@ global_asm! {
     ret
     "
 }
+*/
 
 #[cfg(all(not(jit), not(no_unstable_features), target_arch = "x86_64"))]
 #[naked]

src/driver/jit.rs

@@ -160,12 +160,18 @@ fn create_jit_module(
     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_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info));
     jit_builder.symbol("__clif_jit_fn", clif_jit_fn as *const u8);
     let jit_module = cranelift_jit::JITModule::new(jit_builder);
     let unwind_context = UnwindContext::new(jit_module.isa(), false);
     let mut jit_module = JITModule { jit_module, unwind_context };
 
+    for (_name, module) in
+        super::lto::load_lto_modules(tcx, &CrateInfo::new(tcx, "dummy".to_owned()), &backend_config)
+    {
+        module.apply_to(&mut jit_module);
+    }
+
     let cx = crate::CodegenCx::new(tcx, jit_module.isa(), false, Symbol::intern("dummy_cgu_name"));
 
     crate::allocator::codegen(tcx, &mut jit_module);

src/driver/lto.rs

@@ -15,7 +15,7 @@ use object::{Object, ObjectSection};
 use crate::debuginfo::TypeDebugContext;
 use crate::prelude::*;
 
-struct SerializeModule {
+pub(super) struct SerializeModule {
     isa: Arc<dyn isa::TargetIsa>,
     inner: SerializeModuleInner,
 }
@@ -33,11 +33,11 @@ impl SerializeModule {
         postcard::to_stdvec(&self.inner).unwrap()
     }
 
-    fn deserialize(blob: &[u8], isa: Arc<dyn isa::TargetIsa>) -> SerializeModule {
+    pub(super) fn deserialize(blob: &[u8], isa: Arc<dyn isa::TargetIsa>) -> SerializeModule {
         SerializeModule { isa, inner: postcard::from_bytes(blob).unwrap() }
     }
 
-    fn apply_to(self, module: &mut dyn Module) {
+    pub(super) fn apply_to(self, module: &mut dyn Module) {
         //println!("{:#?}", self.inner);
 
         let mut function_map: SecondaryMap<FuncId, Option<FuncId>> = SecondaryMap::new();
@@ -492,34 +492,12 @@ fn module_codegen(
     }
 }
 
-pub(crate) fn run_aot(
+pub(super) fn load_lto_modules(
     tcx: TyCtxt<'_>,
-    backend_config: BackendConfig,
-    metadata: EncodedMetadata,
-    need_metadata_module: bool,
-) -> Box<OngoingCodegen> {
-    // FIXME handle `-Ctarget-cpu=native`
-    let target_cpu = match tcx.sess.opts.cg.target_cpu {
-        Some(ref name) => name,
-        None => tcx.sess.target.cpu.as_ref(),
-    }
-    .to_owned();
-
-    let crate_info = CrateInfo::new(tcx, target_cpu);
-
-    let cgus = if tcx.sess.opts.output_types.should_codegen() {
-        tcx.collect_and_partition_mono_items(()).1
-    } else {
-        // If only `--emit metadata` is used, we shouldn't perform any codegen.
-        // Also `tcx.collect_and_partition_mono_items` may panic in that case.
-        &[]
-    };
-
-    let mut modules = tcx.sess.time("codegen mono items", || {
-        cgus.iter()
-            .map(|cgu| module_codegen(tcx, (backend_config.clone(), cgu.name())))
-            .collect::<Vec<_>>()
-    });
+    crate_info: &CrateInfo,
+    backend_config: &BackendConfig,
+) -> Vec<(String, SerializeModule)> {
+    let mut modules = vec![];
 
     if !(tcx.crate_types().len() == 1
         && tcx.crate_types()[0] == rustc_session::config::CrateType::Rlib)
@@ -561,14 +539,47 @@ pub(crate) fn run_aot(
                     lto_object.section_by_name(".rodata.cgclif_lto").unwrap().data().unwrap(),
                     crate::build_isa(tcx.sess, &backend_config),
                 );
-                modules.push(
-                    emit_module(tcx, &backend_config, module, ModuleKind::Regular, name.to_owned())
-                        .unwrap(),
-                );
+                modules.push((name.to_owned(), module));
             }
         }
     }
 
+    modules
+}
+
+pub(crate) fn run_aot(
+    tcx: TyCtxt<'_>,
+    backend_config: BackendConfig,
+    metadata: EncodedMetadata,
+    need_metadata_module: bool,
+) -> Box<OngoingCodegen> {
+    // FIXME handle `-Ctarget-cpu=native`
+    let target_cpu = match tcx.sess.opts.cg.target_cpu {
+        Some(ref name) => name,
+        None => tcx.sess.target.cpu.as_ref(),
+    }
+    .to_owned();
+
+    let crate_info = CrateInfo::new(tcx, target_cpu);
+
+    let cgus = if tcx.sess.opts.output_types.should_codegen() {
+        tcx.collect_and_partition_mono_items(()).1
+    } else {
+        // If only `--emit metadata` is used, we shouldn't perform any codegen.
+        // Also `tcx.collect_and_partition_mono_items` may panic in that case.
+        &[]
+    };
+
+    let mut modules = tcx.sess.time("codegen mono items", || {
+        cgus.iter()
+            .map(|cgu| module_codegen(tcx, (backend_config.clone(), cgu.name())))
+            .collect::<Vec<_>>()
+    });
+
+    for (name, module) in load_lto_modules(tcx, &crate_info, &backend_config) {
+        modules.push(emit_module(tcx, &backend_config, module, ModuleKind::Regular, name).unwrap());
+    }
+
     let mut allocator_module = make_module(tcx.sess, &backend_config);
     let created_alloc_shim = crate::allocator::codegen(tcx, &mut allocator_module);