fuel_tensor/

device.rs

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Device abstraction for CPU, CUDA, and Metal backends.
//!
//! ```rust
//! use fuel_tensor::Device;
//! let dev = Device::cpu();
//! assert!(dev.is_cpu());
//! assert_eq!(dev.location(), fuel_tensor::DeviceLocation::Cpu);
//! ```
use crate::{DType, HostBuffer, Result, Shape, WithDType};
use fuel_backend_contract::dyn_backend::DynBackendDevice;
use fuel_cpu_backend::dyn_impl::CpuBackendDevice;
use std::sync::Arc;

pub use fuel_ir::DeviceLocation;

/// A device on which tensors can be created and computations performed.
///
/// Internally this is an `Arc<dyn DynBackendDevice>` — a trait-object handle
/// to the concrete backend device.  CPU, CUDA, Metal, and arbitrary third-party
/// backends are all accessed through this single type.
///
/// # Example
///
/// ```rust
/// use fuel_tensor::{Device, DType};
/// use fuel_tensor::lazy::Tensor;
/// let dev = Device::cpu();
/// let t = Tensor::zeros((2, 3), DType::F32, &dev)?;
/// assert_eq!(t.shape().dims(), &[2, 3]);
/// # Ok::<(), fuel_tensor::Error>(())
/// ```
#[derive(Clone, Debug)]
pub struct Device {
    pub(crate) inner: Arc<dyn DynBackendDevice>,
}

/// Trait for host array-like values that describe their shape and yield CPU storage.
///
/// Implemented for scalars, arrays, slices, and nested vecs up to 4 dimensions.
/// It lets array-like Rust values report the [`Shape`] their data describes and
/// materialize into a [`HostBuffer`] — the raw material a lazy tensor is built from.
///
/// # Example
///
/// ```rust
/// use fuel_tensor::NdArray;
/// // Scalars, arrays, and nested arrays all implement `NdArray` and report
/// // the shape their data describes.
/// assert_eq!(3.14f32.shape()?.dims(), &[] as &[usize]);
/// assert_eq!((&[1f32, 2., 3.]).shape()?.dims(), &[3]);
/// assert_eq!((&[[1f32, 2.], [3., 4.]]).shape()?.dims(), &[2, 2]);
/// # Ok::<(), fuel_tensor::Error>(())
/// ```
pub trait NdArray {
    /// Returns the shape determined by this array-like value.
    fn shape(&self) -> Result<Shape>;

    /// Converts this value into CPU storage.
    fn to_cpu_storage(&self) -> HostBuffer;
}

impl<S: WithDType> NdArray for S {
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from(()))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        S::to_cpu_storage(&[*self])
    }
}

impl<S: WithDType, const N: usize> NdArray for &[S; N] {
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from(self.len()))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        S::to_cpu_storage(self.as_slice())
    }
}

impl<S: WithDType> NdArray for &[S] {
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from(self.len()))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        S::to_cpu_storage(self)
    }
}

impl<S: WithDType, const N: usize, const M: usize> NdArray for &[[S; N]; M] {
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from((M, N)))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        S::to_cpu_storage_owned(self.concat())
    }
}

impl<S: WithDType, const N1: usize, const N2: usize, const N3: usize> NdArray
    for &[[[S; N3]; N2]; N1]
{
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from((N1, N2, N3)))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        let mut vec = Vec::with_capacity(N1 * N2 * N3);
        for i1 in 0..N1 {
            for i2 in 0..N2 {
                vec.extend(self[i1][i2])
            }
        }
        S::to_cpu_storage_owned(vec)
    }
}

impl<S: WithDType, const N1: usize, const N2: usize, const N3: usize, const N4: usize> NdArray
    for &[[[[S; N4]; N3]; N2]; N1]
{
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from((N1, N2, N3, N4)))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        let mut vec = Vec::with_capacity(N1 * N2 * N3 * N4);
        for i1 in 0..N1 {
            for i2 in 0..N2 {
                for i3 in 0..N3 {
                    vec.extend(self[i1][i2][i3])
                }
            }
        }
        S::to_cpu_storage_owned(vec)
    }
}

impl<S: WithDType> NdArray for Vec<S> {
    fn shape(&self) -> Result<Shape> {
        Ok(Shape::from(self.len()))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        S::to_cpu_storage(self.as_slice())
    }
}

impl<S: WithDType> NdArray for Vec<&[S]> {
    fn shape(&self) -> Result<Shape> {
        if self.is_empty() {
            crate::bail!("empty array")
        }
        let n = self.len();
        let m = self[0].len();
        for v in self.iter() {
            if v.len() != m {
                crate::bail!("two elements have different len {m} {}", v.len())
            }
        }
        Ok(Shape::from((n, m)))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        let data = self.iter().copied().flatten().copied().collect::<Vec<_>>();
        S::to_cpu_storage_owned(data)
    }
}

impl<S: WithDType> NdArray for Vec<Vec<S>> {
    fn shape(&self) -> Result<Shape> {
        if self.is_empty() {
            crate::bail!("empty array")
        }
        let n = self.len();
        let m = self[0].len();
        for v in self.iter() {
            if v.len() != m {
                crate::bail!("two elements have different len {m} {}", v.len())
            }
        }
        Ok(Shape::from((n, m)))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        let len: usize = self.iter().map(|v| v.len()).sum();
        let mut dst = Vec::with_capacity(len);
        for v in self.iter() {
            dst.extend(v.iter().copied());
        }
        S::to_cpu_storage_owned(dst)
    }
}

impl<S: WithDType> NdArray for Vec<Vec<Vec<S>>> {
    fn shape(&self) -> Result<Shape> {
        if self.is_empty() {
            crate::bail!("empty array")
        }
        let shape0 = self[0].shape()?;
        let n = self.len();
        for v in self.iter() {
            let shape = v.shape()?;
            if shape != shape0 {
                crate::bail!("two elements have different shapes {shape:?} {shape0:?}")
            }
        }
        Ok(Shape::from([[n].as_slice(), shape0.dims()].concat()))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        if self.is_empty() {
            return S::to_cpu_storage_owned(vec![]);
        }
        let len: usize = self
            .iter()
            .map(|v| v.iter().map(|v| v.len()).sum::<usize>())
            .sum();
        let mut dst = Vec::with_capacity(len);
        for v1 in self.iter() {
            for v2 in v1.iter() {
                dst.extend(v2.iter().copied());
            }
        }
        S::to_cpu_storage_owned(dst)
    }
}

impl<S: WithDType> NdArray for Vec<Vec<Vec<Vec<S>>>> {
    fn shape(&self) -> Result<Shape> {
        if self.is_empty() {
            crate::bail!("empty array")
        }
        let shape0 = self[0].shape()?;
        let n = self.len();
        for v in self.iter() {
            let shape = v.shape()?;
            if shape != shape0 {
                crate::bail!("two elements have different shapes {shape:?} {shape0:?}")
            }
        }
        Ok(Shape::from([[n].as_slice(), shape0.dims()].concat()))
    }

    fn to_cpu_storage(&self) -> HostBuffer {
        let len: usize = self
            .iter()
            .map(|v| {
                v.iter()
                    .map(|v| v.iter().map(|v| v.len()).sum::<usize>())
                    .sum::<usize>()
            })
            .sum();
        let mut dst = Vec::with_capacity(len);
        for v1 in self.iter() {
            for v2 in v1.iter() {
                for v3 in v2.iter() {
                    dst.extend(v3.iter().copied());
                }
            }
        }
        S::to_cpu_storage_owned(dst)
    }
}

impl Device {
    /// Returns a CPU device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// let dev = Device::cpu();
    /// assert!(dev.is_cpu());
    /// ```
    pub fn cpu() -> Self {
        Device {
            inner: Arc::new(CpuBackendDevice),
        }
    }

    /// Creates a device backed by a custom [`DynBackendDevice`].
    ///
    /// The device uses dynamic dispatch for all operations, enabling
    /// third-party backends without modifying `fuel-core`.
    pub fn custom(device: Arc<dyn DynBackendDevice>) -> Self {
        Device { inner: device }
    }

    /// Borrow the underlying `Arc<dyn DynBackendDevice>`.
    ///
    /// Phase 7.5 work item G2: fuel-graph's slot-populating
    /// constructors (`Tensor::from_f32`, `const_f32_like`, etc.) take
    /// a `&Arc<dyn DynBackendDevice>` so they can allocate Storage on
    /// the right device without depending on fuel-core. Callers that
    /// hold a `&fuel_tensor::Device` thread it through with
    /// `dev.as_dyn()`.
    pub fn as_dyn(&self) -> &Arc<dyn DynBackendDevice> {
        &self.inner
    }

    /// Returns `true` if this is a custom (third-party) device.
    pub fn is_custom(&self) -> bool {
        !self.is_cpu() && !self.is_cuda() && !self.is_metal()
    }

    /// Sets the random seed for this device's random number generator.
    pub fn set_seed(&self, seed: u64) -> Result<()> {
        self.inner.set_seed_dyn(seed)
    }

    /// Returns the current random seed for this device.
    pub fn get_current_seed(&self) -> Result<u64> {
        self.inner.get_current_seed_dyn()
    }

    /// Returns `true` if both devices refer to the same physical device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// assert!(Device::cpu().same_device(&Device::cpu()));
    /// ```
    pub fn same_device(&self, rhs: &Self) -> bool {
        self.inner.same_device_dyn(&*rhs.inner)
    }

    /// Returns the physical [`DeviceLocation`] for this device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::{Device, DeviceLocation};
    /// assert_eq!(Device::cpu().location(), DeviceLocation::Cpu);
    /// ```
    pub fn location(&self) -> DeviceLocation {
        self.inner.location_dyn()
    }

    /// Returns `true` if this is the CPU device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// assert!(Device::cpu().is_cpu());
    /// ```
    pub fn is_cpu(&self) -> bool {
        matches!(self.location(), DeviceLocation::Cpu)
    }

    /// Returns `true` if this is a CUDA device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// assert!(!Device::cpu().is_cuda());
    /// ```
    pub fn is_cuda(&self) -> bool {
        matches!(self.location(), DeviceLocation::Cuda { .. })
    }

    /// Returns `true` if this is a Metal device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// assert!(!Device::cpu().is_metal());
    /// ```
    pub fn is_metal(&self) -> bool {
        matches!(self.location(), DeviceLocation::Metal { .. })
    }

    /// Returns `true` if this device has native BF16 support.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// // CPU does not have native BF16 support
    /// assert!(!Device::cpu().supports_bf16());
    /// ```
    pub fn supports_bf16(&self) -> bool {
        self.inner.supports_bf16()
    }

    /// Returns [`DType::BF16`] if supported, otherwise [`DType::F32`].
    ///
    /// # Example
    ///
    /// ```rust
    /// use fuel_tensor::{Device, DType};
    /// assert_eq!(Device::cpu().bf16_default_to_f32(), DType::F32);
    /// ```
    pub fn bf16_default_to_f32(&self) -> DType {
        if self.supports_bf16() {
            DType::BF16
        } else {
            DType::F32
        }
    }

    /// Synchronizes the device, waiting for all pending operations to complete.
    ///
    /// This is a no-op on CPU.
    ///
    /// ```rust
    /// use fuel_tensor::Device;
    /// Device::cpu().synchronize()?;
    /// # Ok::<(), fuel_tensor::Error>(())
    /// ```
    pub fn synchronize(&self) -> Result<()> {
        self.inner.synchronize_dyn()
    }
}