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 =
8    for<'ll> unsafe extern "C" fn(Value<'ll>, Value<'ll>, *const Value<'ll>);
9type LLVMRustOffloadWrapImagesFn =
10    for<'ll> unsafe extern "C" fn(Module<'ll>, *const c_char) -> bool;
11
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    // Keep the dynamic library loaded while the function pointers are used.
21    _lib: libloading::Library,
22}
23
24#[derive(#[automatically_derived]impl ::core::fmt::Debug for RustOffloadLibraryError {
    #[inline]
    #[rustc_trivial_field_reads]
    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)]
25pub(crate) enum RustOffloadLibraryError {
26    NotFound { err: String },
27    LoadFailed { err: String },
28}
29
30impl From<libloading::Error> for RustOffloadLibraryError {
31    fn from(err: libloading::Error) -> Self {
32        Self::LoadFailed { err: ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("{0:?}", err))
    })format!("{err:?}") }
33    }
34}
35
36static OFFLOAD_INSTANCE: OnceLock<RustOffloadWrapper> = OnceLock::new();
37
38impl RustOffloadWrapper {
39    pub(crate) fn get_or_init(
40        sysroot: &rustc_session::config::Sysroot,
41    ) -> Result<&'static RustOffloadWrapper, RustOffloadLibraryError> {
42        OFFLOAD_INSTANCE.get_or_try_init(|| {
43            let w = Self::call_dynamic(sysroot)?;
44            Ok(w)
45        })
46    }
47
48    pub(crate) fn get_instance() -> &'static RustOffloadWrapper {
49        OFFLOAD_INSTANCE
50            .get()
51            .expect("RustOffloadWrapper not initialized. Call get_or_init with sysroot first.")
52    }
53
54    pub(crate) unsafe fn llvm_rust_offload_wrapper<'ll>(
55        &self,
56        v1: Value<'ll>,
57        v2: Value<'ll>,
58        vs: &[Value<'ll>],
59    ) {
60        unsafe { (self.LLVMRustOffloadMapper)(v1, v2, vs.as_ptr()) }
61    }
62
63    pub(crate) unsafe fn llvm_rust_offload_wrap_images<'ll>(
64        &self,
65        host_m: Module<'ll>,
66        device_image_path: &CStr,
67    ) -> bool {
68        unsafe { (self.LLVMRustOffloadWrapImages)(host_m, device_image_path.as_ptr()) }
69    }
70
71    fn call_dynamic(
72        sysroot: &rustc_session::config::Sysroot,
73    ) -> Result<Self, RustOffloadLibraryError> {
74        let rust_offload_path = Self::get_offload_path(sysroot)?;
75        let lib = unsafe { libloading::Library::new(rust_offload_path)? };
76
77        let llvm_rust_offload_wrapper =
78            *unsafe { lib.get::<LLVMRustOffloadMapperFn>(b"LLVMRustOffloadMapper\0")? };
79        let llvm_rust_offload_wrap_images =
80            *unsafe { lib.get::<LLVMRustOffloadWrapImagesFn>(b"LLVMRustOffloadWrapImages\0")? };
81
82        Ok(Self {
83            LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
84            LLVMRustOffloadWrapImages: llvm_rust_offload_wrap_images,
85            _lib: lib,
86        })
87    }
88
89    fn get_offload_path(
90        sysroot: &rustc_session::config::Sysroot,
91    ) -> Result<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            return Ok(rust_offload_path);
106        }
107
108        Err(RustOffloadLibraryError::NotFound {
109            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!(
110                "could not find libRustOffload-{llvm_version_major} in the sysroot candidates:\n* {}",
111                searched.iter().map(|p| p.display().to_string()).collect::<Vec<_>>().join("\n* ")
112            ),
113        })
114    }
115}
rustc_codegen_llvm/llvm/offload_ffi.rs

rustc_codegen_llvm/llvm/
offload_ffi.rs
