cuda_interposer/
lib.rs

1use std::{env, ffi::CString, os::raw::c_void, sync::OnceLock};
2use tracing::{debug, warn};
3
4// Re-exports for macros
5pub use libc;
6pub use once_cell;
7pub use paste;
8pub use tracing;
9
10// ─── Library Loading ─────────────────────────────────────────────────────────
11
12struct DlHandle(*mut c_void);
13unsafe impl Send for DlHandle {}
14unsafe impl Sync for DlHandle {}
15
16static CUDA_LIB: OnceLock<DlHandle> = OnceLock::new();
17static CUDART_LIB: OnceLock<DlHandle> = OnceLock::new();
18
19fn get_libcuda() -> *mut c_void {
20    let handle_wrapper = CUDA_LIB.get_or_init(|| unsafe {
21        let mut paths = vec![
22            "/usr/local/cuda/compat/libcuda.so".to_string(),
23            "/usr/lib/x86_64-linux-gnu/libcuda.so".to_string(),
24            "/usr/lib64/libcuda.so".to_string(),
25            "/usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so".to_string(),
26        ];
27
28        if let Some(cuda_home) = env::var_os("CUDA_HOME") {
29            let path = format!("{}/compat/libcuda.so", cuda_home.to_string_lossy());
30            paths.insert(0, path);
31        }
32
33        for path in paths.iter() {
34            let s = CString::new(path.clone()).unwrap();
35            let flags = libc::RTLD_NOW | libc::RTLD_LOCAL | libc::RTLD_NODELETE;
36            let handle = libc::dlopen(s.as_ptr(), flags);
37            if !handle.is_null() {
38                debug!("Loaded real CUDA driver from: {}", path);
39                return DlHandle(handle);
40            }
41        }
42
43        panic!("Failed to find/load libcuda.so. Ensure it is in LD_LIBRARY_PATH.");
44    });
45    handle_wrapper.0
46}
47
48fn get_libcudart() -> *mut c_void {
49    let handle_wrapper = CUDART_LIB.get_or_init(|| unsafe {
50        let mut paths = vec![
51            "/usr/local/cuda/targets/x86_64-linux/lib/libcudart.so".to_string(),
52            "/usr/lib/x86_64-linux-gnu/libcudart.so".to_string(),
53            "/usr/lib64/libcudart.so".to_string(),
54        ];
55
56        // If CUDA_PATH or CUDA_HOME are set, prioritize those
57        if let Some(cuda_home) = env::var_os("CUDA_HOME") {
58            let path = format!("{}/targets/x86_64-linux/lib/libcudart.so", cuda_home.to_string_lossy());
59            paths.insert(0, path);
60        }
61
62        for path in paths.iter() {
63            let s = CString::new(path.clone()).unwrap();
64            let flags = libc::RTLD_NOW | libc::RTLD_LOCAL | libc::RTLD_NODELETE;
65            let handle = libc::dlopen(s.as_ptr(), flags);
66            if !handle.is_null() {
67                debug!("Loaded real CUDA runtime from: {}", path);
68                return DlHandle(handle);
69            }
70        }
71
72        panic!("Failed to find/load libcudart.so. Ensure CUDA Toolkit is installed and in LD_LIBRARY_PATH.");
73    });
74    handle_wrapper.0
75}
76
77pub fn dlsym_next(symbol: &[u8]) -> *mut c_void {
78    let sym_str = std::str::from_utf8(symbol).unwrap_or("");
79
80    // Route to the correct library based on prefix
81    let handle = if sym_str.starts_with("cuda") || sym_str.starts_with("__cuda") {
82        get_libcudart()
83    } else {
84        get_libcuda()
85    };
86
87    let ptr = unsafe { libc::dlsym(handle, symbol.as_ptr() as *const _) };
88    if ptr.is_null() {
89        warn!("dlsym_next fail for symbol: {sym_str}");
90    }
91    ptr
92}
93
94// ─── Macros ──────────────────────────────────────────────────────────────────
95
96/// Installs the `cuGetProcAddress` hooks required for the interposer to function.
97/// This macro automatically includes the `hook_map.rs` generated by `cuda-interposer-build`.
98#[macro_export]
99macro_rules! install_hooks {
100    () => {
101        use std::{
102            ffi::CStr,
103            os::raw::{c_char, c_int, c_void},
104        };
105        use tracing::debug;
106
107        #[inline(always)]
108        fn get_local_hook(name: &str) -> Option<*mut $crate::libc::c_void> {
109            // Include returns the closure expression from hook_map.rs
110            let hook_fn = include!(concat!(env!("OUT_DIR"), "/hook_map.rs"));
111            hook_fn(name)
112        }
113
114        type CUresult = u32; // enum
115        type CUdriverProcAddressQueryResult = u32; // enum
116
117        $crate::cuda_hook! {
118            pub unsafe extern "C" fn cuGetProcAddress_v2(
119                symbol: *const $crate::libc::c_char,
120                pfn: *mut *mut $crate::libc::c_void,
121                cuda_version: $crate::libc::c_int,
122                flags: u64,
123                symbol_status: *mut CUdriverProcAddressQueryResult
124            ) -> CUresult {
125                let sym_name_c = unsafe { ::std::ffi::CStr::from_ptr(symbol) };
126                let sym_name = sym_name_c.to_string_lossy();
127
128                // A. Call real implementation
129                let real_fn = *__real_cuGetProcAddress_v2;
130                let ret = unsafe { real_fn(symbol, pfn, cuda_version, flags, symbol_status) };
131
132                if let Some(our_ptr) = get_local_hook(&sym_name) {
133                    $crate::tracing::debug!("Hooking symbol via cuGetProcAddress_v2: {}", sym_name);
134                    unsafe { *pfn = our_ptr };
135                    return 0; // CUDA_SUCCESS
136                }
137
138                ret
139            }
140        }
141
142        $crate::cuda_hook! {
143            pub unsafe extern "C" fn cuGetProcAddress(
144                symbol: *const $crate::libc::c_char,
145                pfn: *mut *mut $crate::libc::c_void,
146                cuda_version: $crate::libc::c_int,
147                flags: u64,
148                symbol_status: *mut CUdriverProcAddressQueryResult
149            ) -> CUresult {
150               cuGetProcAddress_v2(symbol, pfn, cuda_version, flags, symbol_status)
151            }
152        }
153    };
154}
155
156#[macro_export]
157macro_rules! cuda_hook {
158    (
159        pub unsafe extern "C" fn $fname:ident( $($arg:ident : $arg_ty:ty),* $(,)? )
160        -> $ret:ty
161        $body:block
162    ) => {
163        $crate::paste::paste! {
164            #[allow(non_upper_case_globals)]
165            pub static [<__real_ $fname>]: $crate::once_cell::sync::Lazy<
166                unsafe extern "C" fn($($arg_ty),*) -> $ret
167            > = $crate::once_cell::sync::Lazy::new(|| {
168                let name = concat!(stringify!($fname), "\0");
169                let sym = $crate::dlsym_next(name.as_bytes());
170                if sym.is_null() {
171                    panic!("Missing symbol: {}", stringify!($fname));
172                }
173                unsafe { std::mem::transmute(sym) }
174            });
175
176            #[unsafe(no_mangle)]
177            pub unsafe extern "C" fn $fname( $($arg : $arg_ty),* ) -> $ret {
178                 $body
179            }
180        }
181    };
182}
183
184#[macro_export]
185macro_rules! generate_proxy {
186    // Internal: Generate specific alias function
187    (
188        @generate_alias
189        alias: $alias:ident,
190        target_fn: $fname:ident,
191        args: ( [ $( ($arg:ident : $arg_ty:ty) ),* ] ),
192        ret: $ret:ty
193    ) => {
194        $crate::paste::paste! {
195            #[unsafe(no_mangle)]
196            pub unsafe extern "C" fn $alias( $( $arg : $arg_ty ),* ) -> $ret {
197                let f = *[<__REAL_ $fname:upper>];
198                f( $( $arg ),* )
199            }
200        }
201    };
202
203    // Internal: Recurse over aliases
204    (
205        @recurse_aliases
206        target_fn: $fname:ident,
207        ret: $ret:ty,
208        args_tt: $args_tt:tt,
209        aliases: [ $($alias:ident),* ]
210    ) => {
211        $(
212            $crate::generate_proxy!(
213                @generate_alias
214                alias: $alias,
215                target_fn: $fname,
216                args: $args_tt,
217                ret: $ret
218            );
219        )*
220    };
221
222    // Internal: Generate Main Function and Lazy static
223    (
224        @generate_main
225        fn $fname:ident ( [ $( ($arg:ident : $arg_ty:ty) ),* ] ) -> $ret:ty;
226        target_symbol: $real_sym:ident
227    ) => {
228        $crate::paste::paste! {
229            static [<__REAL_ $fname:upper>]: $crate::once_cell::sync::Lazy<
230                  unsafe extern "C" fn( $($arg_ty),* ) -> $ret
231            > = $crate::once_cell::sync::Lazy::new(|| {
232                let name = concat!(stringify!($real_sym), "\0");
233                let ptr = $crate::dlsym_next(name.as_bytes());
234                if ptr.is_null() {
235                    eprintln!("fatal: symbol '{}' not found in underlying library", name);
236                    std::process::abort();
237                }
238                unsafe { std::mem::transmute(ptr) }
239            });
240
241            #[unsafe(no_mangle)]
242            pub unsafe extern "C" fn $fname( $( $arg : $arg_ty ),* ) -> $ret {
243                let f = *[<__REAL_ $fname:upper>];
244                f( $( $arg ),* )
245            }
246        }
247    };
248
249    // Entry Point
250    (
251        fn $fname:ident $args_tt:tt -> $ret:ty;
252        name: $real_sym:ident
253        $(, aliases: $($alias:ident),* )?
254    ) => {
255        $crate::generate_proxy!(
256            @generate_main
257            fn $fname $args_tt -> $ret;
258            target_symbol: $real_sym
259        );
260        $(
261            $crate::generate_proxy!(
262                @recurse_aliases
263                target_fn: $fname,
264                ret: $ret,
265                args_tt: $args_tt,
266                aliases: [ $($alias),* ]
267            );
268        )?
269    };
270}
cuda_interposer/lib.rs

cuda_interposer/
lib.rs
