tensor: minimal Tensor crate over xtrain-cuda
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>
This commit is contained in:
58
crates/xtrain-tensor/tests/integration.rs
Normal file
58
crates/xtrain-tensor/tests/integration.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
// 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()
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user