xtrain/crates/xtrain-tensor/src/tensor.rs

//! 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
    }

    // --- GEMM (the T3 kernels) ---

    /// Matrix multiply: `C = self @ other`. `self`:[M,K], `other`:[K,N] → [M,N].
    ///
    /// Runs the tiled `gemm_tiled_f32` CUDA kernel. Requires contiguous F32
    /// tensors on the same GPU. Available only when CUDA is compiled in.
    #[cfg(not(no_cuda))]
    pub fn matmul(&self, other: &Tensor) -> Self {
        assert_eq!(self.dtype, DType::F32, "matmul only supports F32");
        assert_eq!(other.dtype, DType::F32, "matmul only supports F32");
        assert_eq!(self.ndim(), 2, "matmul requires 2D lhs");
        assert_eq!(other.ndim(), 2, "matmul requires 2D rhs");
        assert_eq!(
            self.shape[1], other.shape[0],
            "inner dimension mismatch: {:?} @ {:?}",
            self.shape, other.shape
        );
        assert!(
            self.is_contiguous() && other.is_contiguous(),
            "matmul requires contiguous tensors"
        );
        assert_eq!(self.device(), other.device(), "matmul device mismatch");
        assert!(
            matches!(self.device(), Device::Cuda(_)),
            "matmul requires CUDA tensors"
        );

        let m = self.shape[0];
        let k = self.shape[1];
        let n = other.shape[1];
        let out = Tensor::zeros(&[m, n], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_gemm_tiled_f32(
                self.data_ptr() as *const f32,
                other.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                m as i32,
                n as i32,
                k as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("matmul kernel sync failed");
        out
    }

    /// Out-of-place 2D transpose: returns a new contiguous tensor `out[j,i] =
    /// self[i,j]`. Requires a contiguous F32 CUDA tensor.
    #[cfg(not(no_cuda))]
    pub fn transpose_2d(&self) -> Self {
        assert_eq!(self.dtype, DType::F32, "transpose only supports F32");
        assert_eq!(self.ndim(), 2, "transpose_2d requires 2D tensor");
        assert!(self.is_contiguous(), "transpose requires contiguous tensor");
        assert!(
            matches!(self.device(), Device::Cuda(_)),
            "transpose requires a CUDA tensor"
        );

        let rows = self.shape[0];
        let cols = self.shape[1];
        let out = Tensor::zeros(&[cols, rows], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_transpose_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("transpose kernel sync failed");
        out
    }

    /// Backward of `C = A @ B` given the upstream gradient `dC` (shape [M,N]).
    /// Returns `(dA, dB)` where `dA = dC @ Bᵀ` ([M,K]) and `dB = Aᵀ @ dC`
    /// ([K,N]). All tensors contiguous F32 on the same GPU.
    #[cfg(not(no_cuda))]
    pub fn matmul_backward(a: &Tensor, b: &Tensor, dc: &Tensor) -> (Tensor, Tensor) {
        assert_eq!(a.ndim(), 2, "matmul_backward requires 2D A");
        assert_eq!(b.ndim(), 2, "matmul_backward requires 2D B");
        assert_eq!(dc.ndim(), 2, "matmul_backward requires 2D dC");
        assert_eq!(a.shape[1], b.shape[0], "A/B inner dim mismatch");
        assert_eq!(dc.shape[0], a.shape[0], "dC rows != A rows (M)");
        assert_eq!(dc.shape[1], b.shape[1], "dC cols != B cols (N)");

        let da = dc.matmul(&b.transpose_2d()); // [M,N] @ [N,K] = [M,K]
        let db = a.transpose_2d().matmul(dc); // [K,M] @ [M,N] = [K,N]
        (da, db)
    }

    // --- Transformer / autograd op primitives (the T4 kernels) ---
    //
    // Each is a thin, contiguous-F32-on-GPU wrapper over a kernel in
    // csrc/ops/nn.cu. The autograd `Var` layer (xtrain-autodiff) builds nodes on
    // top of these; the analytic backwards are derived in docs/03-autograd-engine.md.

    /// Elementwise `out = self + other` (same shape).
    #[cfg(not(no_cuda))]
    pub fn add(&self, other: &Tensor) -> Self {
        self.check_binary(other, "add");
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_add_f32(
                self.data_ptr() as *const f32,
                other.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                self.numel() as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("add sync failed");
        out
    }

    /// Elementwise `out = self * other` (same shape, Hadamard product).
    #[cfg(not(no_cuda))]
    pub fn mul(&self, other: &Tensor) -> Self {
        self.check_binary(other, "mul");
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_mul_f32(
                self.data_ptr() as *const f32,
                other.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                self.numel() as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("mul sync failed");
        out
    }

    /// Broadcast bias add: `out[r,c] = self[r,c] + bias[c]`.
    /// `self`:[rows,cols], `bias`:[cols].
    #[cfg(not(no_cuda))]
    pub fn add_bias(&self, bias: &Tensor) -> Self {
        assert_eq!(self.ndim(), 2, "add_bias requires 2D input");
        assert_eq!(bias.ndim(), 1, "bias must be 1D");
        assert_eq!(self.shape[1], bias.shape[0], "bias len != cols");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_add_bias_f32(
                self.data_ptr() as *const f32,
                bias.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("add_bias sync failed");
        out
    }

    /// Column-sum over rows: `out[c] = sum_r self[r,c]`. This is the bias
    /// backward (sum the upstream grad over the broadcast dim). `self`:[rows,cols]
    /// → [cols].
    #[cfg(not(no_cuda))]
    pub fn sum_rows(&self) -> Self {
        assert_eq!(self.ndim(), 2, "sum_rows requires 2D input");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let out = Tensor::zeros(&[cols], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_sum_rows_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("sum_rows sync failed");
        out
    }

    /// RMSNorm forward: `y[r,c] = x[r,c] * inv_rms[r] * gamma[c]` with
    /// `inv_rms = rsqrt(mean(x²) + eps)`. `self`:[rows,cols], `gamma`:[cols].
    /// Returns `(y, inv_rms)`; `inv_rms`:[rows] is cached for backward.
    #[cfg(not(no_cuda))]
    pub fn rms_norm(&self, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) {
        assert_eq!(self.ndim(), 2, "rms_norm requires 2D input");
        assert_eq!(gamma.ndim(), 1, "gamma must be 1D");
        assert_eq!(self.shape[1], gamma.shape[0], "gamma len != cols");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let y = Tensor::zeros(&self.shape, DType::F32, self.device());
        let inv_rms = Tensor::zeros(&[rows], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_rms_norm_f32(
                self.data_ptr() as *const f32,
                gamma.data_ptr() as *const f32,
                y.data_ptr() as *mut f32,
                inv_rms.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                eps,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("rms_norm sync failed");
        (y, inv_rms)
    }

    /// RMSNorm backward. Inputs are the forward `x`, `gamma`, upstream `dy`, and
    /// the cached `inv_rms`. Returns `(dx, dgamma)`.
    #[cfg(not(no_cuda))]
    pub fn rms_norm_backward(
        x: &Tensor,
        gamma: &Tensor,
        dy: &Tensor,
        inv_rms: &Tensor,
    ) -> (Tensor, Tensor) {
        let (rows, cols) = (x.shape[0], x.shape[1]);
        let dx = Tensor::zeros(&[rows, cols], DType::F32, x.device());
        let dgamma = Tensor::zeros(&[cols], DType::F32, x.device());
        unsafe {
            xtrain_cuda::ffi::launch_rms_norm_dx_f32(
                x.data_ptr() as *const f32,
                gamma.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                inv_rms.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
            xtrain_cuda::ffi::launch_rms_norm_dgamma_f32(
                x.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                inv_rms.data_ptr() as *const f32,
                dgamma.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("rms_norm_backward sync failed");
        (dx, dgamma)
    }

    /// SiLU forward: `y = x * sigmoid(x)`, elementwise.
    #[cfg(not(no_cuda))]
    pub fn silu(&self) -> Self {
        assert_eq!(self.dtype, DType::F32, "silu only supports F32");
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_silu_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                self.numel() as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("silu sync failed");
        out
    }

    /// SiLU backward: `dx = dy * (sig + x*sig*(1-sig))`, `sig = sigmoid(x)`.
    /// Inputs are the forward `x` and upstream `dy`.
    #[cfg(not(no_cuda))]
    pub fn silu_backward(x: &Tensor, dy: &Tensor) -> Self {
        let dx = Tensor::zeros(&x.shape, DType::F32, x.device());
        unsafe {
            xtrain_cuda::ffi::launch_silu_dx_f32(
                x.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                x.numel() as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("silu_backward sync failed");
        dx
    }

    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position
    /// of each token is its row index. Returns the rotated tensor.
    #[cfg(not(no_cuda))]
    pub fn rope(&self, theta: f32) -> Self {
        assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
        let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
        assert_eq!(head_dim % 2, 0, "head_dim must be even");
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_rope_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                tokens as i32,
                heads as i32,
                head_dim as i32,
                theta,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("rope sync failed");
        out
    }

    /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
    /// orthogonal map, so it needs no cached forward values, only `theta`.
    #[cfg(not(no_cuda))]
    pub fn rope_backward(dy: &Tensor, theta: f32) -> Self {
        let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
        let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
        unsafe {
            xtrain_cuda::ffi::launch_rope_dx_f32(
                dy.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                tokens as i32,
                heads as i32,
                head_dim as i32,
                theta,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("rope_backward sync failed");
        dx
    }

    /// Row-wise safe softmax over the last dim. `self`:[rows,cols].
    #[cfg(not(no_cuda))]
    pub fn softmax(&self) -> Self {
        assert_eq!(self.ndim(), 2, "softmax requires 2D input");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_softmax_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("softmax sync failed");
        out
    }

    /// Softmax backward (Jacobian): `dx[r,c] = y[r,c]*(dy[r,c] - sum_c'(dy*y))`.
    /// Inputs are the forward output `y` and upstream `dy`.
    #[cfg(not(no_cuda))]
    pub fn softmax_backward(y: &Tensor, dy: &Tensor) -> Self {
        let (rows, cols) = (y.shape[0], y.shape[1]);
        let dx = Tensor::zeros(&y.shape, DType::F32, y.device());
        unsafe {
            xtrain_cuda::ffi::launch_softmax_dx_f32(
                y.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("softmax_backward sync failed");
        dx
    }

    /// Cross-entropy forward over logits `self`:[rows,cols] with one I32 target
    /// per row. Returns `(probs, loss)` where `probs`:[rows,cols] is the softmax
    /// (cached for backward) and `loss`:[rows] is the per-row negative log-likelihood.
    #[cfg(not(no_cuda))]
    pub fn cross_entropy(&self, target: &Tensor) -> (Tensor, Tensor) {
        assert_eq!(self.ndim(), 2, "cross_entropy requires 2D logits");
        assert_eq!(target.dtype, DType::I32, "target must be I32");
        assert_eq!(target.numel(), self.shape[0], "one target per row");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let probs = Tensor::zeros(&self.shape, DType::F32, self.device());
        let loss = Tensor::zeros(&[rows], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_cross_entropy_fwd_f32(
                self.data_ptr() as *const f32,
                target.data_ptr() as *const i32,
                probs.data_ptr() as *mut f32,
                loss.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("cross_entropy sync failed");
        (probs, loss)
    }

    /// Cross-entropy backward: `dx = scale * (probs - onehot(target))`. With
    /// `scale = upstream / rows`, this is the gradient of the mean per-row loss.
    #[cfg(not(no_cuda))]
    pub fn cross_entropy_backward(probs: &Tensor, target: &Tensor, scale: f32) -> Self {
        let (rows, cols) = (probs.shape[0], probs.shape[1]);
        let dx = Tensor::zeros(&probs.shape, DType::F32, probs.device());
        unsafe {
            xtrain_cuda::ffi::launch_cross_entropy_dx_f32(
                probs.data_ptr() as *const f32,
                target.data_ptr() as *const i32,
                dx.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                scale,
                std::ptr::null_mut(),
            );
        }
        xtrain_cuda::device::synchronize().expect("cross_entropy_backward sync failed");
        dx
    }

    // Shared validation for same-shape binary elementwise ops.
    #[cfg(not(no_cuda))]
    fn check_binary(&self, other: &Tensor, op: &str) {
        assert_eq!(self.dtype, DType::F32, "{op} only supports F32");
        assert_eq!(other.dtype, DType::F32, "{op} only supports F32");
        assert_eq!(self.shape(), other.shape(), "{op} shape mismatch");
        assert_eq!(self.device(), other.device(), "{op} device mismatch");
        assert!(
            self.is_contiguous() && other.is_contiguous(),
            "{op} requires contiguous tensors"
        );
    }
}

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]);
    }
}