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:
194
crates/xtrain-tensor/src/tensor.rs
Normal file
194
crates/xtrain-tensor/src/tensor.rs
Normal file
@@ -0,0 +1,194 @@
|
||||
//! The `Tensor` type: shape/strides/dtype over reference-counted [`Storage`],
|
||||
//! with host↔device transfer and one elementwise op (`scale`) wired end-to-end
|
||||
//! through a CUDA kernel.
|
||||
|
||||
use crate::dtype::{DType, TensorDType};
|
||||
use crate::shape::{self, Dims};
|
||||
use crate::storage::{Device, Storage};
|
||||
|
||||
/// Multi-dimensional array backed by CPU or GPU storage.
|
||||
///
|
||||
/// Strides are in elements (row-major). T2 tensors created here are always
|
||||
/// contiguous; the `strides`/`offset` fields exist so later phases can add
|
||||
/// zero-copy views without changing this type's shape.
|
||||
#[derive(Clone)]
|
||||
pub struct Tensor {
|
||||
storage: Storage,
|
||||
shape: Dims,
|
||||
strides: Dims,
|
||||
offset: usize,
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
impl Tensor {
|
||||
// --- Creation ---
|
||||
|
||||
/// Build a contiguous CPU tensor from a typed host slice.
|
||||
pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
|
||||
let numel = shape::num_elements(shape);
|
||||
assert_eq!(
|
||||
data.len(),
|
||||
numel,
|
||||
"data length {} != shape numel {numel}",
|
||||
data.len()
|
||||
);
|
||||
let bytes = unsafe {
|
||||
std::slice::from_raw_parts(data.as_ptr() as *const u8, numel * T::DTYPE.size_bytes())
|
||||
};
|
||||
Self {
|
||||
storage: Storage::cpu(bytes.to_vec()),
|
||||
shape: Dims::from_slice(shape),
|
||||
strides: shape::contiguous_strides(shape),
|
||||
offset: 0,
|
||||
dtype: T::DTYPE,
|
||||
}
|
||||
}
|
||||
|
||||
/// Zero-filled contiguous tensor on the given device.
|
||||
pub fn zeros(shape: &[usize], dtype: DType, device: Device) -> Self {
|
||||
let len_bytes = shape::num_elements(shape) * dtype.size_bytes();
|
||||
let storage = Storage::zeros(len_bytes, device).expect("zeros alloc failed");
|
||||
Self {
|
||||
storage,
|
||||
shape: Dims::from_slice(shape),
|
||||
strides: shape::contiguous_strides(shape),
|
||||
offset: 0,
|
||||
dtype,
|
||||
}
|
||||
}
|
||||
|
||||
// --- Properties ---
|
||||
|
||||
pub fn shape(&self) -> &[usize] {
|
||||
&self.shape
|
||||
}
|
||||
pub fn strides(&self) -> &[usize] {
|
||||
&self.strides
|
||||
}
|
||||
pub fn dtype(&self) -> DType {
|
||||
self.dtype
|
||||
}
|
||||
pub fn ndim(&self) -> usize {
|
||||
self.shape.len()
|
||||
}
|
||||
pub fn numel(&self) -> usize {
|
||||
shape::num_elements(&self.shape)
|
||||
}
|
||||
pub fn offset(&self) -> usize {
|
||||
self.offset
|
||||
}
|
||||
pub fn device(&self) -> Device {
|
||||
self.storage.device()
|
||||
}
|
||||
pub fn is_contiguous(&self) -> bool {
|
||||
shape::is_contiguous(&self.shape, &self.strides)
|
||||
}
|
||||
pub fn storage(&self) -> &Storage {
|
||||
&self.storage
|
||||
}
|
||||
|
||||
// --- Device transfer ---
|
||||
|
||||
/// Move (copy) the tensor to `device`. Returns a cheap clone if already there.
|
||||
pub fn to_device(&self, device: Device) -> Self {
|
||||
if self.device() == device {
|
||||
return self.clone();
|
||||
}
|
||||
let storage = self
|
||||
.storage
|
||||
.to_device(device)
|
||||
.expect("device transfer failed");
|
||||
Self {
|
||||
storage,
|
||||
shape: self.shape.clone(),
|
||||
strides: self.strides.clone(),
|
||||
offset: self.offset,
|
||||
dtype: self.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
// --- Host data access (CPU only) ---
|
||||
|
||||
/// Typed read-only view of the data. Requires a contiguous CPU tensor.
|
||||
pub fn as_slice<T: TensorDType>(&self) -> &[T] {
|
||||
assert_eq!(T::DTYPE, self.dtype, "dtype mismatch");
|
||||
assert_eq!(self.device(), Device::Cpu, "as_slice requires CPU tensor");
|
||||
assert!(self.is_contiguous(), "as_slice requires contiguous tensor");
|
||||
let bytes = self.storage.as_cpu_bytes();
|
||||
let start = self.offset * self.dtype.size_bytes();
|
||||
unsafe { std::slice::from_raw_parts(bytes[start..].as_ptr() as *const T, self.numel()) }
|
||||
}
|
||||
|
||||
/// Raw element pointer at the tensor's offset (for kernel launches).
|
||||
pub fn data_ptr(&self) -> *const u8 {
|
||||
let byte_off = self.offset * self.dtype.size_bytes();
|
||||
match self.device() {
|
||||
Device::Cpu => unsafe { self.storage.as_cpu_bytes().as_ptr().add(byte_off) },
|
||||
Device::Cuda(_) => unsafe { self.storage.gpu_buffer().as_ptr().add(byte_off) },
|
||||
}
|
||||
}
|
||||
|
||||
// --- Elementwise op (the T2 end-to-end kernel) ---
|
||||
|
||||
/// Out-of-place elementwise scale: returns a new tensor `out[i] = self[i] * alpha`.
|
||||
///
|
||||
/// Runs the `scale_f32` CUDA kernel. Requires a contiguous F32 tensor on the
|
||||
/// GPU. Available only when CUDA was compiled in (`not(no_cuda)`).
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn scale(&self, alpha: f32) -> Self {
|
||||
assert_eq!(self.dtype, DType::F32, "scale only supports F32 in T2");
|
||||
assert!(self.is_contiguous(), "scale requires contiguous tensor");
|
||||
assert!(
|
||||
matches!(self.device(), Device::Cuda(_)),
|
||||
"scale requires a CUDA tensor"
|
||||
);
|
||||
|
||||
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_scale_f32(
|
||||
self.data_ptr() as *const f32,
|
||||
out.data_ptr() as *mut f32,
|
||||
alpha,
|
||||
self.numel() as i32,
|
||||
std::ptr::null_mut(), // default stream
|
||||
);
|
||||
}
|
||||
xtrain_cuda::device::synchronize().expect("scale kernel sync failed");
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Tensor {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
|
||||
self.shape.as_slice(),
|
||||
self.dtype,
|
||||
self.device(),
|
||||
self.is_contiguous()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn from_slice_shape_and_data() {
|
||||
let t = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
|
||||
assert_eq!(t.shape(), &[2, 3]);
|
||||
assert_eq!(t.strides(), &[3, 1]);
|
||||
assert_eq!(t.numel(), 6);
|
||||
assert_eq!(t.device(), Device::Cpu);
|
||||
assert!(t.is_contiguous());
|
||||
assert_eq!(t.as_slice::<f32>(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zeros_cpu() {
|
||||
let t = Tensor::zeros(&[4], DType::F32, Device::Cpu);
|
||||
assert_eq!(t.as_slice::<f32>(), &[0.0, 0.0, 0.0, 0.0]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user