sparse

Author: McArthur-Alford

Cargo.lock

@@ -68,6 +68,7 @@ vision = ["burn-dataset?/vision", "burn-common/network"]
 # Backend
 autodiff = ["burn-autodiff"]
 fusion = ["burn-wgpu?/fusion"]
+sparse = ["burn-sparse"]
 
 ## Backend features
 metal = ["burn-candle?/metal"]
@@ -111,6 +112,7 @@ burn-wgpu = { path = "../burn-wgpu", version = "0.14.0", optional = true, defaul
 burn-autodiff = { path = "../burn-autodiff", version = "0.14.0", optional = true }
 burn-tch = { path = "../burn-tch", version = "0.14.0", optional = true }
 burn-candle = { path = "../burn-candle", version = "0.14.0", optional = true }
+burn-sparse = { path = "../burn-sparse", version = "0.14.0", optional = true }
 
 derive-new = { workspace = true }
 log = { workspace = true, optional = true }

crates/burn-core/Cargo.toml

@@ -27,3 +27,6 @@ pub use burn_tch as libtorch;
 
 #[cfg(feature = "tch")]
 pub use burn_tch::LibTorch;
+
+#[cfg(feature = "sparse")]
+pub use burn_sparse as sparse;

crates/burn-core/src/backend.rs

@@ -0,0 +1,43 @@
+[package]
+authors = []
+categories = ["science", "no-std", "embedded", "wasm"]
+description = "Sparse tensor crate that offers a default sparse backend wrapper around burn backends."
+edition.workspace = true
+keywords = ["deep-learning", "machine-learning", "tensor", "sparse"]
+license.workspace = true
+name = "burn-sparse"
+readme.workspace = true
+repository = "https://github.com/tracel-ai/burn/tree/main/burn-sparse"
+version.workspace = true
+
+[features]
+default = ["std"]
+doc = ["default"]
+experimental-named-tensor = []
+std = ["rand/std", "half/std", "num-traits/std"]
+wasm-sync = []
+
+[dependencies]
+burn-common = { path = "../burn-common", version = "0.14.0", default-features = false }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0" }
+
+proc-macro2 = { workspace = true }
+quote = { workspace = true }
+syn = { workspace = true }
+derive-new = { workspace = true }
+half = { workspace = true }
+num-traits = { workspace = true }
+rand = { workspace = true }
+rand_distr = { workspace = true } # use instead of statrs because it supports no_std
+
+# The same implementation of HashMap in std but with no_std support (only needs alloc crate)
+hashbrown = { workspace = true } # no_std compatible
+
+# Serialization
+serde = { workspace = true }
+
+[dev-dependencies]
+rand = { workspace = true, features = ["std", "std_rng"] } # Default enables std
+
+[package.metadata.docs.rs]
+features = ["doc"]

crates/burn-sparse/Cargo.toml

@@ -0,0 +1,4 @@
+use crate::backend::SparseBackend;
+
+/// Sparse tensor primitive type used by the backend.
+pub type SparseTensor<B, const D: usize> = <B as SparseBackend>::SparseTensorPrimitive<D>;

crates/burn-sparse/src/backend/alias.rs

@@ -0,0 +1,33 @@
+use crate::backend::{Sparse, SparseBackend};
+use burn_tensor::{Int, Tensor, TensorPrimitive};
+
+pub trait SparseTensor<const D: usize, B>
+where
+    B: SparseBackend,
+{
+    fn dense_int(self) -> Tensor<B, D, Int>;
+    fn spmm(self, rhs: Tensor<B, D>) -> Tensor<B, D>;
+    fn dense(self) -> Tensor<B, D>;
+}
+
+impl<const D: usize, B> SparseTensor<D, B> for Tensor<B, D, Sparse>
+where
+    B: SparseBackend,
+{
+    fn dense(self) -> Tensor<B, D> {
+        Tensor::new(TensorPrimitive::Float(B::sparse_to_dense(
+            self.into_primitive(),
+        )))
+    }
+
+    fn dense_int(self) -> Tensor<B, D, Int> {
+        self.dense().int()
+    }
+
+    fn spmm(self, rhs: Tensor<B, D>) -> Tensor<B, D> {
+        Tensor::new(TensorPrimitive::Float(B::sparse_spmm(
+            self.into_primitive(),
+            rhs.into_primitive().tensor(),
+        )))
+    }
+}

crates/burn-sparse/src/backend/api.rs

@@ -0,0 +1,61 @@
+use std::{future::Future, ops::Range};
+
+use crate::backend::SparseBackend;
+use burn_tensor::{backend::Backend, BasicOps, Shape, TensorData, TensorKind};
+
+/// A type-level representation of the kind of a sparse (float) tensor.
+#[derive(Clone, Debug)]
+pub struct Sparse;
+
+impl<B: SparseBackend> TensorKind<B> for Sparse {
+    type Primitive<const D: usize> = B::SparseTensorPrimitive<D>;
+    fn name() -> &'static str {
+        "Sparse"
+    }
+}
+
+impl<B: SparseBackend> BasicOps<B> for Sparse {
+    type Elem = B::FloatElem;
+
+    fn into_data_async<const D: usize>(
+        tensor: Self::Primitive<D>,
+    ) -> impl Future<Output = TensorData> + Send {
+        B::sparse_into_data(tensor)
+    }
+
+    fn device<const D: usize>(tensor: &Self::Primitive<D>) -> <B as Backend>::Device {
+        B::sparse_device(tensor)
+    }
+
+    fn to_device<const D: usize>(
+        tensor: Self::Primitive<D>,
+        device: &<B as Backend>::Device,
+    ) -> Self::Primitive<D> {
+        B::sparse_to_device(tensor, device)
+    }
+
+    fn from_data<const D: usize>(
+        data: TensorData,
+        device: &<B as Backend>::Device,
+    ) -> Self::Primitive<D> {
+        B::sparse_from_data(data, device)
+    }
+
+    fn shape<const D: usize>(tensor: &Self::Primitive<D>) -> Shape<D> {
+        B::sparse_shape(tensor)
+    }
+
+    fn empty<const D: usize>(
+        shape: Shape<D>,
+        device: &<B as Backend>::Device,
+    ) -> Self::Primitive<D> {
+        B::sparse_empty(shape, device)
+    }
+
+    fn slice<const D1: usize, const D2: usize>(
+        tensor: Self::Primitive<D1>,
+        ranges: [Range<usize>; D2],
+    ) -> Self::Primitive<D1> {
+        B::sparse_slice(tensor, ranges)
+    }
+}

crates/burn-sparse/src/backend/kind.rs

@@ -0,0 +1,8 @@
+mod alias;
+mod api;
+mod kind;
+mod sparse_backend;
+
+pub use alias::*;
+pub use kind::*;
+pub use sparse_backend::*;

crates/burn-sparse/src/backend/mod.rs

@@ -0,0 +1,110 @@
+use crate::backend::SparseTensor;
+use burn_tensor::{backend::Backend, Device, Shape, TensorData};
+use core::{future::Future, ops::Range};
+
+pub trait SparseBackend: Backend {
+    type SparseTensorPrimitive<const D: usize>: Clone + Send + 'static + core::fmt::Debug;
+
+    fn sparse_empty<const D: usize>(
+        shape: Shape<D>,
+        device: &Device<Self>,
+    ) -> SparseTensor<Self, D>;
+
+    fn sparse_to_sparse<const D: usize>(
+        dense: Self::FloatTensorPrimitive<D>,
+    ) -> Self::SparseTensorPrimitive<D>;
+
+    fn sparse_to_dense<const D: usize>(
+        sparse: Self::SparseTensorPrimitive<D>,
+    ) -> Self::FloatTensorPrimitive<D>;
+
+    fn sparse_spmm<const D: usize>(
+        lhs: Self::SparseTensorPrimitive<D>,
+        rhs: Self::FloatTensorPrimitive<D>,
+    ) -> Self::FloatTensorPrimitive<D>;
+
+    fn sparse_sddmm<const D: usize>(
+        lhs: Self::SparseTensorPrimitive<D>,
+        rhs: Self::FloatTensorPrimitive<D>,
+    ) -> Self::SparseTensorPrimitive<D>;
+
+    /// Gets the element at the given indices.
+    ///
+    /// # Arguments
+    ///
+    /// * `tensor` - The tensor.
+    /// * `indices` - The indices.
+    ///
+    /// # Returns
+    ///
+    /// The elements at the given indices.
+    fn sparse_slice<const D1: usize, const D2: usize>(
+        tensor: SparseTensor<Self, D1>,
+        indices: [Range<usize>; D2],
+    ) -> SparseTensor<Self, D1>;
+
+    /// Gets the device of the tensor.
+    ///
+    /// # Arguments
+    ///
+    /// * `tensor` - The tensor.
+    ///
+    /// # Returns
+    ///
+    /// The device of the tensor.
+    fn sparse_device<const D: usize>(tensor: &SparseTensor<Self, D>) -> Device<Self>;
+
+    /// Moves the tensor to the given device.
+    ///
+    /// # Arguments
+    ///
+    /// * `tensor` - The tensor.
+    /// * `device` - The device to move the tensor to.
+    ///
+    /// # Returns
+    ///
+    /// The tensor on the given device.
+    fn sparse_to_device<const D: usize>(
+        tensor: SparseTensor<Self, D>,
+        device: &Device<Self>,
+    ) -> SparseTensor<Self, D>;
+
+    /// Gets the shape of the tensor.
+    ///
+    /// # Arguments
+    ///
+    /// * `tensor` - The tensor.
+    ///
+    /// # Returns
+    ///
+    /// The shape of the tensor.
+    fn sparse_shape<const D: usize>(tensor: &SparseTensor<Self, D>) -> Shape<D>;
+
+    /// Converts the tensor to a data structure.
+    ///
+    /// # Arguments
+    ///
+    /// * `tensor` - The tensor.
+    ///
+    /// # Returns
+    ///
+    /// The data structure with the tensor's data.
+    fn sparse_into_data<const D: usize>(
+        tensor: SparseTensor<Self, D>,
+    ) -> impl Future<Output = TensorData> + Send;
+
+    /// Creates a tensor from the data structure.
+    ///
+    /// # Arguments
+    ///
+    /// * `data` - The data structure.
+    /// * `device` - The device to create the tensor on.
+    ///
+    /// # Returns
+    ///
+    /// The tensor with the data.
+    fn sparse_from_data<const D: usize>(
+        data: TensorData,
+        device: &Device<Self>,
+    ) -> SparseTensor<Self, D>;
+}

crates/burn-sparse/src/backend/sparse_backend.rs

@@ -0,0 +1,42 @@
+use crate::decorator::FullPrecisionBridge;
+use crate::decorator::SparseRepresentation;
+use burn_tensor::backend::Backend;
+use core::marker::PhantomData;
+use derive_new::new;
+
+/// Tensor backend that extends existing backends with sparse tensor support.
+/// This backend abstracts over all backends, and so lacks the performance of a direct implementation.
+/// Backends implementing SparseDecorator should be used directly where possible.
+#[derive(new, Clone, Copy, Default, Debug)]
+pub struct SparseDecorator<B: Backend, R: SparseRepresentation> {
+    _p: PhantomData<B>,
+    _r: PhantomData<R>,
+}
+
+impl<B: Backend, R: SparseRepresentation> Backend for SparseDecorator<B, R> {
+    type Device = B::Device;
+
+    type FullPrecisionBridge = FullPrecisionBridge<B::FullPrecisionBridge>;
+
+    type FloatTensorPrimitive<const D: usize> = B::FloatTensorPrimitive<D>;
+
+    type FloatElem = B::FloatElem;
+
+    type IntTensorPrimitive<const D: usize> = B::IntTensorPrimitive<D>;
+
+    type IntElem = B::IntElem;
+
+    type BoolTensorPrimitive<const D: usize> = B::BoolTensorPrimitive<D>;
+
+    type QuantizedTensorPrimitive<const D: usize> = B::QuantizedTensorPrimitive<D>;
+
+    fn name() -> String {
+        format!("SparseDecorator<{}>", B::name())
+    }
+
+    fn seed(seed: u64) {
+        B::seed(seed)
+    }
+}
+
+impl<B: Backend, R: SparseRepresentation> SparseDecorator<B, R> {}

crates/burn-sparse/src/decorator/backend.rs

@@ -0,0 +1,10 @@
+mod backend;
+mod ops;
+mod precision_bridge;
+mod representation;
+mod sparse_coo;
+mod sparse_csr;
+
+pub use backend::*;
+pub use precision_bridge::*;
+pub use representation::*;