Skip to main content

rustc_codegen_llvm/llvm/
offload_ffi.rs

1use std::ffi::{CStr, c_char};
2use std::sync::OnceLock;
3
4use super::ffi::{Module, TargetMachine, Value};
5
6type LLVMRustBundleImagesFn = unsafe extern "C" fn(&Module, &TargetMachine, *const c_char) -> bool;
7type LLVMRustOffloadEmbedBufferInModuleFn = unsafe extern "C" fn(&Module, *const c_char) -> bool;
8type LLVMRustOffloadMapperFn = unsafe extern "C" fn(&Value, &Value, *const &Value);
9
10use rustc_session::config::host_tuple;
11use rustc_session::filesearch;
12
13use crate::llvm::LLVMRustVersionMajor;
14
15pub(crate) struct RustOffloadWrapper {
16    LLVMRustBundleImages: LLVMRustBundleImagesFn,
17    LLVMRustOffloadEmbedBufferInModule: LLVMRustOffloadEmbedBufferInModuleFn,
18    LLVMRustOffloadMapper: LLVMRustOffloadMapperFn,
19    // Keep the dynamic library loaded while the function pointers are used.
20    _lib: libloading::Library,
21}
22
23#[derive(#[automatically_derived]
impl ::core::fmt::Debug for RustOffloadLibraryError {
    #[inline]
    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
        match self {
            RustOffloadLibraryError::NotFound { err: __self_0 } =>
                ::core::fmt::Formatter::debug_struct_field1_finish(f,
                    "NotFound", "err", &__self_0),
            RustOffloadLibraryError::LoadFailed { err: __self_0 } =>
                ::core::fmt::Formatter::debug_struct_field1_finish(f,
                    "LoadFailed", "err", &__self_0),
        }
    }
}Debug)]
24pub(crate) enum RustOffloadLibraryError {
25    NotFound { err: String },
26    LoadFailed { err: String },
27}
28
29impl From<libloading::Error> for RustOffloadLibraryError {
30    fn from(err: libloading::Error) -> Self {
31        Self::LoadFailed { err: ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("{0:?}", err))
    })format!("{err:?}") }
32    }
33}
34
35static OFFLOAD_INSTANCE: OnceLock<RustOffloadWrapper> = OnceLock::new();
36
37impl RustOffloadWrapper {
38    pub(crate) fn get_or_init(
39        sysroot: &rustc_session::config::Sysroot,
40    ) -> Result<&'static RustOffloadWrapper, RustOffloadLibraryError> {
41        OFFLOAD_INSTANCE.get_or_try_init(|| {
42            let w = Self::call_dynamic(sysroot)?;
43            Ok(w)
44        })
45    }
46
47    pub(crate) fn get_instance() -> &'static RustOffloadWrapper {
48        OFFLOAD_INSTANCE
49            .get()
50            .expect("RustOffloadWrapper not initialized. Call get_or_init with sysroot first.")
51    }
52
53    pub(crate) unsafe fn llvm_rust_bundle_images(
54        &self,
55        m: &Module,
56        tm: &TargetMachine,
57        c: &CStr,
58    ) -> bool {
59        unsafe { (self.LLVMRustBundleImages)(m, tm, c.as_ptr()) }
60    }
61
62    pub(crate) unsafe fn llvm_rust_offload_embed_buffer_in_module(
63        &self,
64        m: &Module,
65        i: &CStr,
66    ) -> bool {
67        unsafe { (self.LLVMRustOffloadEmbedBufferInModule)(m, i.as_ptr()) }
68    }
69
70    pub(crate) unsafe fn llvm_rust_offload_wrapper(&self, v1: &Value, v2: &Value, vs: &[&Value]) {
71        unsafe { (self.LLVMRustOffloadMapper)(v1, v2, vs.as_ptr()) }
72    }
73
74    fn call_dynamic(
75        sysroot: &rustc_session::config::Sysroot,
76    ) -> Result<Self, RustOffloadLibraryError> {
77        let rust_offload_path = Self::get_rust_offload_path(sysroot)?;
78        let lib = unsafe { libloading::Library::new(rust_offload_path)? };
79
80        let llvm_rust_bundle_images =
81            *unsafe { lib.get::<LLVMRustBundleImagesFn>(b"LLVMRustBundleImages\0")? };
82        let llvm_rust_offload_embed_buffer_in_module = *unsafe {
83            lib.get::<LLVMRustOffloadEmbedBufferInModuleFn>(
84                b"LLVMRustOffloadEmbedBufferInModule\0",
85            )?
86        };
87        let llvm_rust_offload_wrapper =
88            *unsafe { lib.get::<LLVMRustOffloadMapperFn>(b"LLVMRustOffloadMapper\0")? };
89
90        Ok(Self {
91            LLVMRustBundleImages: llvm_rust_bundle_images,
92            LLVMRustOffloadEmbedBufferInModule: llvm_rust_offload_embed_buffer_in_module,
93            LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
94            _lib: lib,
95        })
96    }
97
98    fn get_rust_offload_path(
99        sysroot: &rustc_session::config::Sysroot,
100    ) -> Result<String, RustOffloadLibraryError> {
101        let llvm_version_major = unsafe { LLVMRustVersionMajor() };
102
103        let path_buf = sysroot
104            .all_paths()
105            .find_map(|p| {
106                let candidate = filesearch::make_target_lib_path(p, host_tuple())
107                    .join(::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("libRustOffload-{0}",
                llvm_version_major))
    })format!("libRustOffload-{}", llvm_version_major))
108                    .with_extension(std::env::consts::DLL_EXTENSION);
109
110                candidate.exists().then_some(candidate)
111            })
112            .ok_or_else(|| {
113                let candidates = sysroot
114                    .all_paths()
115                    .map(|p| p.join("lib").display().to_string())
116                    .collect::<Vec<String>>()
117                    .join("\n* ");
118                RustOffloadLibraryError::NotFound {
119                    err: ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("failed to find a `libRustOffload-{0}` in the sysroot candidates:\n* {1}",
                llvm_version_major, candidates))
    })format!(
120                        "failed to find a `libRustOffload-{llvm_version_major}` \
121                    in the sysroot candidates:\n* {candidates}"
122                    ),
123                }
124            })?;
125
126        Ok(path_buf
127            .to_str()
128            .ok_or_else(|| RustOffloadLibraryError::LoadFailed {
129                err: ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("invalid UTF-8 in path: {0}",
                path_buf.display()))
    })format!("invalid UTF-8 in path: {}", path_buf.display()),
130            })?
131            .to_string())
132    }
133}