dlopen offload

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -720,7 +720,11 @@ pub(crate) unsafe fn llvm_optimize(
         // Here we map the old arguments to the new arguments, with an offset of 1 to make sure
         // that we don't use the newly added `%dyn_ptr`.
         unsafe {
-            llvm::LLVMRustOffloadMapper(old_fn, new_fn, old_args_rebuilt.as_ptr());
+            llvm::RustOffloadWrapper::get_instance().llvm_rust_offload_wrapper(
+                old_fn,
+                new_fn,
+                old_args_rebuilt.as_slice(),
+            );
         }
 
         llvm::set_linkage(new_fn, llvm::get_linkage(old_fn));
@@ -814,16 +818,16 @@ pub(crate) unsafe fn llvm_optimize(
         let device_dir = device_path.parent().unwrap();
         let device_out = device_dir.join("device.bin");
         let device_out_c = path_to_c_string(device_out.as_path());
-        unsafe {
-            // 1) Bundle device module into offload image device.bin (device TM)
-            let ok = llvm::LLVMRustBundleImages(
+        // 1) Bundle device module into offload image device.bin (device TM)
+        let ok = unsafe {
+            llvm::RustOffloadWrapper::get_instance().llvm_rust_bundle_images(
                 module.module_llvm.llmod(),
                 module.module_llvm.tm.raw(),
-                device_out_c.as_ptr(),
-            );
-            if !ok || !device_out.exists() {
-                dcx.emit_err(crate::diagnostics::OffloadBundleImagesFailed);
-            }
+                device_out_c.as_c_str(),
+            )
+        };
+        if !ok || !device_out.exists() {
+            dcx.emit_err(crate::diagnostics::OffloadBundleImagesFailed);
         }
     }
 
@@ -859,8 +863,10 @@ pub(crate) unsafe fn llvm_optimize(
             // We create a full clone of our LLVM host module, since we will embed the device IR
             // into it, and this might break caching or incremental compilation otherwise.
             let llmod2 = llvm::LLVMCloneModule(module.module_llvm.llmod());
-            let ok =
-                unsafe { llvm::LLVMRustOffloadEmbedBufferInModule(llmod2, device_bin_c.as_ptr()) };
+            let ok = unsafe {
+                llvm::RustOffloadWrapper::get_instance()
+                    .llvm_rust_offload_embed_buffer_in_module(llmod2, device_bin_c.as_c_str())
+            };
             if !ok {
                 dcx.emit_err(crate::diagnostics::OffloadEmbedFailed);
             }

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -296,7 +296,7 @@ struct KernelArgsTy {
 
 impl KernelArgsTy {
     const OFFLOAD_VERSION: u64 = 3;
-    const FLAGS: u64 = 0;
+    const FLAGS: u64 = 1 << 6; // Enable StrictBlocksAndThreads
     const TRIPCOUNT: u64 = 0;
     fn new_decl<'ll>(cx: &CodegenCx<'ll, '_>) -> &'ll Type {
         let kernel_arguments_ty = cx.type_named_struct("struct.__tgt_kernel_arguments");

compiler/rustc_codegen_llvm/src/diagnostics.rs

@@ -60,6 +60,19 @@ pub(crate) struct AutoDiffWithoutLto;
 #[diag("using the autodiff feature requires -Z autodiff=Enable")]
 pub(crate) struct AutoDiffWithoutEnable;
 
+#[derive(Diagnostic)]
+#[diag("failed to load our rust offload backend: {$err}")]
+pub(crate) struct RustOffloadComponentUnavailable {
+    pub err: String,
+}
+
+#[derive(Diagnostic)]
+#[diag("rust offload backend not found in the sysroot: {$err}")]
+#[note("it will be distributed via rustup in the future")]
+pub(crate) struct RustOffloadComponentMissing {
+    pub err: String,
+}
+
 #[derive(Diagnostic)]
 #[diag(
     "using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/device.bin>"

compiler/rustc_codegen_llvm/src/lib.rs

@@ -373,6 +373,26 @@ impl CodegenBackend for LlvmCodegenBackend {
     }
 
     fn codegen_crate<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Box<dyn Any> {
+        use rustc_session::config::Offload;
+
+        if tcx.sess.opts.unstable_opts.offload.contains(&Offload::Device)
+            || tcx.sess.opts.unstable_opts.offload.iter().any(|o| matches!(o, Offload::Host(_)))
+        {
+            match llvm::RustOffloadWrapper::get_or_init(&tcx.sess.opts.sysroot) {
+                Ok(_) => {}
+                Err(llvm::RustOffloadLibraryError::NotFound { err }) => {
+                    tcx.sess
+                        .dcx()
+                        .emit_fatal(crate::diagnostics::RustOffloadComponentMissing { err });
+                }
+                Err(llvm::RustOffloadLibraryError::LoadFailed { err }) => {
+                    tcx.sess
+                        .dcx()
+                        .emit_fatal(crate::diagnostics::RustOffloadComponentUnavailable { err });
+                }
+            }
+        }
+
         Box::new(rustc_codegen_ssa::base::codegen_crate(LlvmCodegenBackend(()), tcx))
     }
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1713,63 +1713,6 @@ unsafe extern "C" {
     ) -> &'a Value;
 }
 
-#[cfg(feature = "llvm_offload")]
-pub(crate) use self::Offload::*;
-
-#[cfg(feature = "llvm_offload")]
-mod Offload {
-    use super::*;
-    unsafe extern "C" {
-        /// Processes the module and writes it in an offload compatible way into a "device.bin" file.
-        pub(crate) fn LLVMRustBundleImages<'a>(
-            M: &'a Module,
-            TM: &'a TargetMachine,
-            device_bin: *const c_char,
-        ) -> bool;
-        pub(crate) unsafe fn LLVMRustOffloadEmbedBufferInModule<'a>(
-            _M: &'a Module,
-            _device_bin: *const c_char,
-        ) -> bool;
-        pub(crate) fn LLVMRustOffloadMapper<'a>(
-            OldFn: &'a Value,
-            NewFn: &'a Value,
-            RebuiltArgs: *const &Value,
-        );
-    }
-}
-
-#[cfg(not(feature = "llvm_offload"))]
-pub(crate) use self::Offload_fallback::*;
-
-#[cfg(not(feature = "llvm_offload"))]
-mod Offload_fallback {
-    use super::*;
-    /// Processes the module and writes it in an offload compatible way into a "device.bin" file.
-    /// Marked as unsafe to match the real offload wrapper which is unsafe due to FFI.
-    #[allow(unused_unsafe)]
-    pub(crate) unsafe fn LLVMRustBundleImages<'a>(
-        _M: &'a Module,
-        _TM: &'a TargetMachine,
-        _device_bin: *const c_char,
-    ) -> bool {
-        unimplemented!("This rustc version was not built with LLVM Offload support!");
-    }
-    pub(crate) unsafe fn LLVMRustOffloadEmbedBufferInModule<'a>(
-        _M: &'a Module,
-        _device_bin: *const c_char,
-    ) -> bool {
-        unimplemented!("This rustc version was not built with LLVM Offload support!");
-    }
-    #[allow(unused_unsafe)]
-    pub(crate) unsafe fn LLVMRustOffloadMapper<'a>(
-        _OldFn: &'a Value,
-        _NewFn: &'a Value,
-        _RebuiltArgs: *const &Value,
-    ) {
-        unimplemented!("This rustc version was not built with LLVM Offload support!");
-    }
-}
-
 // FFI bindings for `DIBuilder` functions in the LLVM-C API.
 // Try to keep these in the same order as in `llvm/include/llvm-c/DebugInfo.h`.
 //

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -21,8 +21,10 @@ pub(crate) mod diagnostic;
 pub(crate) mod enzyme_ffi;
 mod ffi;
 mod metadata_kind;
+pub(crate) mod offload_ffi;
 
 pub(crate) use self::enzyme_ffi::*;
+pub(crate) use self::offload_ffi::*;
 
 impl LLVMRustResult {
     pub(crate) fn into_result(self) -> Result<(), ()> {

compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs

@@ -0,0 +1,133 @@
+use std::ffi::{CStr, c_char};
+use std::sync::OnceLock;
+
+use super::ffi::{Module, TargetMachine, Value};
+
+type LLVMRustBundleImagesFn = unsafe extern "C" fn(&Module, &TargetMachine, *const c_char) -> bool;
+type LLVMRustOffloadEmbedBufferInModuleFn = unsafe extern "C" fn(&Module, *const c_char) -> bool;
+type LLVMRustOffloadMapperFn = unsafe extern "C" fn(&Value, &Value, *const &Value);
+
+use rustc_session::config::host_tuple;
+use rustc_session::filesearch;
+
+use crate::llvm::LLVMRustVersionMajor;
+
+pub(crate) struct RustOffloadWrapper {
+    LLVMRustBundleImages: LLVMRustBundleImagesFn,
+    LLVMRustOffloadEmbedBufferInModule: LLVMRustOffloadEmbedBufferInModuleFn,
+    LLVMRustOffloadMapper: LLVMRustOffloadMapperFn,
+    // Keep the dynamic library loaded while the function pointers are used.
+    _lib: libloading::Library,
+}
+
+#[derive(Debug)]
+pub(crate) enum RustOffloadLibraryError {
+    NotFound { err: String },
+    LoadFailed { err: String },
+}
+
+impl From<libloading::Error> for RustOffloadLibraryError {
+    fn from(err: libloading::Error) -> Self {
+        Self::LoadFailed { err: format!("{err:?}") }
+    }
+}
+
+static OFFLOAD_INSTANCE: OnceLock<RustOffloadWrapper> = OnceLock::new();
+
+impl RustOffloadWrapper {
+    pub(crate) fn get_or_init(
+        sysroot: &rustc_session::config::Sysroot,
+    ) -> Result<&'static RustOffloadWrapper, RustOffloadLibraryError> {
+        OFFLOAD_INSTANCE.get_or_try_init(|| {
+            let w = Self::call_dynamic(sysroot)?;
+            Ok(w)
+        })
+    }
+
+    pub(crate) fn get_instance() -> &'static RustOffloadWrapper {
+        OFFLOAD_INSTANCE
+            .get()
+            .expect("RustOffloadWrapper not initialized. Call get_or_init with sysroot first.")
+    }
+
+    pub(crate) unsafe fn llvm_rust_bundle_images(
+        &self,
+        m: &Module,
+        tm: &TargetMachine,
+        c: &CStr,
+    ) -> bool {
+        unsafe { (self.LLVMRustBundleImages)(m, tm, c.as_ptr()) }
+    }
+
+    pub(crate) unsafe fn llvm_rust_offload_embed_buffer_in_module(
+        &self,
+        m: &Module,
+        i: &CStr,
+    ) -> bool {
+        unsafe { (self.LLVMRustOffloadEmbedBufferInModule)(m, i.as_ptr()) }
+    }
+
+    pub(crate) unsafe fn llvm_rust_offload_wrapper(&self, v1: &Value, v2: &Value, vs: &[&Value]) {
+        unsafe { (self.LLVMRustOffloadMapper)(v1, v2, vs.as_ptr()) }
+    }
+
+    fn call_dynamic(
+        sysroot: &rustc_session::config::Sysroot,
+    ) -> Result<Self, RustOffloadLibraryError> {
+        let rust_offload_path = Self::get_rust_offload_path(sysroot)?;
+        let lib = unsafe { libloading::Library::new(rust_offload_path)? };
+
+        let llvm_rust_bundle_images =
+            *unsafe { lib.get::<LLVMRustBundleImagesFn>(b"LLVMRustBundleImages\0")? };
+        let llvm_rust_offload_embed_buffer_in_module = *unsafe {
+            lib.get::<LLVMRustOffloadEmbedBufferInModuleFn>(
+                b"LLVMRustOffloadEmbedBufferInModule\0",
+            )?
+        };
+        let llvm_rust_offload_wrapper =
+            *unsafe { lib.get::<LLVMRustOffloadMapperFn>(b"LLVMRustOffloadMapper\0")? };
+
+        Ok(Self {
+            LLVMRustBundleImages: llvm_rust_bundle_images,
+            LLVMRustOffloadEmbedBufferInModule: llvm_rust_offload_embed_buffer_in_module,
+            LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
+            _lib: lib,
+        })
+    }
+
+    fn get_rust_offload_path(
+        sysroot: &rustc_session::config::Sysroot,
+    ) -> Result<String, RustOffloadLibraryError> {
+        let llvm_version_major = unsafe { LLVMRustVersionMajor() };
+
+        let path_buf = sysroot
+            .all_paths()
+            .find_map(|p| {
+                let candidate = filesearch::make_target_lib_path(p, host_tuple())
+                    .join(format!("libRustOffload-{}", llvm_version_major))
+                    .with_extension(std::env::consts::DLL_EXTENSION);
+
+                candidate.exists().then_some(candidate)
+            })
+            .ok_or_else(|| {
+                let candidates = sysroot
+                    .all_paths()
+                    .map(|p| p.join("lib").display().to_string())
+                    .collect::<Vec<String>>()
+                    .join("\n* ");
+                RustOffloadLibraryError::NotFound {
+                    err: format!(
+                        "failed to find a `libRustOffload-{llvm_version_major}` \
+                    in the sysroot candidates:\n* {candidates}"
+                    ),
+                }
+            })?;
+
+        Ok(path_buf
+            .to_str()
+            .ok_or_else(|| RustOffloadLibraryError::LoadFailed {
+                err: format!("invalid UTF-8 in path: {}", path_buf.display()),
+            })?
+            .to_string())
+    }
+}