Gahow Wang
fbd07a578c
New xtrain-tensor crate: DType (F32), shape/stride helpers, Arc-counted host/device Storage with CPU↔CUDA copy, and a contiguous Tensor with creation, host↔device transfer, and a scale() op driving the elementwise kernel. GPU integration tests (host↔device roundtrip + scale correctness) gated behind not(no_cuda); a thin build.rs emits the no_cuda cfg so the kernel call sites compile out locally. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
59 lines
1.9 KiB
Rust
59 lines
1.9 KiB
Rust
// GPU integration tests for the tensor abstraction. Both require nvcc + a GPU,
|
|
// so they are gated behind `not(no_cuda)`. On a GPU-less machine build.rs sets
|
|
// the `no_cuda` cfg and these compile out, keeping host `cargo check` green.
|
|
#![cfg(not(no_cuda))]
|
|
|
|
use xtrain_cuda::device;
|
|
use xtrain_tensor::{Device, Tensor};
|
|
|
|
/// (a) Host → device → host roundtrip preserves the data exactly.
|
|
#[test]
|
|
fn host_device_roundtrip() {
|
|
assert!(
|
|
device::device_count().expect("device count") > 0,
|
|
"no CUDA device"
|
|
);
|
|
device::set_device(0).unwrap();
|
|
|
|
let host: Vec<f32> = (0..1024).map(|i| i as f32 * 0.5).collect();
|
|
let cpu = Tensor::from_slice(&host, &[1024]);
|
|
|
|
let gpu = cpu.to_device(Device::Cuda(0));
|
|
assert_eq!(gpu.device(), Device::Cuda(0));
|
|
assert_eq!(gpu.shape(), &[1024]);
|
|
|
|
let back = gpu.to_device(Device::Cpu);
|
|
assert_eq!(back.device(), Device::Cpu);
|
|
assert_eq!(back.as_slice::<f32>(), host.as_slice());
|
|
println!("roundtrip OK: {} elems preserved", host.len());
|
|
}
|
|
|
|
/// (b) The elementwise `scale` kernel produces correct results.
|
|
#[test]
|
|
fn elementwise_scale_kernel() {
|
|
assert!(
|
|
device::device_count().expect("device count") > 0,
|
|
"no CUDA device"
|
|
);
|
|
device::set_device(0).unwrap();
|
|
|
|
let host: Vec<f32> = (0..2048).map(|i| i as f32).collect();
|
|
let alpha = 3.0f32;
|
|
let expected: Vec<f32> = host.iter().map(|x| x * alpha).collect();
|
|
|
|
let gpu = Tensor::from_slice(&host, &[2048]).to_device(Device::Cuda(0));
|
|
let scaled = gpu.scale(alpha);
|
|
let result = scaled.to_device(Device::Cpu);
|
|
|
|
assert_eq!(result.shape(), &[2048]);
|
|
assert_eq!(result.as_slice::<f32>(), expected.as_slice());
|
|
let r = result.as_slice::<f32>();
|
|
println!(
|
|
"scale OK (alpha={alpha}): first={} mid={} last={} ({} elems)",
|
|
r[0],
|
|
r[r.len() / 2],
|
|
r[r.len() - 1],
|
|
r.len()
|
|
);
|
|
}
|