1use cudarc::driver::sys;
4use cudarc::driver::DeviceRepr;
5use fs2::FileExt;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex, OnceLock};
8use std::{fs::OpenOptions, path::Path};
9use xlog_core::{resolve_bool, MemoryBudget, Result, XlogError};
10use xlog_cuda::device_runtime::{LogRecord, LoggingSink, SinkError, XlogDeviceRuntime};
11use xlog_cuda::memory::{DeviceRead, DeviceWrite};
12use xlog_cuda::{
13 CudaBuffer, CudaDevice, CudaKernelProvider, CudaProviderBuilder, GpuMemoryManager,
14};
15
16#[derive(Default)]
17struct TransferCounters {
18 htod_bytes: AtomicU64,
19 dtoh_bytes: AtomicU64,
20}
21
22impl TransferCounters {
23 fn reset(&self) {
24 self.htod_bytes.store(0, Ordering::SeqCst);
25 self.dtoh_bytes.store(0, Ordering::SeqCst);
26 }
27
28 fn add_htod(&self, bytes: u64) {
29 self.htod_bytes.fetch_add(bytes, Ordering::SeqCst);
30 }
31
32 fn add_dtoh(&self, bytes: u64) {
33 self.dtoh_bytes.fetch_add(bytes, Ordering::SeqCst);
34 }
35
36 fn snapshot(&self) -> (u64, u64) {
37 (
38 self.htod_bytes.load(Ordering::SeqCst),
39 self.dtoh_bytes.load(Ordering::SeqCst),
40 )
41 }
42}
43
44const GPU_TEST_LOCK_PATH: &str = "/tmp/xlog-gpu-tests.lock";
45
46fn gpu_test_lock() -> std::sync::MutexGuard<'static, ()> {
47 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
48 LOCK.get_or_init(|| Mutex::new(()))
49 .lock()
50 .unwrap_or_else(|e| e.into_inner())
51}
52
53fn gpu_test_lock_file() -> Result<std::fs::File> {
54 let path = Path::new(GPU_TEST_LOCK_PATH);
55 let file = OpenOptions::new()
56 .create(true)
57 .truncate(false)
58 .read(true)
59 .write(true)
60 .open(path)
61 .map_err(|e| XlogError::Kernel(format!("Failed to open GPU test lock file: {}", e)))?;
62 file.lock_exclusive()
63 .map_err(|e| XlogError::Kernel(format!("Failed to acquire GPU test lock file: {}", e)))?;
64 Ok(file)
65}
66
67struct DiscardSink;
73
74impl LoggingSink for DiscardSink {
75 fn emit(&self, _record: LogRecord) -> std::result::Result<(), SinkError> {
76 Ok(())
77 }
78}
79
80pub struct TestContext {
82 _lock: std::sync::MutexGuard<'static, ()>,
86 _file_lock: std::fs::File,
87 pub provider: Arc<CudaKernelProvider>,
88 pub device: Arc<CudaDevice>,
89 pub memory: Arc<GpuMemoryManager>,
90 _runtime: Arc<XlogDeviceRuntime>,
92 transfer: Arc<TransferCounters>,
93}
94
95pub fn enforce_cuda_required(context: &str, err: &XlogError) {
101 if resolve_bool(None, "XLOG_REQUIRE_CUDA", false)
102 .unwrap_or_else(|error| panic!("invalid CUDA test configuration: {error}"))
103 {
104 panic!("XLOG_REQUIRE_CUDA=1 but CUDA is unavailable ({context}): {err}");
105 }
106}
107
108impl TestContext {
109 pub fn with_budget(budget_bytes: usize) -> Result<Self> {
115 let result = Self::with_budget_inner(budget_bytes);
116 if let Err(err) = &result {
117 enforce_cuda_required("TestContext::with_budget", err);
118 }
119 result
120 }
121
122 fn with_budget_inner(budget_bytes: usize) -> Result<Self> {
123 let lock = gpu_test_lock();
124 let file_lock = gpu_test_lock_file()?;
125
126 let device_count = std::panic::catch_unwind(CudaDevice::count)
128 .map_err(|_| {
129 XlogError::Kernel(
130 "Failed to get device count: cudarc panicked during driver initialization"
131 .to_string(),
132 )
133 })?
134 .map_err(|e| XlogError::Kernel(format!("Failed to get device count: {}", e)))?;
135
136 if device_count == 0 {
137 return Err(XlogError::Kernel("No CUDA devices available".to_string()));
138 }
139
140 let budget = MemoryBudget::with_limit(budget_bytes as u64);
141 let transfer = Arc::new(TransferCounters::default());
142 let provider = Arc::new(
143 CudaProviderBuilder::new(0, budget)
144 .with_logging_sink(Arc::new(DiscardSink) as Arc<dyn LoggingSink>)
145 .build()?,
146 );
147 let device = Arc::clone(provider.device());
148 let memory = Arc::clone(provider.memory());
149 let runtime = Arc::clone(memory.runtime().ok_or_else(|| {
150 XlogError::Kernel("provider builder did not attach a runtime".to_string())
151 })?);
152
153 Ok(Self {
154 _lock: lock,
155 _file_lock: file_lock,
156 provider,
157 device,
158 memory,
159 _runtime: runtime,
160 transfer,
161 })
162 }
163
164 pub fn reap_pending(&self) {
174 let _ = self._runtime.reap_pending();
178 }
179
180 pub fn new() -> Result<Self> {
182 Self::with_budget(1024 * 1024 * 1024)
183 }
184
185 pub fn large() -> Result<Self> {
187 Self::with_budget(4 * 1024 * 1024 * 1024)
188 }
189
190 pub fn sync_and_check(&self) -> Result<()> {
192 self.device
193 .inner()
194 .synchronize()
195 .map_err(|e| XlogError::Kernel(format!("Sync failed: {}", e)))?;
196 Ok(())
197 }
198
199 pub fn memory_used(&self) -> usize {
201 self.memory.allocated_bytes() as usize
202 }
203
204 pub fn memory_budget(&self) -> usize {
206 self.memory.budget().device_bytes as usize
207 }
208
209 pub fn reset_transfer_counters(&self) {
211 self.transfer.reset();
212 }
213
214 pub fn transfer_counters(&self) -> (u64, u64) {
216 self.transfer.snapshot()
217 }
218
219 pub fn htod_sync_copy_into<T: DeviceRepr, Dst: DeviceWrite<T>>(
221 &self,
222 src: &[T],
223 dst: &mut Dst,
224 ) -> Result<()> {
225 let bytes = std::mem::size_of::<T>()
226 .checked_mul(src.len())
227 .ok_or_else(|| XlogError::Kernel("htod byte count overflow".to_string()))?;
228 self.transfer.add_htod(bytes as u64);
229 self.device
230 .inner()
231 .htod_sync_copy_into(src, dst)
232 .map_err(|e| XlogError::Kernel(format!("Failed htod copy: {}", e)))?;
233 Ok(())
234 }
235
236 pub fn dtoh_sync_copy<T: DeviceRepr, Src: DeviceRead<T>>(&self, src: &Src) -> Result<Vec<T>> {
238 let bytes = std::mem::size_of::<T>()
239 .checked_mul(src.device_view().len())
240 .ok_or_else(|| XlogError::Kernel("dtoh byte count overflow".to_string()))?;
241 self.transfer.add_dtoh(bytes as u64);
242 self.device
243 .inner()
244 .dtoh_sync_copy(src)
245 .map_err(|e| XlogError::Kernel(format!("Failed dtoh copy: {}", e)))
246 }
247
248 pub fn dtoh_sync_copy_into<T: DeviceRepr, Src: DeviceRead<T>>(
250 &self,
251 src: &Src,
252 dst: &mut [T],
253 ) -> Result<()> {
254 let bytes = std::mem::size_of::<T>()
255 .checked_mul(dst.len())
256 .ok_or_else(|| XlogError::Kernel("dtoh byte count overflow".to_string()))?;
257 self.transfer.add_dtoh(bytes as u64);
258 self.device
259 .inner()
260 .dtoh_sync_copy_into(src, dst)
261 .map_err(|e| XlogError::Kernel(format!("Failed dtoh copy: {}", e)))?;
262 Ok(())
263 }
264
265 pub fn measure_gpu_ms<F>(&self, f: F) -> Result<f32>
267 where
268 F: FnOnce() -> Result<()>,
269 {
270 let stream = self.device.inner().stream();
271 let start = stream
272 .context()
273 .new_event(Some(sys::CUevent_flags::CU_EVENT_DEFAULT))
274 .map_err(|e| XlogError::Kernel(format!("Failed to create start event: {}", e)))?;
275 let end = stream
276 .context()
277 .new_event(Some(sys::CUevent_flags::CU_EVENT_DEFAULT))
278 .map_err(|e| XlogError::Kernel(format!("Failed to create end event: {}", e)))?;
279
280 start
281 .record(stream)
282 .map_err(|e| XlogError::Kernel(format!("Failed to record start event: {}", e)))?;
283 f()?;
284 end.record(stream)
285 .map_err(|e| XlogError::Kernel(format!("Failed to record end event: {}", e)))?;
286 self.sync_and_check()?;
287 let elapsed = start.elapsed_ms(&end).map_err(|e| {
288 XlogError::Kernel(format!("Failed to measure event elapsed time: {}", e))
289 })?;
290 Ok(elapsed)
291 }
292
293 pub fn multi_gpu_available(&self) -> bool {
295 match std::panic::catch_unwind(CudaDevice::count) {
296 Ok(Ok(n)) => n > 1,
297 _ => false,
298 }
299 }
300
301 pub fn compute_capability(&self) -> Result<(u32, u32)> {
303 use cudarc::driver::sys::CUdevice_attribute;
304
305 let major = self
306 .device
307 .inner()
308 .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
309 .map_err(|e| {
310 XlogError::Kernel(format!("Failed to query compute capability major: {}", e))
311 })?;
312 let minor = self
313 .device
314 .inner()
315 .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)
316 .map_err(|e| {
317 XlogError::Kernel(format!("Failed to query compute capability minor: {}", e))
318 })?;
319
320 let major_u32: u32 = major.try_into().map_err(|_| {
321 XlogError::Kernel(format!(
322 "Failed to convert compute capability major {} to u32",
323 major
324 ))
325 })?;
326 let minor_u32: u32 = minor.try_into().map_err(|_| {
327 XlogError::Kernel(format!(
328 "Failed to convert compute capability minor {} to u32",
329 minor
330 ))
331 })?;
332
333 Ok((major_u32, minor_u32))
334 }
335
336 pub fn device_row_count(&self, buffer: &CudaBuffer) -> u64 {
340 let mut host_rows = [0u32];
341 self.dtoh_sync_copy_into(buffer.num_rows_device(), &mut host_rows)
342 .unwrap_or_else(|e| panic!("Failed to read device row count: {}", e));
343 host_rows[0] as u64
344 }
345}
346
347#[macro_export]
349macro_rules! gpu_test {
350 ($name:ident, $body:expr) => {
351 #[test]
352 fn $name() {
353 let ctx =
354 $crate::harness::TestContext::new().expect("CUDA device required for this test");
355 $body(&ctx);
356 }
357 };
358}
359
360#[macro_export]
362macro_rules! gpu_test_skip {
363 ($name:ident, $body:expr) => {
364 #[test]
365 fn $name() {
366 let ctx = match $crate::harness::TestContext::new() {
367 Ok(ctx) => ctx,
368 Err(_) => {
369 eprintln!("Skipping {}: no CUDA device", stringify!($name));
370 return;
371 }
372 };
373 $body(&ctx);
374 }
375 };
376}
377
378#[cfg(test)]
379mod tests {
380 use super::*;
381 use std::process::Command;
382
383 #[test]
384 fn test_gpu_test_lock_file_is_exclusive() {
385 let ctx = match TestContext::new() {
386 Ok(ctx) => ctx,
387 Err(_) => {
388 eprintln!("Skipping test: no CUDA device available");
389 return;
390 }
391 };
392
393 let status = Command::new("flock")
394 .arg("-n")
395 .arg(GPU_TEST_LOCK_PATH)
396 .arg("-c")
397 .arg("true")
398 .status();
399
400 match status {
401 Ok(status) => assert!(
402 !status.success(),
403 "GPU test lock file should be held exclusively while TestContext is alive"
404 ),
405 Err(_) => {
406 eprintln!("Skipping test: flock command not available");
407 }
408 }
409
410 drop(ctx);
411 }
412}