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

rustc_codegen_llvm/llvm/
offload_ffi.rs
