rustytorch/crates/core/rtx-backend/src/lib.rs

//! # RustyTorch++ Backend Abstraction
//!
//! Burn-inspired compile-time backend dispatch for RustyTorch++.
//!
//! This crate provides the core `Backend` trait that enables:
//! - Compile-time backend selection (no runtime overhead)
//! - Clean separation of backend-specific code
//! - Easy addition of new backends
//! - Type-safe operations
//!
//! ## Architecture
//!
//! The backend system uses associated types for compile-time polymorphism:
//!
//! ```text
//! Backend Trait
//! ├── CudaBackend      - NVIDIA GPU (hand-optimized kernels)
//! ├── MetalBackend     - Apple Silicon (Metal Performance Shaders)
//! ├── RocmBackend      - AMD GPU (HIP runtime)
//! ├── WebGpuBackend    - Browser/WASM (WebGPU shaders)
//! └── CpuBackend       - CPU fallback (BLAS/SIMD)
//! ```
//!
//! ## Example
//!
//! ```rust,ignore
//! use rtx_backend::{Backend, CudaBackend};
//!
//! // Training uses autodiff wrapper
//! type TrainingBackend = Autodiff<CudaBackend>;
//!
//! // Inference uses raw backend (zero overhead)
//! type InferenceBackend = CudaBackend;
//!
//! fn train<B: Backend>(model: &Model<B>, data: &Tensor<B, 2>) {
//!     // Backend-agnostic training code
//! }
//! ```

#![warn(missing_docs)]

pub mod auto_select;
mod device;
mod element;
pub mod error;
mod ops;
mod tensor;

pub use device::{DeviceId, DeviceOps};
pub use element::{BoolElement, BoolU8, FloatElement, IntElement};
pub use error::{BackendError, BackendResult};
pub use ops::{
    ActivationOps,
    // LLM-specific ops (keep hand-optimized)
    AttentionOps,
    KVCacheOps,
    ModuleOps,
    TensorOps,
};
pub use tensor::TensorHandle;

use std::fmt::Debug;

/// Core backend trait for RustyTorch++ - Burn-inspired compile-time dispatch.
///
/// Each backend implementation provides its own tensor primitive type and
/// operations. This enables:
/// - Zero-cost abstraction through monomorphization
/// - Backend-specific optimizations
/// - Type-safe device management
///
/// # Design Philosophy
///
/// Unlike the previous Device enum approach, the Backend trait:
/// 1. Uses associated types for compile-time dispatch (no match arms)
/// 2. Allows each backend to define its own tensor storage
/// 3. Enables composable wrappers (Autodiff, Quantized, etc.)
/// 4. Provides better type safety
///
/// # LLM Performance
///
/// For LLM-critical operations, backends delegate to hand-optimized kernels:
/// - `flash_attention` - FlashAttention-3 for CUDA, Metal MSL for Apple
/// - `ring_attention` - Distributed attention for 16M+ context
/// - `kv_cache_ops` - Entropy-guided eviction
pub trait Backend: Clone + Send + Sync + Debug + Default + 'static {
    /// The tensor primitive type for this backend.
    /// This is the actual storage/handle that holds tensor data.
    type TensorPrimitive<const D: usize>: Clone + Send + Sync + Debug;

    /// The device type for this backend.
    type Device: DeviceOps<Self>;

    /// The floating-point element type (f32, f16, bf16, fp8).
    type FloatElem: FloatElement;

    /// The integer element type (i32, i64, u32, u64).
    type IntElem: IntElement;

    /// The boolean element type.
    type BoolElem: BoolElement;

    /// Backend name for debugging and logging.
    fn name() -> &'static str;

    /// Seed the random number generator.
    fn seed(seed: u64);

    // ==================== Tensor Creation ====================

    /// Create a tensor filled with zeros.
    fn zeros<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D>;

    /// Create a tensor filled with ones.
    fn ones<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D>;

    /// Create a tensor filled with a scalar value.
    fn full<const D: usize>(
        shape: [usize; D],
        fill_value: Self::FloatElem,
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D>;

    /// Create a tensor with random uniform values in [0, 1).
    fn rand<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D>;

    /// Create a tensor with random normal values (mean=0, std=1).
    fn randn<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D>;

    /// Create a tensor from raw data.
    fn from_data<const D: usize>(
        data: &[Self::FloatElem],
        shape: [usize; D],
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D>;

    // ==================== Basic Operations ====================
    // NOTE: All operations take OWNED tensors (not references) to enable
    // ownership-based kernel fusion. Use .clone() when a tensor is needed
    // multiple times - the clone count informs the fusion system.

    /// Element-wise addition.
    fn add<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D>;

    /// Element-wise subtraction.
    fn sub<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D>;

    /// Element-wise multiplication.
    fn mul<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D>;

    /// Element-wise division.
    fn div<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D>;

    /// Matrix multiplication.
    fn matmul(
        lhs: Self::TensorPrimitive<2>,
        rhs: Self::TensorPrimitive<2>,
    ) -> Self::TensorPrimitive<2>;

    /// Batched matrix multiplication.
    fn bmm(
        lhs: Self::TensorPrimitive<3>,
        rhs: Self::TensorPrimitive<3>,
    ) -> Self::TensorPrimitive<3>;

    // ==================== Unary Operations ====================
    // NOTE: Unary ops also take owned tensors for fusion opportunities.

    /// Element-wise negation.
    fn neg<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise exponential.
    fn exp<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise natural logarithm.
    fn log<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise square root.
    fn sqrt<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise absolute value.
    fn abs<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise sine.
    fn sin<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise cosine.
    fn cos<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Element-wise power.
    fn pow<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        exp: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    /// Clamp tensor values to a range.
    fn clamp<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        min: Self::FloatElem,
        max: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    // ==================== Activation Functions ====================

    /// ReLU activation: max(0, x).
    fn relu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Sigmoid activation: 1 / (1 + exp(-x)).
    fn sigmoid<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Tanh activation.
    fn tanh<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    // ==================== Reduction Operations ====================
    // NOTE: Reductions consume input tensor for fusion with preceding ops.

    /// Sum all elements.
    fn sum<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1>;

    /// Sum along a dimension.
    fn sum_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D>;

    /// Mean of all elements.
    fn mean<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1>;

    /// Mean along a dimension.
    fn mean_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D>;

    /// Variance of all elements.
    fn var<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1>;

    /// Variance along a dimension.
    fn var_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D>;

    /// Maximum element.
    fn max<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1>;

    /// Minimum element.
    fn min<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1>;

    // ==================== Shape Operations ====================

    /// Get the shape of a tensor.
    fn shape<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> [usize; D];

    /// Reshape a tensor (no data copy if contiguous).
    fn reshape<const D1: usize, const D2: usize>(
        tensor: Self::TensorPrimitive<D1>,
        shape: [usize; D2],
    ) -> Self::TensorPrimitive<D2>;

    /// Transpose last two dimensions.
    fn transpose<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Swap two dimensions.
    fn swap_dims<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim1: usize,
        dim2: usize,
    ) -> Self::TensorPrimitive<D>;

    // ==================== Row Indexing (gather / scatter-add) ====================
    // Differentiable row indexing along dim 0. These two ops are each other's
    // adjoint, which is exactly what message passing on a graph needs:
    //   d/dx index_select(x, idx)        = index_add(grad, idx, rows(x))
    //   d/dx index_add(x, idx, num_rows) = index_select(grad, idx)
    //
    // Both have default bodies that round-trip through host memory via
    // `to_data` / `from_data`, so every backend is correct out of the box;
    // backends override them with native kernels for speed.

    /// Gather rows along dim 0: `out[i, ..] = tensor[indices[i], ..]`.
    ///
    /// Output shape is `[indices.len(), shape[1..]]`. Indices may repeat.
    ///
    /// # Panics
    /// Panics if any index is `>= shape[0]`, or if `D == 0`.
    fn index_select<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        indices: &[usize],
    ) -> Self::TensorPrimitive<D> {
        let shape = Self::shape(&tensor);
        assert!(D >= 1, "index_select requires at least one dimension");
        let num_rows = shape[0];
        let row_len: usize = shape[1..].iter().product();
        let src = Self::to_data(&tensor);
        let mut out = Vec::with_capacity(indices.len() * row_len);
        for &idx in indices {
            assert!(
                idx < num_rows,
                "index_select: index {idx} out of range for {num_rows} rows"
            );
            out.extend_from_slice(&src[idx * row_len..(idx + 1) * row_len]);
        }
        let mut out_shape = shape;
        out_shape[0] = indices.len();
        Self::from_data(&out, out_shape, &Self::device(&tensor))
    }

    /// Scatter-add rows along dim 0 into a zero tensor with `num_rows` rows:
    /// `out = zeros([num_rows, shape[1..]]); out[indices[i], ..] += tensor[i, ..]`.
    ///
    /// Indices may repeat (contributions accumulate); rows never referenced
    /// stay zero. This is the adjoint of [`Backend::index_select`].
    ///
    /// # Panics
    /// Panics if `indices.len() != shape[0]`, if any index is `>= num_rows`,
    /// or if `D == 0`.
    fn index_add<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        indices: &[usize],
        num_rows: usize,
    ) -> Self::TensorPrimitive<D> {
        let shape = Self::shape(&tensor);
        assert!(D >= 1, "index_add requires at least one dimension");
        assert_eq!(
            indices.len(),
            shape[0],
            "index_add: indices.len() must equal the number of input rows"
        );
        let row_len: usize = shape[1..].iter().product();
        let src = Self::to_data(&tensor);
        let mut out = vec![Self::FloatElem::zero(); num_rows * row_len];
        for (i, &idx) in indices.iter().enumerate() {
            assert!(
                idx < num_rows,
                "index_add: index {idx} out of range for {num_rows} rows"
            );
            let dst = &mut out[idx * row_len..(idx + 1) * row_len];
            let row = &src[i * row_len..(i + 1) * row_len];
            for (d, &s) in dst.iter_mut().zip(row) {
                *d = Self::FloatElem::from_f64(d.to_f64() + s.to_f64());
            }
        }
        let mut out_shape = shape;
        out_shape[0] = num_rows;
        Self::from_data(&out, out_shape, &Self::device(&tensor))
    }

    // ==================== LLM-Specific Operations ====================
    // These delegate to hand-optimized kernels for maximum performance.

    /// Flash Attention (optimized for each backend).
    ///
    /// Uses FlashAttention-3 on CUDA, Metal MSL kernels on Apple Silicon,
    /// or efficient fallbacks on other backends.
    /// NOTE: Q/K/V are owned to enable fusion; mask stays borrowed (read-only).
    fn flash_attention(
        query: Self::TensorPrimitive<4>, // [batch, heads, seq_len, head_dim]
        key: Self::TensorPrimitive<4>,
        value: Self::TensorPrimitive<4>,
        mask: Option<&Self::TensorPrimitive<4>>,
        scale: Self::FloatElem,
        causal: bool,
    ) -> Self::TensorPrimitive<4>;

    /// Softmax along the last dimension (numerically stable).
    fn softmax<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D>;

    /// Layer normalization.
    /// NOTE: Input tensor is owned; weight/bias stay borrowed (shared params).
    fn layer_norm<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        weight: &Self::TensorPrimitive<1>,
        bias: Option<&Self::TensorPrimitive<1>>,
        eps: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    /// RMS normalization (used in LLaMA, etc.).
    /// NOTE: Input tensor is owned; weight stays borrowed (shared param).
    fn rms_norm<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        weight: &Self::TensorPrimitive<1>,
        eps: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    /// Rotary Position Embeddings (RoPE).
    /// NOTE: Input tensor is owned; cos/sin stay borrowed (precomputed).
    fn rope<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        cos: &Self::TensorPrimitive<2>,
        sin: &Self::TensorPrimitive<2>,
    ) -> Self::TensorPrimitive<D>;

    /// GELU activation.
    fn gelu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// SiLU (Swish) activation.
    fn silu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Leaky ReLU activation: max(negative_slope * x, x).
    fn leaky_relu<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        negative_slope: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    /// ELU activation: x if x > 0, else alpha * (exp(x) - 1).
    fn elu<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        alpha: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    // ==================== Comparison Operations ====================

    /// Greater than scalar comparison.
    /// Returns a tensor with 1.0 where elements are > value, 0.0 otherwise.
    fn gt_scalar<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        value: Self::FloatElem,
    ) -> Self::TensorPrimitive<D>;

    // ==================== Convolution Operations ====================

    /// 2D Convolution.
    ///
    /// # Arguments
    /// * `input` - Input tensor `[batch, in_channels, height, width]`
    /// * `weight` - Kernel weights `[out_channels, in_channels/groups, kernel_h, kernel_w]`
    /// * `bias` - Optional bias `[out_channels]`
    /// * `stride` - Stride (height, width)
    /// * `padding` - Padding (height, width)
    /// * `dilation` - Dilation (height, width)
    /// * `groups` - Number of groups for grouped convolution
    fn conv2d(
        input: Self::TensorPrimitive<4>,
        weight: &Self::TensorPrimitive<4>,
        bias: Option<&Self::TensorPrimitive<1>>,
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        groups: usize,
    ) -> Self::TensorPrimitive<4>;

    // ==================== Pooling Operations ====================

    /// 2D max pooling.
    ///
    /// # Arguments
    /// * `input` - Input tensor `[batch, channels, height, width]`
    /// * `kernel_size` - Size of the pooling window `[kH, kW]`
    /// * `stride` - Stride of the pooling window `[sH, sW]`
    /// * `padding` - Padding added to both sides `[pH, pW]`
    ///
    /// # Returns
    /// Output tensor `[batch, channels, out_height, out_width]`
    fn max_pool2d(
        input: Self::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
    ) -> Self::TensorPrimitive<4>;

    /// 2D average pooling.
    ///
    /// # Arguments
    /// * `input` - Input tensor `[batch, channels, height, width]`
    /// * `kernel_size` - Size of the pooling window `[kH, kW]`
    /// * `stride` - Stride of the pooling window `[sH, sW]`
    /// * `padding` - Padding added to both sides `[pH, pW]`
    /// * `count_include_pad` - Whether to include padding in the averaging calculation
    ///
    /// # Returns
    /// Output tensor `[batch, channels, out_height, out_width]`
    fn avg_pool2d(
        input: Self::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
    ) -> Self::TensorPrimitive<4>;

    // ==================== Device Management ====================

    /// Get the device of a tensor.
    fn device<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Self::Device;

    /// Move tensor to a device.
    fn to_device<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D>;

    /// Copy data to CPU for inspection.
    fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem>;

    /// Synchronize the backend (wait for all operations to complete).
    fn sync(device: &Self::Device);
}

/// Marker trait for backends that support automatic differentiation.
///
/// This trait is implemented by autodiff wrappers like `Autodiff<B>`.
pub trait AutodiffBackend: Backend {
    /// The inner backend (without autodiff).
    type InnerBackend: Backend;

    /// Get a reference to the inner backend.
    fn inner(&self) -> &Self::InnerBackend;

    /// Convert a tensor to require gradients.
    fn require_grad<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;

    /// Compute gradients via backpropagation.
    ///
    /// This method panics if:
    /// - The tensor dimension is > 6 (compile-time prevented by const generics)
    /// - The backward pass encounters an error
    ///
    /// For fallible gradient computation, use [`try_backward`](Self::try_backward).
    fn backward<const D: usize>(
        tensor: &Self::TensorPrimitive<D>,
    ) -> GradientMap<Self::InnerBackend>;

    /// Compute gradients via backpropagation, returning errors instead of panicking.
    ///
    /// This is the fallible version of [`backward`](Self::backward). It returns:
    /// - `Ok(GradientMap)` on success
    /// - `Err(AutogradError)` if the backward pass fails
    ///
    /// # Errors
    ///
    /// Returns `AutogradError::UnsupportedDimension` if the tensor dimension is > 6.
    /// Returns `AutogradError::BackwardError` if gradient computation fails.
    fn try_backward<const D: usize>(
        tensor: &Self::TensorPrimitive<D>,
    ) -> Result<GradientMap<Self::InnerBackend>, crate::error::BackendError>;
}

/// Map from tensor IDs to their gradients.
///
/// This structure stores computed gradients during backpropagation.
/// The `gradients` field will be used when `AutodiffBackend` implementations
/// populate gradients during the backward pass.
pub struct GradientMap<B: Backend> {
    #[allow(dead_code)] // Used by AutodiffBackend implementations
    gradients: std::collections::HashMap<usize, Box<dyn std::any::Any + Send>>,
    _marker: std::marker::PhantomData<B>,
}

impl<B: Backend> GradientMap<B> {
    /// Create a new empty gradient map.
    pub fn new() -> Self {
        Self {
            gradients: std::collections::HashMap::new(),
            _marker: std::marker::PhantomData,
        }
    }

    /// Get the gradient for a tensor by ID.
    pub fn get<const D: usize>(&self, _id: usize) -> Option<&B::TensorPrimitive<D>> {
        // Implementation would downcast from Any
        None
    }
}

impl<B: Backend> Default for GradientMap<B> {
    fn default() -> Self {
        Self::new()
    }
}

/// Marker trait for backends that support quantization.
pub trait QuantizedBackend: Backend {
    /// The quantization scheme.
    type QuantScheme;

    /// Quantize a tensor.
    fn quantize<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        scheme: &Self::QuantScheme,
    ) -> Self::TensorPrimitive<D>;

    /// Dequantize a tensor.
    fn dequantize<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D>;
}