offload: pass the linked device image to the host pass

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -654,22 +654,17 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
         None
     };
 
-    fn handle_offload<'ll>(cx: &SimpleCx<'ll>, old_fn: &'ll llvm::Value) {
+    // Wraps `old_fn` into a new kernel that takes the implicit `dyn_ptr` as its first argument.
+    // Returns the new function, or `None` if nothing was wrapped
+    fn handle_offload<'ll>(
+        cx: &SimpleCx<'ll>,
+        old_fn: &'ll llvm::Value,
+    ) -> Option<&'ll llvm::Value> {
         let old_fn_ty = cx.get_type_of_global(old_fn);
         let old_param_types = cx.func_params_types(old_fn_ty);
         let old_param_count = old_param_types.len();
         if old_param_count == 0 {
-            return;
-        }
-
-        let first_param = llvm::get_param(old_fn, 0);
-        let c_name = llvm::get_value_name(first_param);
-        let first_arg_name = str::from_utf8(&c_name).unwrap();
-        // We might call llvm_optimize (and thus this code) multiple times on the same IR,
-        // but we shouldn't add this helper ptr multiple times.
-        // FIXME(offload): This could break if the user calls his first argument `dyn_ptr`.
-        if first_arg_name == "dyn_ptr" {
-            return;
+            return None;
         }
 
         // Create the new parameter list, with ptr as the first argument
@@ -744,6 +739,7 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
         }
         // Now we can re-use the old name, without name collision.
         llvm::set_value_name(new_fn, &name);
+        Some(new_fn)
     }
 
     if cgcx.target_is_like_gpu
@@ -752,10 +748,13 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
         let cx = SimpleCx::new(llmod, llcx, cgcx.pointer_size);
         for func in cx.get_functions() {
             let offload_kernel = "offload-kernel";
-            if attributes::has_string_attr(func, offload_kernel) {
-                handle_offload(&cx, func);
+            if !attributes::has_string_attr(func, offload_kernel) {
+                continue;
+            }
+            match handle_offload(&cx, func) {
+                Some(new_fn) => attributes::remove_string_attr_from_llfn(new_fn, offload_kernel),
+                None => attributes::remove_string_attr_from_llfn(func, offload_kernel),
             }
-            attributes::remove_string_attr_from_llfn(func, offload_kernel);
         }
     }
 
@@ -819,8 +818,9 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
     };
 
     // This assumes that we previously compiled our kernels for a gpu target, which created a
-    // `<crate>.bc` artifact. The user is supposed to provide us with a path to this artifact, we
-    // don't need any other artifacts from the previous run.
+    // `<crate>.ptx` (NVPTX) or `<crate>.elf` (AMDGPU) artifact.
+    // The user is supposed to provide us with a path to this artifact,
+    // we don't need any other artifacts from the previous run.
     if !cgcx.target_is_like_gpu && is_final_stage {
         if let Some(device_path) = config
             .offload
@@ -830,15 +830,18 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
             let device_pathbuf = PathBuf::from(device_path);
             if device_pathbuf.is_relative() {
                 dcx.emit_err(crate::diagnostics::OffloadWithoutAbsPath);
-            } else if device_pathbuf.extension().and_then(|e| e.to_str()) != Some("bc") {
+            } else if !matches!(
+                device_pathbuf.extension().and_then(|e| e.to_str()),
+                Some("ptx" | "elf")
+            ) {
                 dcx.emit_err(crate::diagnostics::OffloadWrongFileName);
             } else if !device_pathbuf.exists() {
                 dcx.emit_err(crate::diagnostics::OffloadNonexistingPath);
             }
-            let device_bitcode_c = path_to_c_string(device_pathbuf.as_path());
+            let device_image_c = path_to_c_string(device_pathbuf.as_path());
             let ok = unsafe {
                 llvm::RustOffloadWrapper::get_instance()
-                    .llvm_rust_offload_wrap_images(llmod, device_bitcode_c.as_c_str())
+                    .llvm_rust_offload_wrap_images(llmod, device_image_c.as_c_str())
             };
             if !ok {
                 dcx.emit_err(crate::diagnostics::OffloadWrapImagesFailed);

compiler/rustc_codegen_llvm/src/diagnostics.rs

@@ -73,7 +73,7 @@ pub(crate) struct RustOffloadComponentMissing {
 
 #[derive(Diagnostic)]
 #[diag(
-    "using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/<crate>.bc>"
+    "using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/<crate>.ptx or <crate>.elf>"
 )]
 pub(crate) struct OffloadWithoutEnable;
 
@@ -82,18 +82,20 @@ pub(crate) struct OffloadWithoutEnable;
 pub(crate) struct OffloadWithoutFatLTO;
 
 #[derive(Diagnostic)]
-#[diag("using the `-Z offload=Host=/absolute/path/to/<crate>.bc` flag requires an absolute path")]
+#[diag(
+    "using the `-Z offload=Host=/absolute/path/to/<crate>.ptx` or `<crate>.elf` flag requires an absolute path"
+)]
 pub(crate) struct OffloadWithoutAbsPath;
 
 #[derive(Diagnostic)]
 #[diag(
-    "using the `-Z offload=Host=/absolute/path/to/<crate>.bc` flag must point to a .bc file produced by the device compilation (--emit=llvm-bc)"
+    "using the `-Z offload=Host=/absolute/path/to/<crate>.ptx` or `<crate>.elf` flag must point to a .ptx or .elf file produced by the device compilation"
 )]
 pub(crate) struct OffloadWrongFileName;
 
 #[derive(Diagnostic)]
 #[diag(
-    "the given path/file to `<crate>.bc` does not exist. Did you forget to run the device compilation first?"
+    "the given path/file to `<crate>.ptx` or `<crate>.elf` does not exist. Did you forget to run the device compilation first?"
 )]
 pub(crate) struct OffloadNonexistingPath;
 

compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs

@@ -5,10 +5,8 @@ use std::sync::OnceLock;
 use super::ffi::{Module, Value};
 
 type LLVMRustOffloadMapperFn = unsafe extern "C" fn(&Value, &Value, *const &Value);
-type LLVMRustOffloadWrapImagesFn =
-    unsafe extern "C" fn(&Module, *const c_char, *const c_char) -> bool;
+type LLVMRustOffloadWrapImagesFn = unsafe extern "C" fn(&Module, *const c_char) -> bool;
 
-use rustc_fs_util::path_to_c_string;
 use rustc_session::config::host_tuple;
 use rustc_session::filesearch;
 
@@ -17,7 +15,6 @@ use crate::llvm;
 pub(crate) struct RustOffloadWrapper {
     LLVMRustOffloadMapper: LLVMRustOffloadMapperFn,
     LLVMRustOffloadWrapImages: LLVMRustOffloadWrapImagesFn,
-    lld_path: Option<PathBuf>,
     // Keep the dynamic library loaded while the function pointers are used.
     _lib: libloading::Library,
 }
@@ -59,18 +56,15 @@ impl RustOffloadWrapper {
     pub(crate) unsafe fn llvm_rust_offload_wrap_images(
         &self,
         host_m: &Module,
-        device_bitcode_path: &CStr,
+        device_image_path: &CStr,
     ) -> bool {
-        let lld_c = self.lld_path.as_deref().map(path_to_c_string).unwrap_or_default();
-        unsafe {
-            (self.LLVMRustOffloadWrapImages)(host_m, lld_c.as_ptr(), device_bitcode_path.as_ptr())
-        }
+        unsafe { (self.LLVMRustOffloadWrapImages)(host_m, device_image_path.as_ptr()) }
     }
 
     fn call_dynamic(
         sysroot: &rustc_session::config::Sysroot,
     ) -> Result<Self, RustOffloadLibraryError> {
-        let (rust_offload_path, lld_path) = Self::get_offload_and_lld_paths(sysroot)?;
+        let rust_offload_path = Self::get_offload_path(sysroot)?;
         let lib = unsafe { libloading::Library::new(rust_offload_path)? };
 
         let llvm_rust_offload_wrapper =
@@ -81,14 +75,13 @@ impl RustOffloadWrapper {
         Ok(Self {
             LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
             LLVMRustOffloadWrapImages: llvm_rust_offload_wrap_images,
-            lld_path,
             _lib: lib,
         })
     }
 
-    fn get_offload_and_lld_paths(
+    fn get_offload_path(
         sysroot: &rustc_session::config::Sysroot,
-    ) -> Result<(PathBuf, Option<PathBuf>), RustOffloadLibraryError> {
+    ) -> Result<PathBuf, RustOffloadLibraryError> {
         let llvm_version_major = llvm::LLVMRustVersionMajor();
         let mut searched = Vec::new();
 
@@ -102,11 +95,7 @@ impl RustOffloadWrapper {
                 continue;
             }
 
-            let lld_path = filesearch::make_target_bin_path(root, host_tuple())
-                .join(format!("rust-lld{}", std::env::consts::EXE_SUFFIX));
-            let lld_path = lld_path.is_file().then_some(lld_path);
-
-            return Ok((rust_offload_path, lld_path));
+            return Ok(rust_offload_path);
         }
 
         Err(RustOffloadLibraryError::NotFound {

compiler/rustc_interface/src/passes.rs

@@ -36,7 +36,7 @@ use rustc_parse::lexer::StripTokens;
 use rustc_parse::{new_parser_from_file, new_parser_from_source_str, unwrap_or_emit_fatal};
 use rustc_passes::{abi_test, input_stats, layout_test};
 use rustc_resolve::{Resolver, ResolverOutputs};
-use rustc_session::config::{Input, OutFileName, OutputFilenames, OutputType};
+use rustc_session::config::{Input, Offload, OutFileName, OutputFilenames, OutputType};
 use rustc_session::diagnostics::feature_err;
 use rustc_session::output::{filename_for_input, invalid_output_for_target};
 use rustc_session::search_paths::PathKind;
@@ -1491,6 +1491,16 @@ pub fn collect_crate_types(
         }
     }
 
+    // GPU code is never run as a binary; the offload runtime looks up kernels by
+    // name, and only cdylibs export their symbols.
+    if session.opts.unstable_opts.offload.iter().any(|o| matches!(o, Offload::Device(_))) {
+        for crate_type in &mut base {
+            if *crate_type == CrateType::Executable {
+                *crate_type = CrateType::Cdylib;
+            }
+        }
+    }
+
     base.retain(|crate_type| {
         if invalid_output_for_target(session, *crate_type) {
             session.dcx().emit_warn(diagnostics::UnsupportedCrateTypeForTarget {