solana_perf/
perf_libs.rs

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