diff --git a/crates/cuda_builder/src/lib.rs b/crates/cuda_builder/src/lib.rs index 1e47138..e5b1e60 100644 --- a/crates/cuda_builder/src/lib.rs +++ b/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, + /// 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, } 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) -> 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 { @@ -404,6 +415,11 @@ fn invoke_rustc(builder: &CudaBuilder) -> Result { 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); diff --git a/crates/rustc_codegen_nvvm/CHANGELOG.md b/crates/rustc_codegen_nvvm/CHANGELOG.md index e7e7a99..327a566 100644 --- a/crates/rustc_codegen_nvvm/CHANGELOG.md +++ b/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 diff --git a/crates/rustc_codegen_nvvm/src/context.rs b/crates/rustc_codegen_nvvm/src/context.rs index ddecf89..49ea650 100644 --- a/crates/rustc_codegen_nvvm/src/context.rs +++ b/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, pub override_libm: bool, + pub final_module_path: Option, } 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"), + )); } } diff --git a/crates/rustc_codegen_nvvm/src/link.rs b/crates/rustc_codegen_nvvm/src/link.rs index 2c3a94c..0d5e484 100644 --- a/crates/rustc_codegen_nvvm/src/link.rs +++ b/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 diff --git a/crates/rustc_codegen_nvvm/src/nvvm.rs b/crates/rustc_codegen_nvvm/src/nvvm.rs index 1b2814a..00f2c84 100644 --- a/crates/rustc_codegen_nvvm/src/nvvm.rs +++ b/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>, 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