src/driver/jit.rs

   1 //! The JIT driver uses [`cranelift_simplejit`] to JIT execute programs without writing any object
   2 //! files.
   3
   4 use std::cell::RefCell;
   5 use std::ffi::CString;
   6 use std::os::raw::{c_char, c_int};
   7
   8 use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
   9 use rustc_codegen_ssa::CrateInfo;
  10 use rustc_middle::mir::mono::MonoItem;
  11
  12 use cranelift_jit::{JITBuilder, JITModule};
  13
  14 use crate::{prelude::*, BackendConfig};
  15 use crate::{CodegenCx, CodegenMode};
  16
  17 thread_local! {
  18     pub static BACKEND_CONFIG: RefCell<Option<BackendConfig>> = RefCell::new(None);
  19     pub static CURRENT_MODULE: RefCell<Option<JITModule>> = RefCell::new(None);
  20 }
  21
  22 pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
  23     if !tcx.sess.opts.output_types.should_codegen() {
  24         tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
  25     }
  26
  27     if !tcx.sess.crate_types().contains(&rustc_session::config::CrateType::Executable) {
  28         tcx.sess.fatal("can't jit non-executable crate");
  29     }
  30
  31     let imported_symbols = load_imported_symbols_for_jit(tcx);
  32
  33     let isa = crate::build_isa(tcx.sess, &backend_config);
  34     let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
  35     jit_builder.hotswap(matches!(backend_config.codegen_mode, CodegenMode::JitLazy));
  36     crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
  37     jit_builder.symbols(imported_symbols);
  38     let mut jit_module = JITModule::new(jit_builder);
  39     assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
  40
  41     let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
  42     let mono_items = cgus
  43         .iter()
  44         .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
  45         .flatten()
  46         .collect::<FxHashMap<_, (_, _)>>()
  47         .into_iter()
  48         .collect::<Vec<(_, (_, _))>>();
  49
  50     let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);
  51
  52     super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
  53         super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
  54         for (mono_item, _) in mono_items {
  55             match mono_item {
  56                 MonoItem::Fn(inst) => match backend_config.codegen_mode {
  57                     CodegenMode::Aot => unreachable!(),
  58                     CodegenMode::Jit => {
  59                         cx.tcx.sess.time("codegen fn", || {
  60                             crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
  61                         });
  62                     }
  63                     CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
  64                 },
  65                 MonoItem::Static(def_id) => {
  66                     crate::constant::codegen_static(tcx, &mut jit_module, def_id);
  67                 }
  68                 MonoItem::GlobalAsm(item_id) => {
  69                     let item = tcx.hir().item(item_id);
  70                     tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
  71                 }
  72             }
  73         }
  74     });
  75
  76     jit_module.finalize_definitions();
  77
  78     if !cx.global_asm.is_empty() {
  79         tcx.sess.fatal("Inline asm is not supported in JIT mode");
  80     }
  81
  82     crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
  83     crate::main_shim::maybe_create_entry_wrapper(
  84         tcx,
  85         &mut jit_module,
  86         &mut cx.unwind_context,
  87         true,
  88     );
  89
  90     tcx.sess.abort_if_errors();
  91
  92     jit_module.finalize_definitions();
  93     unsafe { cx.unwind_context.register_jit(&jit_module) };
  94
  95     println!(
  96         "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
  97     );
  98
  99     let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
 100         .chain(backend_config.jit_args.iter().map(|arg| &**arg))
 101         .map(|arg| CString::new(arg).unwrap())
 102         .collect::<Vec<_>>();
 103     let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
 104
 105     // Push a null pointer as a terminating argument. This is required by POSIX and
 106     // useful as some dynamic linkers use it as a marker to jump over.
 107     argv.push(std::ptr::null());
 108
 109     BACKEND_CONFIG.with(|tls_backend_config| {
 110         assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
 111     });
 112
 113     let start_sig = Signature {
 114         params: vec![
 115             AbiParam::new(jit_module.target_config().pointer_type()),
 116             AbiParam::new(jit_module.target_config().pointer_type()),
 117         ],
 118         returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
 119         call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
 120     };
 121     let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
 122     let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);
 123
 124     CURRENT_MODULE
 125         .with(|current_module| assert!(current_module.borrow_mut().replace(jit_module).is_none()));
 126
 127     let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
 128         unsafe { ::std::mem::transmute(finalized_start) };
 129     let ret = f(args.len() as c_int, argv.as_ptr());
 130     std::process::exit(ret);
 131 }
 132
 133 #[no_mangle]
 134 extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
 135     rustc_middle::ty::tls::with(|tcx| {
 136         // lift is used to ensure the correct lifetime for instance.
 137         let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();
 138
 139         CURRENT_MODULE.with(|jit_module| {
 140             let mut jit_module = jit_module.borrow_mut();
 141             let jit_module = jit_module.as_mut().unwrap();
 142             let backend_config =
 143                 BACKEND_CONFIG.with(|backend_config| backend_config.borrow().clone().unwrap());
 144
 145             let name = tcx.symbol_name(instance).name.to_string();
 146             let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
 147             let func_id = jit_module.declare_function(&name, Linkage::Export, &sig).unwrap();
 148             jit_module.prepare_for_function_redefine(func_id).unwrap();
 149
 150             let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module.isa(), false);
 151             tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));
 152
 153             assert!(cx.global_asm.is_empty());
 154             jit_module.finalize_definitions();
 155             unsafe { cx.unwind_context.register_jit(&jit_module) };
 156             jit_module.get_finalized_function(func_id)
 157         })
 158     })
 159 }
 160
 161 fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
 162     use rustc_middle::middle::dependency_format::Linkage;
 163
 164     let mut dylib_paths = Vec::new();
 165
 166     let crate_info = CrateInfo::new(tcx);
 167     let formats = tcx.dependency_formats(LOCAL_CRATE);
 168     let data = &formats
 169         .iter()
 170         .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
 171         .unwrap()
 172         .1;
 173     for &(cnum, _) in &crate_info.used_crates_dynamic {
 174         let src = &crate_info.used_crate_source[&cnum];
 175         match data[cnum.as_usize() - 1] {
 176             Linkage::NotLinked | Linkage::IncludedFromDylib => {}
 177             Linkage::Static => {
 178                 let name = tcx.crate_name(cnum);
 179                 let mut err =
 180                     tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
 181                 err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
 182                 err.emit();
 183             }
 184             Linkage::Dynamic => {
 185                 dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
 186             }
 187         }
 188     }
 189
 190     let mut imported_symbols = Vec::new();
 191     for path in dylib_paths {
 192         use object::{Object, ObjectSymbol};
 193         let lib = libloading::Library::new(&path).unwrap();
 194         let obj = std::fs::read(path).unwrap();
 195         let obj = object::File::parse(&obj).unwrap();
 196         imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
 197             let name = symbol.name().unwrap().to_string();
 198             if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
 199                 return None;
 200             }
 201             if name.starts_with("rust_metadata_") {
 202                 // The metadata is part of a section that is not loaded by the dynamic linker in
 203                 // case of cg_llvm.
 204                 return None;
 205             }
 206             let dlsym_name = if cfg!(target_os = "macos") {
 207                 // On macOS `dlsym` expects the name without leading `_`.
 208                 assert!(name.starts_with('_'), "{:?}", name);
 209                 &name[1..]
 210             } else {
 211                 &name
 212             };
 213             let symbol: libloading::Symbol<'_, *const u8> =
 214                 unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
 215             Some((name, *symbol))
 216         }));
 217         std::mem::forget(lib)
 218     }
 219
 220     tcx.sess.abort_if_errors();
 221
 222     imported_symbols
 223 }
 224
 225 fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
 226     let tcx = cx.tcx;
 227
 228     let pointer_type = module.target_config().pointer_type();
 229
 230     let name = tcx.symbol_name(inst).name.to_string();
 231     let sig = crate::abi::get_function_sig(tcx, module.isa().triple(), inst);
 232     let func_id = module.declare_function(&name, Linkage::Export, &sig).unwrap();
 233
 234     let instance_ptr = Box::into_raw(Box::new(inst));
 235
 236     let jit_fn = module
 237         .declare_function(
 238             "__clif_jit_fn",
 239             Linkage::Import,
 240             &Signature {
 241                 call_conv: module.target_config().default_call_conv,
 242                 params: vec![AbiParam::new(pointer_type)],
 243                 returns: vec![AbiParam::new(pointer_type)],
 244             },
 245         )
 246         .unwrap();
 247
 248     cx.cached_context.clear();
 249     let trampoline = &mut cx.cached_context.func;
 250     trampoline.signature = sig.clone();
 251
 252     let mut builder_ctx = FunctionBuilderContext::new();
 253     let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);
 254
 255     let jit_fn = module.declare_func_in_func(jit_fn, trampoline_builder.func);
 256     let sig_ref = trampoline_builder.func.import_signature(sig);
 257
 258     let entry_block = trampoline_builder.create_block();
 259     trampoline_builder.append_block_params_for_function_params(entry_block);
 260     let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();
 261
 262     trampoline_builder.switch_to_block(entry_block);
 263     let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
 264     let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
 265     let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
 266     let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
 267     let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
 268     trampoline_builder.ins().return_(&ret_vals);
 269
 270     module
 271         .define_function(
 272             func_id,
 273             &mut cx.cached_context,
 274             &mut NullTrapSink {},
 275             &mut NullStackMapSink {},
 276         )
 277         .unwrap();
 278 }