xlog-cuda 0.9.2

CUDA kernel provider, buffers, and interop for XLOG
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

#[cfg(feature = "arrow-device-import")]
mod tests {
    use std::sync::Arc;

    use xlog_core::{MemoryBudget, ScalarType, Schema};
    use xlog_cuda::{CudaDevice, CudaKernelProvider, GpuMemoryManager};

    fn setup_provider() -> Option<Arc<CudaKernelProvider>> {
        let device = match CudaDevice::new(0) {
            Ok(d) => Arc::new(d),
            Err(e) => {
                eprintln!("Skipping: CUDA runtime unavailable: {}", e);
                return None;
            }
        };
        let memory = Arc::new(GpuMemoryManager::new(
            device.clone(),
            MemoryBudget::with_limit(512 * 1024 * 1024),
        ));
        CudaKernelProvider::new(device, memory).ok().map(Arc::new)
    }

    #[test]
    fn arrow_device_roundtrip_import_export() {
        let Some(provider) = setup_provider() else {
            eprintln!("Skipping: no CUDA device");
            return;
        };

        let schema = Schema::new(vec![("x".to_string(), ScalarType::U32)]);
        let buffer = provider
            .create_buffer_from_slice::<u32>(&[1, 2, 3], schema.clone())
            .expect("create buffer");

        let arrow_dev = provider
            .to_arrow_device_record_batch(buffer)
            .expect("export device record batch");

        let imported = provider
            .from_arrow_device_record_batch(arrow_dev)
            .expect("import device record batch");

        assert_eq!(imported.schema(), &schema);
        assert_eq!(imported.num_rows(), 3);

        let batch = provider
            .to_arrow_record_batch(&imported)
            .expect("to arrow record batch");
        let array = batch
            .column(0)
            .as_any()
            .downcast_ref::<arrow::array::UInt32Array>()
            .expect("u32 array");
        assert_eq!(array.len(), 3);
        assert_eq!(array.value(0), 1);
        assert_eq!(array.value(1), 2);
        assert_eq!(array.value(2), 3);
    }
}

#[cfg(not(feature = "arrow-device-import"))]
#[test]
fn arrow_device_import_feature_disabled() {
    eprintln!("arrow-device-import feature disabled; skipping");
}