(crudely) implement MIR-only rlibs

saethlin · saethlin · commit 1a13a122cb20 · 2024-12-08T14:45:06.000-05:00
diff --git a/compiler/rustc_codegen_ssa/src/back/symbol_export.rs b/compiler/rustc_codegen_ssa/src/back/symbol_export.rs
@@ -209,8 +209,14 @@ fn exported_symbols_provider_local(
     if allocator_kind_for_codegen(tcx).is_some() {
         for symbol_name in ALLOCATOR_METHODS
             .iter()
-            .map(|method| format!("__rust_{}", method.name))
-            .chain(["__rust_alloc_error_handler".to_string(), OomStrategy::SYMBOL.to_string()])
+            .flat_map(|method| {
+                [format!("__rust_{}", method.name), format!("__rdl_{}", method.name)]
+            })
+            .chain([
+                "__rust_alloc_error_handler".to_string(),
+                OomStrategy::SYMBOL.to_string(),
+                "__rg_oom".to_string(),
+            ])
         {
             let exported_symbol = ExportedSymbol::NoDefId(SymbolName::new(tcx, &symbol_name));
 
@@ -359,6 +365,27 @@ fn exported_symbols_provider_local(
         }
     }
 
+    if tcx.building_mir_only_rlib() {
+        for def_id in tcx.mir_keys(()) {
+            if !matches!(tcx.def_kind(def_id.to_def_id()), DefKind::Static { .. }) {
+                continue;
+            }
+            if tcx.is_reachable_non_generic(def_id.to_def_id()) {
+                continue;
+            }
+            let codegen_attrs = tcx.codegen_fn_attrs(def_id.to_def_id());
+            symbols.push((ExportedSymbol::NonGeneric(def_id.to_def_id()), SymbolExportInfo {
+                level: symbol_export_level(tcx, def_id.to_def_id()),
+                kind: if codegen_attrs.flags.contains(CodegenFnAttrFlags::THREAD_LOCAL) {
+                    SymbolExportKind::Tls
+                } else {
+                    SymbolExportKind::Data
+                },
+                used: true,
+            }));
+        }
+    }
+
     // Sort so we get a stable incr. comp. hash.
     symbols.sort_by_cached_key(|s| s.0.symbol_name_for_local_instance(tcx));
 
diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs
@@ -806,6 +806,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(mir_emit_retag, true);
     tracked!(mir_enable_passes, vec![("DestProp".to_string(), false)]);
     tracked!(mir_keep_place_mention, true);
+    tracked!(mir_only_rlibs, true);
     tracked!(mir_opt_level, Some(4));
     tracked!(move_size_limit, Some(4096));
     tracked!(mutable_noalias, false);
diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -572,6 +572,14 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
                     .filter_map(|(cnum, data)| data.used().then_some(cnum)),
             )
         },
+        mir_only_crates: |tcx, ()| {
+            tcx.untracked().cstore.freeze();
+            let store = CStore::from_tcx(tcx);
+            let crates = store
+                .iter_crate_data()
+                .filter_map(|(cnum, data)| if data.root.is_mir_only { Some(cnum) } else { None });
+            tcx.arena.alloc_from_iter(crates)
+        },
         ..providers.queries
     };
     provide_extern(&mut providers.extern_queries);
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -736,6 +736,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 impls,
                 incoherent_impls,
                 exported_symbols,
+                is_mir_only: tcx.building_mir_only_rlib(),
                 interpret_alloc_index,
                 tables,
                 syntax_contexts,
@@ -1058,12 +1059,13 @@ fn should_encode_mir(
     reachable_set: &LocalDefIdSet,
     def_id: LocalDefId,
 ) -> (bool, bool) {
+    let opts = &tcx.sess.opts;
+    let mir_required = opts.unstable_opts.always_encode_mir || tcx.building_mir_only_rlib();
     match tcx.def_kind(def_id) {
         // Constructors
         DefKind::Ctor(_, _) => {
-            let mir_opt_base = tcx.sess.opts.output_types.should_codegen()
-                || tcx.sess.opts.unstable_opts.always_encode_mir;
-            (true, mir_opt_base)
+            let opt = mir_required || opts.output_types.should_codegen();
+            (true, opt)
         }
         // Constants
         DefKind::AnonConst | DefKind::InlineConst | DefKind::AssocConst | DefKind::Const => {
@@ -1075,7 +1077,7 @@ fn should_encode_mir(
         // Full-fledged functions + closures
         DefKind::AssocFn | DefKind::Fn | DefKind::Closure => {
             let generics = tcx.generics_of(def_id);
-            let opt = tcx.sess.opts.unstable_opts.always_encode_mir
+            let opt = mir_required
                 || (tcx.sess.opts.output_types.should_codegen()
                     && reachable_set.contains(&def_id)
                     && (generics.requires_monomorphization(tcx)
diff --git a/compiler/rustc_metadata/src/rmeta/mod.rs b/compiler/rustc_metadata/src/rmeta/mod.rs
@@ -275,6 +275,7 @@ pub(crate) struct CrateRoot {
     debugger_visualizers: LazyArray<DebuggerVisualizerFile>,
 
     exported_symbols: LazyArray<(ExportedSymbol<'static>, SymbolExportInfo)>,
+    is_mir_only: bool,
 
     syntax_contexts: SyntaxContextTable,
     expn_data: ExpnDataTable,
diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
@@ -91,6 +91,17 @@ impl<'tcx> MonoItem<'tcx> {
     }
 
     pub fn instantiation_mode(&self, tcx: TyCtxt<'tcx>) -> InstantiationMode {
+        // Always do LocalCopy codegen when building a MIR-only rlib
+        if tcx.building_mir_only_rlib() {
+            return InstantiationMode::LocalCopy;
+        }
+        // If this is a monomorphization from a MIR-only rlib and we are building another lib, do
+        // local codegen.
+        if tcx.mir_only_crates(()).iter().any(|c| *c == self.def_id().krate)
+            && tcx.crate_types() == &[rustc_session::config::CrateType::Rlib]
+        {
+            return InstantiationMode::LocalCopy;
+        }
         let generate_cgu_internal_copies = tcx
             .sess
             .opts
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -2371,6 +2371,11 @@ rustc_queries! {
         desc { "estimating codegen size of `{}`", key }
         cache_on_disk_if { true }
     }
+
+    query mir_only_crates(_: ()) -> &'tcx [CrateNum] {
+        eval_always
+        desc { "fetching all foreign crates built in mir-only mode" }
+    }
 }
 
 rustc_query_append! { define_callbacks! }
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -1799,6 +1799,10 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn dcx(self) -> DiagCtxtHandle<'tcx> {
         self.sess.dcx()
     }
+
+    pub fn building_mir_only_rlib(self) -> bool {
+        self.sess.opts.unstable_opts.mir_only_rlibs && self.crate_types() == &[CrateType::Rlib]
+    }
 }
 
 impl<'tcx> TyCtxtAt<'tcx> {
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
@@ -215,6 +215,7 @@ use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::middle::exported_symbols::ExportedSymbol;
 use rustc_middle::mir::interpret::{AllocId, ErrorHandled, GlobalAlloc, Scalar};
 use rustc_middle::mir::mono::{CollectionMode, InstantiationMode, MonoItem};
 use rustc_middle::mir::visit::Visitor as MirVisitor;
@@ -230,7 +231,7 @@ use rustc_middle::ty::{
 use rustc_middle::util::Providers;
 use rustc_middle::{bug, span_bug};
 use rustc_session::Limit;
-use rustc_session::config::EntryFnType;
+use rustc_session::config::{CrateType, EntryFnType};
 use rustc_span::source_map::{Spanned, dummy_spanned, respan};
 use rustc_span::symbol::sym;
 use rustc_span::{DUMMY_SP, Span};
@@ -955,23 +956,40 @@ fn should_codegen_locally<'tcx>(tcx: TyCtxtAt<'tcx>, instance: Instance<'tcx>) -
         return true;
     };
 
+    let def_is_for_mir_only_rlib = if def_id.krate == rustc_hir::def_id::LOCAL_CRATE {
+        tcx.building_mir_only_rlib()
+    } else {
+        tcx.mir_only_crates(()).iter().any(|c| *c == def_id.krate)
+    };
+
     if tcx.is_foreign_item(def_id) {
-        // Foreign items are always linked against, there's no way of instantiating them.
-        return false;
+        if def_is_for_mir_only_rlib {
+            return tcx.is_mir_available(instance.def_id());
+        } else {
+            // Foreign items are always linked against, there's no way of instantiating them.
+            return false;
+        }
+    }
+
+    if def_is_for_mir_only_rlib {
+        let has_mir = tcx.is_mir_available(instance.def_id());
+        return has_mir || matches!(tcx.def_kind(instance.def_id()), DefKind::Static { .. });
     }
 
     if def_id.is_local() {
         // Local items cannot be referred to locally without monomorphizing them locally.
         return true;
     }
 
-    if tcx.is_reachable_non_generic(def_id) || instance.upstream_monomorphization(*tcx).is_some() {
-        // We can link to the item in question, no instance needed in this crate.
-        return false;
+    if !def_is_for_mir_only_rlib {
+        if let DefKind::Static { .. } = tcx.def_kind(def_id) {
+            // We cannot monomorphize statics from upstream crates.
+            return false;
+        }
     }
 
-    if let DefKind::Static { .. } = tcx.def_kind(def_id) {
-        // We cannot monomorphize statics from upstream crates.
+    if tcx.is_reachable_non_generic(def_id) || instance.upstream_monomorphization(*tcx).is_some() {
+        // We can link to the item in question, no instance needed in this crate.
         return false;
     }
 
@@ -1370,6 +1388,7 @@ fn collect_roots(tcx: TyCtxt<'_>, mode: MonoItemCollectionStrategy) -> Vec<MonoI
         }
 
         collector.push_extra_entry_roots();
+        collector.push_extra_roots_from_mir_only_rlibs();
     }
 
     // We can only codegen items that are instantiable - items all of
@@ -1514,6 +1533,50 @@ impl<'v> RootCollector<'_, 'v> {
 
         self.output.push(create_fn_mono_item(self.tcx, start_instance, DUMMY_SP));
     }
+
+    fn push_extra_roots_from_mir_only_rlibs(&mut self) {
+        // An upstream extern function may be used anywhere in the dependency tree, so we
+        // cannot do any reachability analysis on them. We blindly monomorphize every
+        // extern function declared anywhere in our dependency tree. We must give them
+        // GloballyShared codegen because we don't know if the only call to an upstream
+        // extern function is also upstream: We don't have reachability information. All we
+        // can do is codegen all extern functions and pray for the linker to delete the
+        // ones that are reachable.
+        if !self.tcx.crate_types().iter().any(|c| !matches!(c, CrateType::Rlib)) {
+            return;
+        }
+
+        for (symbol, _info) in self
+            .tcx
+            .mir_only_crates(())
+            .into_iter()
+            .flat_map(|krate| self.tcx.exported_symbols(*krate))
+        {
+            let def_id = match symbol {
+                ExportedSymbol::NonGeneric(def_id) => def_id,
+                ExportedSymbol::ThreadLocalShim(def_id) => {
+                    let item = MonoItem::Fn(Instance {
+                        def: InstanceKind::ThreadLocalShim(*def_id),
+                        args: GenericArgs::empty(),
+                    });
+                    self.output.push(dummy_spanned(item));
+                    continue;
+                }
+                _ => continue,
+            };
+            match self.tcx.def_kind(def_id) {
+                DefKind::Fn | DefKind::AssocFn => {
+                    let instance = Instance::mono(self.tcx, *def_id);
+                    let item = create_fn_mono_item(self.tcx, instance, DUMMY_SP);
+                    self.output.push(item);
+                }
+                DefKind::Static { .. } => {
+                    self.output.push(dummy_spanned(MonoItem::Static(*def_id)));
+                }
+                _ => {}
+            }
+        }
+    }
 }
 
 #[instrument(level = "debug", skip(tcx, output))]
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -145,6 +145,12 @@ fn partition<'tcx, I>(
 where
     I: Iterator<Item = MonoItem<'tcx>>,
 {
+    if tcx.building_mir_only_rlib() {
+        let cgu_name_builder = &mut CodegenUnitNameBuilder::new(tcx);
+        let cgu_name = fallback_cgu_name(cgu_name_builder);
+        return vec![CodegenUnit::new(cgu_name)];
+    }
+
     let _prof_timer = tcx.prof.generic_activity("cgu_partitioning");
 
     let cx = &PartitioningCx { tcx, usage_map };
@@ -169,6 +175,10 @@ where
         debug_dump(tcx, "MERGE", &codegen_units);
     }
 
+    if !codegen_units.is_sorted_by(|a, b| a.name().as_str() < b.name().as_str()) {
+        bug!("unsorted CGUs");
+    }
+
     // Make as many symbols "internal" as possible, so LLVM has more freedom to
     // optimize.
     if !tcx.sess.link_dead_code() {
@@ -189,7 +199,12 @@ where
         for cgu in codegen_units.iter() {
             names += &format!("- {}\n", cgu.name());
         }
-        bug!("unsorted CGUs:\n{names}");
+        codegen_units.sort_by(|a, b| a.name().as_str().cmp(b.name().as_str()));
+        let mut sorted_names = String::new();
+        for cgu in codegen_units.iter() {
+            sorted_names += &format!("- {}\n", cgu.name());
+        }
+        bug!("unsorted CGUs:\n{names}\n{sorted_names}");
     }
 
     codegen_units
@@ -213,6 +228,9 @@ where
     let cgu_name_builder = &mut CodegenUnitNameBuilder::new(cx.tcx);
     let cgu_name_cache = &mut UnordMap::default();
 
+    let start_fn = cx.tcx.lang_items().start_fn();
+    let entry_fn = cx.tcx.entry_fn(()).map(|(id, _)| id);
+
     for mono_item in mono_items {
         // Handle only root (GloballyShared) items directly here. Inlined (LocalCopy) items
         // are handled at the bottom of the loop based on reachability, with one exception.
@@ -221,7 +239,8 @@ where
         match mono_item.instantiation_mode(cx.tcx) {
             InstantiationMode::GloballyShared { .. } => {}
             InstantiationMode::LocalCopy => {
-                if Some(mono_item.def_id()) != cx.tcx.lang_items().start_fn() {
+                let def_id = mono_item.def_id();
+                if ![start_fn, entry_fn].contains(&Some(def_id)) {
                     continue;
                 }
             }
@@ -243,7 +262,7 @@ where
 
         let cgu = codegen_units.entry(cgu_name).or_insert_with(|| CodegenUnit::new(cgu_name));
 
-        let mut can_be_internalized = true;
+        let mut can_be_internalized = false;
         let (linkage, visibility) = mono_item_linkage_and_visibility(
             cx.tcx,
             &mono_item,
@@ -486,7 +505,7 @@ fn merge_codegen_units<'tcx>(
         // If we didn't zero-pad the sorted-by-name order would be `XYZ-cgu.0`,
         // `XYZ-cgu.1`, `XYZ-cgu.10`, `XYZ-cgu.11`, ..., `XYZ-cgu.2`, etc.
         codegen_units.sort_by_key(|cgu| cmp::Reverse(cgu.size_estimate()));
-        let num_digits = codegen_units.len().ilog10() as usize + 1;
+        let num_digits = std::hint::black_box(codegen_units.len().ilog10() as usize + 1);
         for (index, cgu) in codegen_units.iter_mut().enumerate() {
             // Note: `WorkItem::short_description` depends on this name ending
             // with `-cgu.` followed by a numeric suffix. Please keep it in
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
@@ -1890,6 +1890,8 @@ options! {
     mir_keep_place_mention: bool = (false, parse_bool, [TRACKED],
         "keep place mention MIR statements, interpreted e.g., by miri; implies -Zmir-opt-level=0 \
         (default: no)"),
+    mir_only_rlibs: bool = (false, parse_bool, [TRACKED],
+        "only generate MIR when building rlibs (default: no)"),
     #[rustc_lint_opt_deny_field_access("use `Session::mir_opt_level` instead of this field")]
     mir_opt_level: Option<usize> = (None, parse_opt_number, [TRACKED],
         "MIR optimization level (0-4; default: 1 in non optimized builds and 2 in optimized builds)"),

Original file line number	Diff line number	Diff line change
`@@ -2371,6 +2371,11 @@ rustc_queries! {`
`2371`	`2371`	desc { "estimating codegen size of `{}`", key }
`2372`	`2372`	`cache_on_disk_if { true }`
`2373`	`2373`	`}`
	`2374`	`+`
	`2375`	`+ query mir_only_crates(_: ()) -> &'tcx [CrateNum] {`
	`2376`	`+ eval_always`
	`2377`	`+ desc { "fetching all foreign crates built in mir-only mode" }`
	`2378`	`+ }`
`2374`	`2379`	`}`
`2375`	`2380`
`2376`	`2381`	`rustc_query_append! { define_callbacks! }`
Original file line number	Diff line number	Diff line change
`@@ -1799,6 +1799,10 @@ impl<'tcx> TyCtxt<'tcx> {`
`1799`	`1799`	`pub fn dcx(self) -> DiagCtxtHandle<'tcx> {`
`1800`	`1800`	`self.sess.dcx()`
`1801`	`1801`	`}`
	`1802`	`+`
	`1803`	`+ pub fn building_mir_only_rlib(self) -> bool {`
	`1804`	`+ self.sess.opts.unstable_opts.mir_only_rlibs && self.crate_types() == &[CrateType::Rlib]`
	`1805`	`+ }`
`1802`	`1806`	`}`
`1803`	`1807`
`1804`	`1808`	`impl<'tcx> TyCtxtAt<'tcx> {`