rustc_codegen_llvm/llvm/
offload_ffi.rs

1use std::ffi::{CStr, c_char};
2use std::path::PathBuf;
3use std::sync::OnceLock;
4
5use super::ffi::{Module, Value};
6
7type LLVMRustOffloadMapperFn = unsafe extern "C" fn(&Value, &Value, *const &Value);
8type LLVMRustOffloadWrapImagesFn =
9    unsafe extern "C" fn(&Module, *const c_char, *const c_char) -> bool;
10
11use rustc_fs_util::path_to_c_string;
12use rustc_session::config::host_tuple;
13use rustc_session::filesearch;
14
15use crate::llvm;
16
17pub(crate) struct RustOffloadWrapper {
18    LLVMRustOffloadMapper: LLVMRustOffloadMapperFn,
19    LLVMRustOffloadWrapImages: LLVMRustOffloadWrapImagesFn,
20    lld_path: Option<PathBuf>,
21    // Keep the dynamic library loaded while the function pointers are used.
22    _lib: libloading::Library,
23}
24
25#[derive(#[automatically_derived]
impl ::core::fmt::Debug for RustOffloadLibraryError {
    #[inline]
    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
        match self {
            Self::NotFound { err: __self_0 } =>
                ::core::fmt::Formatter::debug_struct_field1_finish(f,
                    "NotFound", "err", &__self_0),
            Self::LoadFailed { err: __self_0 } =>
                ::core::fmt::Formatter::debug_struct_field1_finish(f,
                    "LoadFailed", "err", &__self_0),
        }
    }
}Debug)]
26pub(crate) enum RustOffloadLibraryError {
27    NotFound { err: String },
28    LoadFailed { err: String },
29}
30
31impl From<libloading::Error> for RustOffloadLibraryError {
32    fn from(err: libloading::Error) -> Self {
33        Self::LoadFailed { err: ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("{0:?}", err))
    })format!("{err:?}") }
34    }
35}
36
37static OFFLOAD_INSTANCE: OnceLock<RustOffloadWrapper> = OnceLock::new();
38
39impl RustOffloadWrapper {
40    pub(crate) fn get_or_init(
41        sysroot: &rustc_session::config::Sysroot,
42    ) -> Result<&'static RustOffloadWrapper, RustOffloadLibraryError> {
43        OFFLOAD_INSTANCE.get_or_try_init(|| {
44            let w = Self::call_dynamic(sysroot)?;
45            Ok(w)
46        })
47    }
48
49    pub(crate) fn get_instance() -> &'static RustOffloadWrapper {
50        OFFLOAD_INSTANCE
51            .get()
52            .expect("RustOffloadWrapper not initialized. Call get_or_init with sysroot first.")
53    }
54
55    pub(crate) unsafe fn llvm_rust_offload_wrapper(&self, v1: &Value, v2: &Value, vs: &[&Value]) {
56        unsafe { (self.LLVMRustOffloadMapper)(v1, v2, vs.as_ptr()) }
57    }
58
59    pub(crate) unsafe fn llvm_rust_offload_wrap_images(
60        &self,
61        host_m: &Module,
62        device_bitcode_path: &CStr,
63    ) -> bool {
64        let lld_c = self.lld_path.as_deref().map(path_to_c_string).unwrap_or_default();
65        unsafe {
66            (self.LLVMRustOffloadWrapImages)(host_m, lld_c.as_ptr(), device_bitcode_path.as_ptr())
67        }
68    }
69
70    fn call_dynamic(
71        sysroot: &rustc_session::config::Sysroot,
72    ) -> Result<Self, RustOffloadLibraryError> {
73        let (rust_offload_path, lld_path) = Self::get_offload_and_lld_paths(sysroot)?;
74        let lib = unsafe { libloading::Library::new(rust_offload_path)? };
75
76        let llvm_rust_offload_wrapper =
77            *unsafe { lib.get::<LLVMRustOffloadMapperFn>(b"LLVMRustOffloadMapper\0")? };
78        let llvm_rust_offload_wrap_images =
79            *unsafe { lib.get::<LLVMRustOffloadWrapImagesFn>(b"LLVMRustOffloadWrapImages\0")? };
80
81        Ok(Self {
82            LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
83            LLVMRustOffloadWrapImages: llvm_rust_offload_wrap_images,
84            lld_path,
85            _lib: lib,
86        })
87    }
88
89    fn get_offload_and_lld_paths(
90        sysroot: &rustc_session::config::Sysroot,
91    ) -> Result<(PathBuf, Option<PathBuf>), RustOffloadLibraryError> {
92        let llvm_version_major = llvm::LLVMRustVersionMajor();
93        let mut searched = Vec::new();
94
95        for root in sysroot.all_paths() {
96            let rust_offload_path = filesearch::make_target_lib_path(root, host_tuple())
97                .join(::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("libRustOffload-{0}",
                llvm_version_major))
    })format!("libRustOffload-{llvm_version_major}"))
98                .with_extension(std::env::consts::DLL_EXTENSION);
99
100            if !rust_offload_path.is_file() {
101                searched.push(rust_offload_path);
102                continue;
103            }
104
105            let lld_path = filesearch::make_target_bin_path(root, host_tuple())
106                .join(::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("rust-lld{0}",
                std::env::consts::EXE_SUFFIX))
    })format!("rust-lld{}", std::env::consts::EXE_SUFFIX));
107            let lld_path = lld_path.is_file().then_some(lld_path);
108
109            return Ok((rust_offload_path, lld_path));
110        }
111
112        Err(RustOffloadLibraryError::NotFound {
113            err: ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("could not find libRustOffload-{1} in the sysroot candidates:\n* {0}",
                searched.iter().map(|p|
                                p.display().to_string()).collect::<Vec<_>>().join("\n* "),
                llvm_version_major))
    })format!(
114                "could not find libRustOffload-{llvm_version_major} in the sysroot candidates:\n* {}",
115                searched.iter().map(|p| p.display().to_string()).collect::<Vec<_>>().join("\n* ")
116            ),
117        })
118    }
119}
rustc_codegen_llvm/llvm/offload_ffi.rs

rustc_codegen_llvm/llvm/
offload_ffi.rs
