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

rustc_codegen_llvm/llvm/
offload_ffi.rs
