solana_perf/
perf_libs.rs

1use {
2    core::ffi::c_void,
3    dlopen2::symbor::{Container, SymBorApi, Symbol},
4    log::*,
5    std::{
6        env,
7        ffi::OsStr,
8        fs,
9        os::raw::{c_int, c_uint},
10        path::{Path, PathBuf},
11        sync::Once,
12    },
13};
14
15#[repr(C)]
16pub struct Elems {
17    pub elems: *const u8,
18    pub num: u32,
19}
20
21#[derive(SymBorApi)]
22pub struct Api<'a> {
23    pub ed25519_init: Symbol<'a, unsafe extern "C" fn() -> bool>,
24    pub ed25519_set_verbose: Symbol<'a, unsafe extern "C" fn(val: bool)>,
25
26    #[allow(clippy::type_complexity)]
27    pub ed25519_verify_many: Symbol<
28        'a,
29        unsafe extern "C" fn(
30            vecs: *const Elems,
31            num: u32,          //number of vecs
32            message_size: u32, //size of each element inside the elems field of the vec
33            total_packets: u32,
34            total_signatures: u32,
35            message_lens: *const u32,
36            pubkey_offsets: *const u32,
37            signature_offsets: *const u32,
38            signed_message_offsets: *const u32,
39            out: *mut u8, //combined length of all the items in vecs
40            use_non_default_stream: u8,
41        ) -> u32,
42    >,
43
44    #[allow(clippy::type_complexity)]
45    pub ed25519_sign_many: Symbol<
46        'a,
47        unsafe extern "C" fn(
48            vecs: *mut Elems,
49            num: u32,          //number of vecs
50            message_size: u32, //size of each element inside the elems field of the vec
51            total_packets: u32,
52            total_signatures: u32,
53            message_lens: *const u32,
54            pubkey_offsets: *const u32,
55            privkey_offsets: *const u32,
56            signed_message_offsets: *const u32,
57            sgnatures_out: *mut u8, //combined length of all the items in vecs
58            use_non_default_stream: u8,
59        ) -> u32,
60    >,
61
62    pub poh_verify_many: Symbol<
63        'a,
64        unsafe extern "C" fn(
65            hashes: *mut u8,
66            num_hashes_arr: *const u64,
67            num_elems: usize,
68            use_non_default_stream: u8,
69        ) -> c_int,
70    >,
71
72    pub cuda_host_register:
73        Symbol<'a, unsafe extern "C" fn(ptr: *mut c_void, size: usize, flags: c_uint) -> c_int>,
74
75    pub cuda_host_unregister: Symbol<'a, unsafe extern "C" fn(ptr: *mut c_void) -> c_int>,
76
77    pub ed25519_get_checked_scalar:
78        Symbol<'a, unsafe extern "C" fn(out_scalar: *mut u8, in_scalar: *const u8) -> c_int>,
79
80    pub ed25519_check_packed_ge_small_order:
81        Symbol<'a, unsafe extern "C" fn(packed_ge: *const u8) -> c_int>,
82}
83
84static mut API: Option<Container<Api>> = None;
85
86fn init(name: &OsStr) {
87    static INIT_HOOK: Once = Once::new();
88
89    info!("Loading {:?}", name);
90    unsafe {
91        INIT_HOOK.call_once(|| {
92            API = Some(Container::load(name).unwrap_or_else(|err| {
93                error!("Unable to load {:?}: {}", name, err);
94                std::process::exit(1);
95            }));
96        })
97    }
98}
99
100pub fn locate_perf_libs() -> Option<PathBuf> {
101    let exe = env::current_exe().expect("Unable to get executable path");
102    let perf_libs = exe.parent().unwrap().join("perf-libs");
103    if perf_libs.is_dir() {
104        info!("perf-libs found at {:?}", perf_libs);
105        return Some(perf_libs);
106    }
107    warn!("{:?} does not exist", perf_libs);
108    None
109}
110
111fn find_cuda_home(perf_libs_path: &Path) -> Option<PathBuf> {
112    if let Ok(cuda_home) = env::var("CUDA_HOME") {
113        let path = PathBuf::from(cuda_home);
114        if path.is_dir() {
115            info!("Using CUDA_HOME: {:?}", path);
116            return Some(path);
117        }
118        warn!("Ignoring CUDA_HOME, not a path: {:?}", path);
119    }
120
121    // Search /usr/local for a `cuda-` directory that matches a perf-libs subdirectory
122    for entry in fs::read_dir(perf_libs_path).unwrap().flatten() {
123        let path = entry.path();
124        if !path.is_dir() {
125            continue;
126        }
127        let dir_name = path.file_name().unwrap().to_str().unwrap_or("");
128        if !dir_name.starts_with("cuda-") {
129            continue;
130        }
131
132        let cuda_home: PathBuf = ["/", "usr", "local", dir_name].iter().collect();
133        if !cuda_home.is_dir() {
134            continue;
135        }
136
137        info!("CUDA installation found at {:?}", cuda_home);
138        return Some(cuda_home);
139    }
140    None
141}
142
143pub fn append_to_ld_library_path(mut ld_library_path: String) {
144    if let Ok(env_value) = env::var("LD_LIBRARY_PATH") {
145        ld_library_path.push(':');
146        ld_library_path.push_str(&env_value);
147    }
148    info!("setting ld_library_path to: {:?}", ld_library_path);
149    env::set_var("LD_LIBRARY_PATH", ld_library_path);
150}
151
152pub fn init_cuda() {
153    if let Some(perf_libs_path) = locate_perf_libs() {
154        if let Some(cuda_home) = find_cuda_home(&perf_libs_path) {
155            let cuda_lib64_dir = cuda_home.join("lib64");
156            if cuda_lib64_dir.is_dir() {
157                // Prefix LD_LIBRARY_PATH with $CUDA_HOME/lib64 directory
158                // to ensure the correct CUDA version is used
159                append_to_ld_library_path(cuda_lib64_dir.to_str().unwrap_or("").to_string())
160            } else {
161                warn!("CUDA lib64 directory does not exist: {:?}", cuda_lib64_dir);
162            }
163
164            let libcuda_crypt = perf_libs_path
165                .join(cuda_home.file_name().unwrap())
166                .join("libcuda-crypt.so");
167            return init(libcuda_crypt.as_os_str());
168        } else {
169            warn!("CUDA installation not found");
170        }
171    }
172
173    // Last resort!  Blindly load the shared object and hope it all works out
174    init(OsStr::new("libcuda-crypt.so"))
175}
176
177pub fn api() -> Option<&'static Container<Api<'static>>> {
178    {
179        static INIT_HOOK: Once = Once::new();
180        INIT_HOOK.call_once(|| {
181            if std::env::var("TEST_PERF_LIBS_CUDA").is_ok() {
182                init_cuda();
183            }
184        })
185    }
186
187    unsafe { API.as_ref() }
188}