Feat: add a way to export the final llvm ir fed to libnvvm

Author: RDambrosio016

crates/cuda_builder/src/lib.rs

@@ -136,6 +136,9 @@ pub struct CudaBuilder {
     pub debug: DebugInfo,
     /// Additional arguments passed to cargo during `cargo build`.
     pub build_args: Vec<String>,
+    /// An optional path where to dump LLVM IR of the final output the codegen will feed to libnvvm. Usually
+    /// used for debugging.
+    pub final_module_path: Option<PathBuf>,
 }
 
 impl CudaBuilder {
@@ -156,6 +159,7 @@ impl CudaBuilder {
             override_libm: true,
             debug: DebugInfo::None,
             build_args: vec![],
+            final_module_path: None,
         }
     }
 
@@ -282,6 +286,13 @@ impl CudaBuilder {
         self
     }
 
+    /// An optional path where to dump LLVM IR of the final output the codegen will feed to libnvvm. Usually
+    /// used for debugging.
+    pub fn final_module_path(mut self, path: impl AsRef<Path>) -> Self {
+        self.final_module_path = Some(path.as_ref().to_path_buf());
+        self
+    }
+
     /// Runs rustc to build the codegen and codegens the gpu crate, returning the path of the final
     /// ptx file. If [`ptx_file_copy_path`](Self::ptx_file_copy_path) is set, this returns the copied path.
     pub fn build(self) -> Result<PathBuf, CudaBuilderError> {
@@ -404,6 +415,11 @@ fn invoke_rustc(builder: &CudaBuilder) -> Result<PathBuf, CudaBuilderError> {
         llvm_args.push("--override-libm".to_string());
     }
 
+    if let Some(path) = &builder.final_module_path {
+        llvm_args.push("--final-module-path".to_string());
+        llvm_args.push(path.to_str().unwrap().to_string());
+    }
+
     if builder.debug != DebugInfo::None {
         let (nvvm_flag, rustc_flag) = builder.debug.into_nvvm_and_rustc_options();
         llvm_args.push(nvvm_flag);

crates/rustc_codegen_nvvm/CHANGELOG.md

@@ -6,6 +6,7 @@ Notable changes to this project will be documented in this file.
 
 - Added symbols for cuda_std to link to for warp intrinsics.
 - Completely remove support for 32-bit CUDA (it was broken and it is essentially unused nowadays).
+- Add a way to export the final llvm ir module fed to libnvvm.
 
 ## 0.2.3 - 1/2/22
 

crates/rustc_codegen_nvvm/src/context.rs

@@ -31,6 +31,7 @@ use std::cell::{Cell, RefCell};
 use std::ffi::CStr;
 use std::hash::BuildHasherDefault;
 use std::os::raw::c_char;
+use std::path::PathBuf;
 use std::ptr::null;
 use std::str::FromStr;
 use tracing::{debug, trace};
@@ -523,6 +524,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 pub struct CodegenArgs {
     pub nvvm_options: Vec<NvvmOption>,
     pub override_libm: bool,
+    pub final_module_path: Option<PathBuf>,
 }
 
 impl CodegenArgs {
@@ -535,11 +537,15 @@ impl CodegenArgs {
         // TODO: replace this with a "proper" arg parser.
         let mut cg_args = Self::default();
 
-        for arg in args {
+        for (idx, arg) in args.iter().enumerate() {
             if let Ok(flag) = NvvmOption::from_str(arg) {
                 cg_args.nvvm_options.push(flag);
             } else if arg == "--override-libm" {
                 cg_args.override_libm = true;
+            } else if arg == "--final-module-path" {
+                cg_args.final_module_path = Some(PathBuf::from(
+                    args.get(idx + 1).expect("No path for --final-module-path"),
+                ));
             }
         }
 

crates/rustc_codegen_nvvm/src/link.rs

@@ -256,9 +256,9 @@ fn codegen_into_ptx_file(
     // 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 args = CodegenArgs::from_session(sess);
 
-    let ptx_bytes = match crate::nvvm::codegen_bitcode_modules(&nvvm_opts, sess, modules, cx.llcx) {
+    let ptx_bytes = match crate::nvvm::codegen_bitcode_modules(&args, sess, modules, cx.llcx) {
         Ok(bytes) => bytes,
         Err(err) => {
             // TODO(RDambrosio016): maybe include the nvvm log with this fatal error

crates/rustc_codegen_nvvm/src/nvvm.rs

@@ -1,11 +1,14 @@
 //! Final steps in codegen, coalescing modules and feeding them to libnvvm.
 
+use crate::back::demangle_callback;
 use crate::builder::unnamed;
+use crate::context::CodegenArgs;
 use crate::llvm::*;
 use crate::lto::ThinBuffer;
 use find_cuda_helper::find_cuda_root;
 use nvvm::*;
 use rustc_codegen_ssa::traits::ThinBufferMethods;
+use rustc_fs_util::path_to_c_string;
 use rustc_session::{config::DebugInfo, Session};
 use std::ffi::OsStr;
 use std::fmt::Display;
@@ -50,7 +53,7 @@ impl Display for CodegenErr {
 /// Note that this will implicitly try to find libdevice and add it, so don't do that
 /// step before this. It will fatal error if it cannot find it.
 pub fn codegen_bitcode_modules(
-    opts: &[NvvmOption],
+    args: &CodegenArgs,
     sess: &Session,
     modules: Vec<Vec<u8>>,
     llcx: &Context,
@@ -89,7 +92,16 @@ pub fn codegen_bitcode_modules(
         let node = LLVMMDNodeInContext(llcx, vals.as_ptr(), vals.len() as u32);
 
         LLVMAddNamedMetadataOperand(module, "nvvmir.version\0".as_ptr().cast(), node);
+
+        if let Some(path) = &args.final_module_path {
+            let out_c = path_to_c_string(path);
+            let result = LLVMRustPrintModule(module, out_c.as_ptr(), demangle_callback);
+            result
+                .into_result()
+                .expect("Failed to write final llvm module output");
+        }
     }
+
     let buf = ThinBuffer::new(module);
 
     prog.add_module(buf.data(), "merged".to_string())?;
@@ -122,7 +134,7 @@ pub fn codegen_bitcode_modules(
         );
     }
 
-    let res = match prog.compile(opts) {
+    let res = match prog.compile(&args.nvvm_options) {
         Ok(b) => b,
         Err(error) => {
             // this should never happen, if it does, something went really bad or its a bug on libnvvm's end