Feat: Implement Dead Code Elimination (no more gigantic PTX files 🎉 ) (#12)

* Feat: start implementing DCE/module merging

* Feat: remove useless deps graph stuff and remove mentions in docs

* Feat: rename path_tracer kernels name and commit ptx files

* Feat: add #[externally_visible], ignore ptx files, disable trace!

Author: RDambrosio016

crates/cuda_builder/src/lib.rs

@@ -378,6 +378,10 @@ fn invoke_rustc(builder: &CudaBuilder) -> Result<PathBuf, CudaBuilderError> {
         cargo.arg("--release");
     }
 
+    // TODO(RDambrosio016): Remove this once we can get meaningful error messages in panic to work.
+    // for now we enable it to remove some useless indirect calls in the ptx.
+    cargo.arg("-Zbuild-std-features=panic_immediate_abort");
+
     if builder.optix {
         cargo.arg("-Zbuild-std-features=panic_immediate_abort");
         cargo.arg("-Zunstable-options");

crates/cuda_std/CHANGELOG.md

@@ -0,0 +1,8 @@
+# Changelog
+
+Notable changes to this project will be documented in this file.
+
+## Unreleased
+
+- Added `#[externally_visible]` in conjunction with cg_nvvm dead code elimination changes to mark that
+a function is externally visible.

crates/cuda_std_macros/src/lib.rs

@@ -184,3 +184,27 @@ pub fn gpu_only(_attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -
 
     output.into()
 }
+
+/// Notifies the codegen that this function is externally visible and should not be
+/// removed if it is not used by a kernel. Usually used for linking with other PTX/cubin files.
+///
+/// # Panics
+///
+/// Panics if the function is not also no_mangle.
+#[proc_macro_attribute]
+pub fn externally_visible(
+    _attr: proc_macro::TokenStream,
+    item: proc_macro::TokenStream,
+) -> TokenStream {
+    let mut func = syn::parse_macro_input!(item as syn::ItemFn);
+
+    assert!(
+        func.attrs.iter().any(|a| a.path.is_ident("no_mangle")),
+        "#[externally_visible] function should also be #[no_mangle]"
+    );
+
+    let new_attr = parse_quote!(#[cfg_attr(target_os = "cuda", nvvm_internal(used))]);
+    func.attrs.push(new_attr);
+
+    func.into_token_stream().into()
+}

crates/rustc_codegen_nvvm/CHANGELOG.md

@@ -4,6 +4,17 @@ Notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Dead Code Elimination 
+
+PTX files no longer include useless functions and globals, we have switched to an alternative
+method of codegen for the final steps of the codegen. We no longer lazily-load modules using dependency graphs, 
+we instead merge all the modules into one then run global DCE on it before giving it to libnvvm.
+
+This means all of the dead code is gone before it gets to the libnvvm stage, drastically lowering the size of 
+the built PTX and improving codegen performance.
+
+- Trace-level debug is compiled out for release now, decreasing the size of the codegen dll and improving compile times.
+
 ## 0.1.1 - 11/26/21
 
 - Fix things using the `bswap` intrinsic panicking.

crates/rustc_codegen_nvvm/Cargo.toml

@@ -16,7 +16,7 @@ libc = "0.2.97"
 tar = "0.4.35"
 once_cell = "1.8.0"
 bitflags = "1.3.2"
-tracing = "0.1.29"
+tracing = { version = "0.1.29", features = ["release_max_level_debug"] }
 find_cuda_helper = { version = "0.1", path = "../find_cuda_helper" }
 tracing-subscriber = { version = "0.3.1", features = ["env-filter"] }
 rustc_codegen_nvvm_macros = { version = "0.1", path = "../rustc_codegen_nvvm_macros" }

crates/rustc_codegen_nvvm/src/attributes.rs

@@ -3,7 +3,7 @@ use rustc_ast::Attribute;
 use rustc_attr::{InlineAttr, OptimizeAttr};
 use rustc_middle::{middle::codegen_fn_attrs::CodegenFnAttrFlags, ty};
 use rustc_session::{config::OptLevel, Session};
-use rustc_span::Symbol;
+use rustc_span::{sym, Symbol};
 
 use crate::context::CodegenCx;
 
@@ -99,6 +99,7 @@ pub struct Symbols {
 #[derive(Default, Clone, PartialEq)]
 pub(crate) struct NvvmAttributes {
     pub kernel: bool,
+    pub used: bool,
 }
 
 impl NvvmAttributes {
@@ -112,6 +113,9 @@ impl NvvmAttributes {
                     if arg.has_name(cx.symbols.kernel) {
                         nvvm_attrs.kernel = true;
                     }
+                    if arg.has_name(sym::used) {
+                        nvvm_attrs.used = true;
+                    }
                 }
             }
         }

crates/rustc_codegen_nvvm/src/back.rs

@@ -14,11 +14,10 @@ use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_errors::{FatalError, Handler};
 use rustc_fs_util::path_to_c_string;
 use rustc_middle::bug;
-use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::{dep_graph, ty::TyCtxt};
 use rustc_session::config::{self, DebugInfo, OutputType};
 use rustc_session::Session;
-use rustc_span::{sym, Symbol};
+use rustc_span::Symbol;
 use rustc_target::spec::{CodeModel, RelocModel};
 use std::ffi::CString;
 use std::sync::Arc;
@@ -265,22 +264,6 @@ pub fn compile_codegen_unit(tcx: TyCtxt<'_>, cgu_name: Symbol) -> (ModuleCodegen
             // ... and now that we have everything pre-defined, fill out those definitions.
             for &(mono_item, _) in &mono_items {
                 mono_item.define::<Builder<'_, '_, '_>>(&cx);
-                if let MonoItem::Fn(inst) = mono_item {
-                    let name = tcx.symbol_name(inst).name;
-                    let attrs = tcx.get_attrs(inst.def_id());
-                    let is_no_mangle = attrs.iter().any(|x| x.has_name(sym::no_mangle));
-
-                    if name == "rust_begin_unwind"
-                        || name.starts_with("__rg")
-                        || name == "rust_oom"
-                        || is_no_mangle
-                    {
-                        let func = cx.get_fn(inst);
-                        let llval =
-                            unsafe { llvm::LLVMConstBitCast(func, cx.type_ptr_to(cx.type_i8())) };
-                        cx.used_statics.borrow_mut().push(llval);
-                    }
-                }
             }
 
             // a main function for gpu kernels really makes no sense but

crates/rustc_codegen_nvvm/src/lib.rs

@@ -69,7 +69,7 @@ use rustc_middle::{
 use rustc_session::{cstore::MetadataLoaderDyn, Session};
 use tracing::debug;
 
-use std::{ffi::CString, sync::Arc, sync::Mutex};
+use std::ffi::CString;
 
 // codegen dylib entrypoint
 #[no_mangle]
@@ -79,13 +79,7 @@ pub fn __rustc_codegen_backend() -> Box<dyn CodegenBackend> {
 }
 
 #[derive(Default, Clone)]
-pub struct NvvmCodegenBackend {
-    /// HACK(RDambrosio016): To lazy load modules, we need to know the depedency graph
-    /// of the crate, we can get this through cstore, but cstore is made using TyCtxt
-    /// and rustc drops TyCtxt before linking to save memory. So we populate this field
-    /// in codegen_crate so we can then use it in link.
-    deps: Arc<Mutex<Option<Vec<String>>>>,
-}
+pub struct NvvmCodegenBackend(());
 
 unsafe impl Send for NvvmCodegenBackend {}
 unsafe impl Sync for NvvmCodegenBackend {}
@@ -124,14 +118,6 @@ impl CodegenBackend for NvvmCodegenBackend {
         need_metadata_module: bool,
     ) -> Box<dyn std::any::Any> {
         debug!("Codegen crate");
-        let mut raw_deps = tcx.postorder_cnums(()).to_vec();
-        raw_deps.reverse();
-        let out = raw_deps
-            .into_iter()
-            .map(|x| tcx.crate_name(x).as_str().to_string())
-            .collect::<Vec<_>>();
-
-        *self.deps.lock().unwrap() = Some(out);
         Box::new(rustc_codegen_ssa::base::codegen_crate(
             NvvmCodegenBackend::default(),
             tcx,
@@ -164,7 +150,6 @@ impl CodegenBackend for NvvmCodegenBackend {
         outputs: &rustc_session::config::OutputFilenames,
     ) -> Result<(), rustc_errors::ErrorReported> {
         link::link(
-            self.deps.lock().unwrap().clone(),
             sess,
             &codegen_results,
             outputs,

crates/rustc_codegen_nvvm/src/link.rs

@@ -1,4 +1,3 @@
-use rustc_codegen_ssa::traits::ThinBufferMethods;
 use rustc_codegen_ssa::CodegenResults;
 use rustc_codegen_ssa::CompiledModule;
 use rustc_codegen_ssa::NativeLib;
@@ -16,7 +15,6 @@ use rustc_session::{
     Session,
 };
 use rustc_target::spec::Target;
-use std::ffi::CString;
 use std::{
     ffi::OsStr,
     fs::File,
@@ -27,11 +25,6 @@ use tar::{Archive, Builder, Header};
 use tracing::{debug, trace};
 
 use crate::context::CodegenArgs;
-use crate::create_module;
-use crate::llvm::Context;
-use crate::llvm::LLVMLinkModules2;
-use crate::llvm::LLVMRustParseBitcodeForLTO;
-use crate::lto::ThinBuffer;
 use crate::LlvmMod;
 
 pub(crate) struct NvvmMetadataLoader;
@@ -73,13 +66,12 @@ fn read_metadata(rlib: &Path) -> Result<MetadataRef, String> {
 }
 
 pub fn link<'tcx>(
-    deps: Option<Vec<String>>,
     sess: &'tcx Session,
     codegen_results: &CodegenResults,
     outputs: &OutputFilenames,
     crate_name: &str,
 ) {
-    debug!("Linking crate `{}`, deps:\n{:?}", crate_name, deps);
+    debug!("Linking crate `{}`", crate_name);
     // largely inspired by rust-gpu
     let output_metadata = sess.opts.output_types.contains_key(&OutputType::Metadata);
     for &crate_type in sess.crate_types().iter() {
@@ -107,7 +99,6 @@ pub fn link<'tcx>(
                 CrateType::Executable | CrateType::Cdylib | CrateType::Dylib => {
                     let _ = link_exe(
                         &codegen_results.allocator_module,
-                        deps.clone(),
                         sess,
                         crate_type,
                         &out_filename,
@@ -171,7 +162,6 @@ fn link_rlib(sess: &Session, codegen_results: &CodegenResults, out_filename: &Pa
 
 fn link_exe(
     allocator: &Option<CompiledModule>,
-    deps: Option<Vec<String>>,
     sess: &Session,
     crate_type: CrateType,
     out_filename: &Path,
@@ -201,22 +191,20 @@ fn link_exe(
         std::fs::create_dir_all(&out_dir)?;
     }
 
-    codegen_into_ptx_file(allocator, deps, sess, &objects, &rlibs, out_filename)
+    codegen_into_ptx_file(allocator, sess, &objects, &rlibs, out_filename)
 }
 
 /// This is the meat of the codegen, taking all of the llvm bitcode modules we have, and giving them to
 /// nvvm to make into a final
 fn codegen_into_ptx_file(
     allocator: &Option<CompiledModule>,
-    deps: Option<Vec<String>>,
     sess: &Session,
     objects: &[PathBuf],
     rlibs: &[PathBuf],
     out_filename: &Path,
 ) -> io::Result<()> {
-    debug!("Codegenning crate into PTX, allocator: {}, deps:\n{:#?}, objects:\n{:#?}, rlibs:\n{:#?}, out_filename:\n{:#?}",
+    debug!("Codegenning crate into PTX, allocator: {}, objects:\n{:#?}, rlibs:\n{:#?}, out_filename:\n{:#?}",
         allocator.is_some(),
-        deps,
         objects,
         rlibs,
         out_filename
@@ -226,27 +214,18 @@ fn codegen_into_ptx_file(
     // but we dont have our original one because rustc drops tyctxt and codegencx before linking.
     let cx = LlvmMod::new("link_tmp");
 
-    let deps = deps.unwrap_or_default();
-    let mut main_modules = Vec::with_capacity(objects.len());
-    let mut rlib_deps = Vec::with_capacity(rlibs.len());
+    let mut modules = Vec::with_capacity(objects.len() + rlibs.len());
 
     // object files (theyre not object files, they are impostors ඞ) are the bitcode modules produced by this codegen session
     // they *should* be the final crate.
     for obj in objects {
         let bitcode = std::fs::read(obj)?;
-        let name = obj
-            .file_name()
-            .unwrap()
-            .to_str()
-            .expect("non-utf8 bitcode file name")
-            .to_string();
-        main_modules.push((bitcode, name));
+        modules.push(bitcode);
     }
 
     // rlibs are archives that we made previously, they are usually made for crates that are referenced
     // in this crate. We must unpack them and devour their bitcode to link in.
     for rlib in rlibs {
-        // every entry will be a CGU, we need to merge those CGUs into a single module so we can give it to libnvvm to load.
         let mut cgus = Vec::with_capacity(16);
         // just pick the first cgu name as the overall name for now.
         let mut name = String::new();
@@ -272,8 +251,7 @@ fn codegen_into_ptx_file(
             }
         }
 
-        let merged = merge_cgus(cgus, cx.llcx, name.clone());
-        rlib_deps.push((merged, name));
+        modules.extend(cgus);
     }
 
     if let Some(alloc) = allocator {
@@ -283,56 +261,24 @@ fn codegen_into_ptx_file(
                 .clone()
                 .expect("expected obj path for allocator module"),
         )?;
-        main_modules.push((bc, String::from("allocator")));
+        modules.push(bc);
     }
 
-    let sorted_deps = deps.into_iter().filter_map(|x| {
-        for (bc, name) in &rlib_deps {
-            let new_name = name.split_once("-").expect("uh oh rustc changed the format of rlib file names, better go make an angry zulip thread.").0;
-            if new_name == x.replace("-", "_") {
-                return Some((bc.clone(), name.to_string()));
-            }
-        }
-        // HACK(RDambrosio016): If a dep cannot be found then it is probably a proc macro crate.
-        // in which case we should just ignore it and move on, but in the future we should filter out those
-        // deps before linking.
-        None
-    }).collect::<Vec<_>>();
-
     // now that we have our nice bitcode modules, we just need to find libdevice and give our
     // modules to nvvm to make a final ptx file
 
     // we need to actually parse the codegen args again, because codegencx is not available at link time.
     let nvvm_opts = CodegenArgs::from_session(sess).nvvm_options;
 
-    let ptx_bytes =
-        match crate::nvvm::codegen_bitcode_modules(&nvvm_opts, sess, main_modules, sorted_deps) {
-            Ok(bytes) => bytes,
-            Err(err) => {
-                // TODO(RDambrosio016): maybe include the nvvm log with this fatal error
-                sess.fatal(&err.to_string())
-            }
-        };
-
-    std::fs::write(out_filename, ptx_bytes)
-}
-
-/// Merges multiple codegen units into a single codegen unit. This is needed because
-/// we lazy-load modules in dependency order, not sub-crate order, so we need to lazy load
-/// entire modules, not just individual CGUs.
-fn merge_cgus(cgus: Vec<Vec<u8>>, llcx: &Context, crate_name: String) -> Vec<u8> {
-    let cstr = CString::new(crate_name.clone()).unwrap();
-    let module = unsafe { create_module(llcx, &crate_name) };
-    for cgu in cgus {
-        unsafe {
-            let tmp = LLVMRustParseBitcodeForLTO(llcx, cgu.as_ptr(), cgu.len(), cstr.as_ptr())
-                .expect("Failed to parse CGU bitcode");
-            LLVMLinkModules2(module, tmp);
+    let ptx_bytes = match crate::nvvm::codegen_bitcode_modules(&nvvm_opts, sess, modules, cx.llcx) {
+        Ok(bytes) => bytes,
+        Err(err) => {
+            // TODO(RDambrosio016): maybe include the nvvm log with this fatal error
+            sess.fatal(&err.to_string())
         }
-    }
+    };
 
-    let thin = ThinBuffer::new(module);
-    thin.data().to_vec()
+    std::fs::write(out_filename, ptx_bytes)
 }
 
 fn create_archive(sess: &Session, files: &[&Path], metadata: &[u8], out_filename: &Path) {

crates/rustc_codegen_nvvm/src/llvm.rs

@@ -652,6 +652,13 @@ extern "C" {
 // use rustc_codegen_nvvm_macros::trace_ffi_calls;
 // #[trace_ffi_calls]
 extern "C" {
+    pub(crate) fn LLVMAddGlobalDCEPass(PM: &mut PassManager);
+    pub(crate) fn LLVMGetNamedMetadataOperands(M: &Module, name: *const c_char, Dest: *mut &Value);
+    pub(crate) fn LLVMGetNamedMetadataNumOperands(M: &Module, name: *const c_char) -> c_uint;
+    pub(crate) fn LLVMGetMDNodeOperands(V: &Value, Dest: *mut &Value);
+    pub(crate) fn LLVMGetMDNodeNumOperands(V: &Value) -> c_uint;
+    pub(crate) fn LLVMGetFirstFunction(M: &Module) -> Option<&Value>;
+    pub(crate) fn LLVMGetNextFunction(Fn: &Value) -> Option<&Value>;
     pub(crate) fn LLVMAddGlobalInAddressSpace<'a>(
         M: &'a Module,
         Ty: &'a Type,