Make DCE mandatory and run it much earlier (even per-crate, shrinking .rlibs).

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -397,7 +397,6 @@ impl CodegenArgs {
         // FIXME(eddyb) should these be handled as `-C linker-args="..."` instead?
         {
             // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
-            opts.optflag("", "no-dce", "disables running dead code elimination");
             opts.optflag(
                 "",
                 "no-compact-ids",
@@ -604,7 +603,6 @@ impl CodegenArgs {
         // FIXME(eddyb) should these be handled as `-C linker-args="..."` instead?
         let linker_opts = crate::linker::Options {
             // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
-            dce: !matches.opt_present("no-dce"),
             compact_ids: !matches.opt_present("no-compact-ids"),
             early_report_zombies: !matches.opt_present("no-early-report-zombies"),
             infer_storage_classes: !matches.opt_present("no-infer-storage-classes"),

crates/rustc_codegen_spirv/src/lib.rs

@@ -167,16 +167,16 @@ use rustc_session::Session;
 use rustc_session::config::{self, OutputFilenames, OutputType};
 use rustc_span::symbol::Symbol;
 use std::any::Any;
-use std::fs::{File, create_dir_all};
+use std::fs;
 use std::io::Cursor;
 use std::io::Write;
 use std::path::Path;
 use std::sync::Arc;
 use tracing::{error, warn};
 
 fn dump_mir(tcx: TyCtxt<'_>, mono_items: &[(MonoItem<'_>, MonoItemData)], path: &Path) {
-    create_dir_all(path.parent().unwrap()).unwrap();
-    let mut file = File::create(path).unwrap();
+    fs::create_dir_all(path.parent().unwrap()).unwrap();
+    let mut file = fs::File::create(path).unwrap();
     for &(mono_item, _) in mono_items {
         if let MonoItem::Fn(instance) = mono_item
             && matches!(instance.def, InstanceKind::Item(_))
@@ -189,27 +189,6 @@ fn dump_mir(tcx: TyCtxt<'_>, mono_items: &[(MonoItem<'_>, MonoItemData)], path:
     }
 }
 
-// TODO: Should this store Vec or Module?
-struct SpirvModuleBuffer(Vec<u32>);
-
-impl ModuleBufferMethods for SpirvModuleBuffer {
-    fn data(&self) -> &[u8] {
-        spirv_tools::binary::from_binary(&self.0)
-    }
-}
-
-// TODO: Should this store Vec or Module?
-struct SpirvThinBuffer(Vec<u32>);
-
-impl ThinBufferMethods for SpirvThinBuffer {
-    fn data(&self) -> &[u8] {
-        spirv_tools::binary::from_binary(&self.0)
-    }
-    fn thin_link_data(&self) -> &[u8] {
-        &[]
-    }
-}
-
 #[derive(Clone)]
 struct SpirvCodegenBackend;
 
@@ -306,28 +285,83 @@ impl CodegenBackend for SpirvCodegenBackend {
     }
 }
 
+struct SpirvModuleBuffer(Vec<u32>);
+
+impl SpirvModuleBuffer {
+    fn as_bytes(&self) -> &[u8] {
+        spirv_tools::binary::from_binary(&self.0)
+    }
+}
+impl ModuleBufferMethods for SpirvModuleBuffer {
+    fn data(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+impl ThinBufferMethods for SpirvModuleBuffer {
+    fn data(&self) -> &[u8] {
+        self.as_bytes()
+    }
+    fn thin_link_data(&self) -> &[u8] {
+        &[]
+    }
+}
+
+impl SpirvCodegenBackend {
+    fn optimize_common(
+        _cgcx: &CodegenContext<Self>,
+        module: &mut ModuleCodegen<<Self as WriteBackendMethods>::Module>,
+    ) -> Result<(), FatalError> {
+        // Apply DCE ("dead code elimination") to modules before ever serializing
+        // them as `.spv` files (technically, `.rcgu.o` files inside `.rlib`s),
+        // that will later get linked (potentially many times, esp. if this is
+        // some big upstream library, e.g. `core` itself), and will therefore
+        // benefit from not having to clean up all sorts of unreachable helpers.
+        linker::dce::dce(&mut module.module_llvm);
+
+        // FIXME(eddyb) run as many optimization passes as possible, not just DCE.
+
+        Ok(())
+    }
+}
+
 impl WriteBackendMethods for SpirvCodegenBackend {
-    type Module = Vec<u32>;
+    type Module = rspirv::dr::Module;
     type TargetMachine = ();
     type TargetMachineError = String;
     type ModuleBuffer = SpirvModuleBuffer;
     type ThinData = ();
-    type ThinBuffer = SpirvThinBuffer;
+    type ThinBuffer = SpirvModuleBuffer;
 
+    // FIXME(eddyb) reuse the "merge" stage of `crate::linker` for this, or even
+    // delegate to `run_fat_lto` (although `-Zcombine-cgu` is much more niche).
     fn run_link(
-        _cgcx: &CodegenContext<Self>,
-        _diag_handler: DiagCtxtHandle<'_>,
+        cgcx: &CodegenContext<Self>,
+        diag_handler: DiagCtxtHandle<'_>,
         _modules: Vec<ModuleCodegen<Self::Module>>,
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
-        todo!()
+        assert!(
+            cgcx.opts.unstable_opts.combine_cgu,
+            "`run_link` (for `WorkItemResult::NeedsLink`) should \
+             only be invoked due to `-Zcombine-cgu`"
+        );
+        diag_handler.fatal("Rust-GPU does not support `-Zcombine-cgu`")
     }
 
+    // FIXME(eddyb) reuse the "merge" stage of `crate::linker` for this, or even
+    // consider setting `requires_lto = true` in the target specs and moving the
+    // entirety of `crate::linker` into this stage (lacking diagnostics may be
+    // an issue - it's surprising `CodegenBackend::link` has `Session` at all).
     fn run_fat_lto(
-        _: &CodegenContext<Self>,
-        _: Vec<FatLtoInput<Self>>,
-        _: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
+        cgcx: &CodegenContext<Self>,
+        _modules: Vec<FatLtoInput<Self>>,
+        _cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
     ) -> Result<LtoModuleCodegen<Self>, FatalError> {
-        todo!()
+        assert!(
+            cgcx.lto == rustc_session::config::Lto::Fat,
+            "`run_fat_lto` (for `WorkItemResult::NeedsFatLto`) should \
+             only be invoked due to `-Clto` (or equivalent)"
+        );
+        unreachable!("Rust-GPU does not support fat LTO")
     }
 
     fn run_thin_lto(
@@ -347,35 +381,39 @@ impl WriteBackendMethods for SpirvCodegenBackend {
     }
 
     fn optimize(
-        _: &CodegenContext<Self>,
-        _: DiagCtxtHandle<'_>,
-        _: &mut ModuleCodegen<Self::Module>,
-        _: &ModuleConfig,
+        cgcx: &CodegenContext<Self>,
+        _dcx: DiagCtxtHandle<'_>,
+        module: &mut ModuleCodegen<Self::Module>,
+        _config: &ModuleConfig,
     ) -> Result<(), FatalError> {
-        // TODO: Implement
-        Ok(())
+        Self::optimize_common(cgcx, module)
     }
 
     fn optimize_thin(
-        _cgcx: &CodegenContext<Self>,
+        cgcx: &CodegenContext<Self>,
         thin_module: ThinModule<Self>,
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
-        let module = ModuleCodegen {
-            module_llvm: spirv_tools::binary::to_binary(thin_module.data())
-                .unwrap()
-                .to_vec(),
+        // FIXME(eddyb) the inefficiency of Module -> [u8] -> Module roundtrips
+        // comes from upstream and it applies to `rustc_codegen_llvm` as well,
+        // eventually it should be properly addressed (for `ThinLocal` at least).
+        let mut module = ModuleCodegen {
+            module_llvm: link::with_rspirv_loader(|loader| {
+                rspirv::binary::parse_bytes(thin_module.data(), loader)
+            })
+            .unwrap(),
             name: thin_module.name().to_string(),
             kind: ModuleKind::Regular,
             thin_lto_buffer: None,
         };
+        Self::optimize_common(cgcx, &mut module)?;
         Ok(module)
     }
 
     fn optimize_fat(
-        _: &CodegenContext<Self>,
-        _: &mut ModuleCodegen<Self::Module>,
+        cgcx: &CodegenContext<Self>,
+        module: &mut ModuleCodegen<Self::Module>,
     ) -> Result<(), FatalError> {
-        todo!()
+        Self::optimize_common(cgcx, module)
     }
 
     fn codegen(
@@ -384,20 +422,19 @@ impl WriteBackendMethods for SpirvCodegenBackend {
         module: ModuleCodegen<Self::Module>,
         _config: &ModuleConfig,
     ) -> Result<CompiledModule, FatalError> {
+        let kind = module.kind;
+        let (name, module_buffer) = Self::serialize_module(module);
+
         let path = cgcx.output_filenames.temp_path_for_cgu(
             OutputType::Object,
-            &module.name,
+            &name,
             cgcx.invocation_temp.as_deref(),
         );
-        // Note: endianness doesn't matter, readers deduce endianness from magic header.
-        let spirv_module = spirv_tools::binary::from_binary(&module.module_llvm);
-        File::create(&path)
-            .unwrap()
-            .write_all(spirv_module)
-            .unwrap();
+        fs::write(&path, module_buffer.as_bytes()).unwrap();
+
         Ok(CompiledModule {
-            name: module.name,
-            kind: module.kind,
+            name,
+            kind,
             object: Some(path),
             dwarf_object: None,
             bytecode: None,
@@ -411,11 +448,14 @@ impl WriteBackendMethods for SpirvCodegenBackend {
         module: ModuleCodegen<Self::Module>,
         _want_summary: bool,
     ) -> (String, Self::ThinBuffer) {
-        (module.name, SpirvThinBuffer(module.module_llvm))
+        Self::serialize_module(module)
     }
 
     fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer) {
-        (module.name, SpirvModuleBuffer(module.module_llvm))
+        (
+            module.name,
+            SpirvModuleBuffer(module.module_llvm.assemble()),
+        )
     }
 
     fn autodiff(
@@ -424,7 +464,7 @@ impl WriteBackendMethods for SpirvCodegenBackend {
         _diff_fncs: Vec<AutoDiffItem>,
         _config: &ModuleConfig,
     ) -> Result<(), FatalError> {
-        todo!()
+        unreachable!("Rust-GPU does not support autodiff")
     }
 }
 
@@ -490,12 +530,11 @@ impl ExtraBackendMethods for SpirvCodegenBackend {
         } else {
             with_no_trimmed_paths!(do_codegen(&mut cx));
         }
-        let spirv_module = cx.finalize_module().assemble();
 
         (
             ModuleCodegen {
                 name: cgu_name.to_string(),
-                module_llvm: spirv_module,
+                module_llvm: cx.finalize_module(),
                 kind: ModuleKind::Regular,
                 thin_lto_buffer: None,
             },

crates/rustc_codegen_spirv/src/link.rs

@@ -2,7 +2,7 @@
 use crate::maybe_pqp_cg_ssa as rustc_codegen_ssa;
 
 use crate::codegen_cx::{CodegenArgs, SpirvMetadata};
-use crate::{SpirvCodegenBackend, SpirvModuleBuffer, SpirvThinBuffer, linker};
+use crate::{SpirvCodegenBackend, SpirvModuleBuffer, linker};
 use ar::{Archive, GnuBuilder, Header};
 use rspirv::binary::Assemble;
 use rspirv::dr::Module;
@@ -548,6 +548,16 @@ fn create_archive(files: &[&Path], metadata: &[u8], out_filename: &Path) {
     builder.into_inner().unwrap();
 }
 
+// HACK(eddyb) hiding the actual implementation to avoid `rspirv::dr::Loader`
+// being hardcoded (as future work may need to customize it for various reasons).
+pub fn with_rspirv_loader<E>(
+    f: impl FnOnce(&mut dyn rspirv::binary::Consumer) -> Result<(), E>,
+) -> Result<rspirv::dr::Module, E> {
+    let mut loader = rspirv::dr::Loader::new();
+    f(&mut loader)?;
+    Ok(loader.module())
+}
+
 /// This is the actual guts of linking: the rest of the link-related functions are just digging through rustc's
 /// shenanigans to collect all the object files we need to link.
 fn do_link(
@@ -562,11 +572,8 @@ fn do_link(
 
     let mut modules = Vec::new();
     let mut add_module = |file_name: &OsStr, bytes: &[u8]| {
-        let module = {
-            let mut loader = rspirv::dr::Loader::new();
-            rspirv::binary::parse_bytes(bytes, &mut loader).unwrap();
-            loader.module()
-        };
+        let module =
+            with_rspirv_loader(|loader| rspirv::binary::parse_bytes(bytes, loader)).unwrap();
         if let Some(dir) = &cg_args.dump_pre_link {
             // FIXME(eddyb) is it a good idea to re-`assemble` the `rspirv::dr`
             // module, or should this just save the original bytes?
@@ -625,7 +632,7 @@ fn do_link(
 // TODO: WorkProduct impl
 pub(crate) fn run_thin(
     cgcx: &CodegenContext<SpirvCodegenBackend>,
-    modules: Vec<(String, SpirvThinBuffer)>,
+    modules: Vec<(String, SpirvModuleBuffer)>,
     cached_modules: Vec<(SerializedModule<SpirvModuleBuffer>, WorkProduct)>,
 ) -> Result<(Vec<LtoModuleCodegen<SpirvCodegenBackend>>, Vec<WorkProduct>), FatalError> {
     if cgcx.opts.cg.linker_plugin_lto.enabled() {