Move HandleContainer and Tensor Ops descriptions from burn-fusion to …

…burn-tensor (#1654)

* Move HandlerContainer and Tensor Ops description to burn-tensor

Move HandleContainer and Tensor operations descriptions to burn-tensor crate.
Removed the FusionDevice and replaced it with a DeviceOps trait bound to Backend::Device.

For now added modules to burn-tensor are excluded from no-std as they rely on Arc.

* [burn-tensor] Flatten module hierarchy for tensor representation

+ Add new repr feature to cargo file.

* Remove prefix on dosctring

* [burn-fusion] Require default features of burn-tensor

crates/burn-candle/src/backend.rs

@@ -1,6 +1,9 @@
 use std::marker::PhantomData;
 
-use burn_tensor::{backend::Backend, Device};
+use burn_tensor::{
+    backend::{Backend, DeviceId, DeviceOps},
+    Device,
+};
 use candle_core::DeviceLocation;
 
 use crate::{
@@ -60,6 +63,16 @@ impl From<candle_core::Device> for CandleDevice {
     }
 }
 
+impl DeviceOps for CandleDevice {
+    fn id(&self) -> burn_tensor::backend::DeviceId {
+        match self {
+            CandleDevice::Cpu => DeviceId::new(0, 0),
+            CandleDevice::Cuda(index) => DeviceId::new(1, *index as u32),
+            CandleDevice::Metal(index) => DeviceId::new(2, *index as u32),
+        }
+    }
+}
+
 impl Default for CandleDevice {
     fn default() -> Self {
         Self::Cpu

crates/burn-fusion/Cargo.toml

@@ -16,7 +16,7 @@ std = ["serde/std"]
 doc = ["default"]
 
 [dependencies]
-burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0" }
 burn-common = { path = "../burn-common", version = "0.14.0" }
 hashbrown = { workspace = true }
 derive-new = {workspace = true }

crates/burn-fusion/src/backend.rs

@@ -1,15 +1,17 @@
 use crate::{
-    client::FusionClient,
-    stream::{Context, OperationDescription},
-    FusionClientLocator, FusionTensor, PrecisionBridge,
+    client::FusionClient, stream::Context, FusionClientLocator, FusionTensor, PrecisionBridge,
+};
+use burn_tensor::{
+    backend::Backend,
+    repr::{OperationDescription, ReprBackend},
+    Device,
 };
-use burn_tensor::{backend::Backend, Device, Shape};
 use serde::{de::DeserializeOwned, Serialize};
 use std::marker::PhantomData;
 
 pub(crate) static CLIENTS: FusionClientLocator = FusionClientLocator::new();
 
-pub(crate) fn get_client<B: FusionBackend>(device: &B::FusionDevice) -> B::FusionClient {
+pub(crate) fn get_client<B: FusionBackend>(device: &B::Device) -> B::FusionClient {
     CLIENTS.client(device)
 }
 
@@ -43,7 +45,7 @@ impl<B: FusionBackend> Backend for Fusion<B> {
     }
 
     fn sync(device: &Self::Device) {
-        let client = CLIENTS.client::<B::FusionClient>(&device.clone().into());
+        let client = CLIENTS.client::<B::FusionClient>(&device.clone());
         client.drain();
         B::sync(device)
     }
@@ -114,62 +116,17 @@ pub trait Optimization<B: FusionBackend>: Send {
     fn from_state(device: &B::Device, state: B::OptimizationState) -> Self;
 }
 
-/// The device id.
-#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy, new)]
-pub struct DeviceId {
-    /// The type id identifies the type of the device.
-    pub type_id: u16,
-    /// The index id identifies the device number.
-    pub index_id: u32,
-}
-
-/// The handle device trait allows to get an id for a backend device.
-pub trait FusionDevice: Clone + Send + Sync + PartialEq {
-    /// Return the [device id](DeviceId).
-    fn id(&self) -> DeviceId;
-}
-
 /// Trait that allows an existing [backend](Backend) to specify graph optimizations using
 /// [operation builder](crate::OptimizationBuilder).
-pub trait FusionBackend: Backend {
+pub trait FusionBackend: Backend + ReprBackend {
     /// The state that can be serialized for an optimization.
     type OptimizationState: Serialize + DeserializeOwned;
     /// Optimization type for the backend.
     type Optimization: Optimization<Self>;
-
-    /// The device type that can return an ID.
-    ///
-    /// It can be the same as (Backend::Device), but must implement (FusionDevice).
-    type FusionDevice: FusionDevice + From<Self::Device> + Into<Self::Device> + core::fmt::Debug;
-    /// The type that can be used to point to a tensor of any kind.
-    type Handle: Sync + Send + Clone;
     /// What kind of client should be used.
     type FusionClient: FusionClient<FusionBackend = Self>;
 
     /// The list of optimizations that will be used to optimize the computational graph.
     fn optimizations(device: Device<Self>)
         -> Vec<Box<dyn OptimizationBuilder<Self::Optimization>>>;
-
-    /// Convert a [handle](FusionBackend::Handle) to a [float tensor](Backend::FloatTensorPrimitive).
-    fn float_tensor<const D: usize>(
-        handle: Self::Handle,
-        shape: Shape<D>,
-    ) -> Self::FloatTensorPrimitive<D>;
-    /// Convert a [handle](FusionBackend::Handle) to an [int tensor](Backend::IntTensorPrimitive).
-    fn int_tensor<const D: usize>(
-        handle: Self::Handle,
-        shape: Shape<D>,
-    ) -> Self::IntTensorPrimitive<D>;
-    /// Convert a [handle](FusionBackend::Handle) to a [bool tensor](Backend::BoolTensorPrimitive).
-    fn bool_tensor<const D: usize>(
-        handle: Self::Handle,
-        shape: Shape<D>,
-    ) -> Self::BoolTensorPrimitive<D>;
-
-    /// Convert a [float tensor](Backend::FloatTensorPrimitive) to a [handle](FusionBackend::Handle).
-    fn float_tensor_handle<const D: usize>(tensor: Self::FloatTensorPrimitive<D>) -> Self::Handle;
-    /// Convert an [int tensor](Backend::IntTensorPrimitive) to a [handle](FusionBackend::Handle).
-    fn int_tensor_handle<const D: usize>(tensor: Self::IntTensorPrimitive<D>) -> Self::Handle;
-    /// Convert a [bool tensor](Backend::BoolTensorPrimitive) to a [handle](FusionBackend::Handle).
-    fn bool_tensor_handle<const D: usize>(tensor: Self::BoolTensorPrimitive<D>) -> Self::Handle;
 }

crates/burn-fusion/src/client/base.rs

@@ -1,19 +1,21 @@
 use crate::{
-    stream::{Operation, OperationDescription, StreamId},
-    FusionBackend, FusionTensor, Handle, TensorDescription, TensorId,
+    stream::{execution::Operation, StreamId},
+    FusionBackend, FusionTensor, Handle,
 };
 use burn_tensor::{
+    backend::Backend,
     ops::{FloatElem, IntElem},
-    Data, Reader,
+    repr::{OperationDescription, TensorDescription, TensorId},
+    Data, Device, Reader,
 };
 
 /// Define how to interact with the fusion server.
 pub trait FusionClient: Send + Sync + Clone {
     /// The [fusion backend](FusionBackend) associated type.
     type FusionBackend: FusionBackend;
 
-    /// Create a new client for the given [fusion device](FusionBackend::FusionDevice).
-    fn new(device: <Self::FusionBackend as FusionBackend>::FusionDevice) -> Self;
+    /// Create a new client for the given [device](Backend::Device).
+    fn new(device: Device<Self::FusionBackend>) -> Self;
     /// Register a new [tensor operation description](OperationDescription).
     fn register<O: Operation<Self::FusionBackend> + 'static>(
         &self,
@@ -24,7 +26,7 @@ pub trait FusionClient: Send + Sync + Clone {
     /// Register all lazy computation.
     fn drain(&self);
     /// Get the current device used by all operations handled by this client.
-    fn device(&self) -> &<Self::FusionBackend as FusionBackend>::FusionDevice;
+    fn device(&self) -> &<Self::FusionBackend as Backend>::Device;
     /// Create a new [fusion tensor](FusionTensor), but with no resources allocated to it.
     fn tensor_uninitialized(&self, shape: Vec<usize>) -> FusionTensor<Self>;
     /// Create a tensor with the given handle and shape.

crates/burn-fusion/src/client/mutex.rs

@@ -1,9 +1,13 @@
 use super::FusionClient;
 use crate::{
-    stream::{Operation, OperationDescription, StreamId},
+    stream::{execution::Operation, StreamId},
     FusionBackend, FusionServer, FusionTensor, Handle,
 };
-use burn_tensor::ops::FloatElem;
+use burn_tensor::{
+    backend::Backend,
+    ops::FloatElem,
+    repr::{OperationDescription, TensorDescription, TensorId},
+};
 use spin::Mutex;
 use std::sync::Arc;
 
@@ -13,7 +17,7 @@ where
     B: FusionBackend,
 {
     server: Arc<Mutex<FusionServer<B>>>,
-    device: B::FusionDevice,
+    device: B::Device,
 }
 
 impl<B> Clone for MutexFusionClient<B>
@@ -34,7 +38,7 @@ where
 {
     type FusionBackend = B;
 
-    fn new(device: B::FusionDevice) -> Self {
+    fn new(device: B::Device) -> Self {
         Self {
             device: device.clone(),
             server: Arc::new(Mutex::new(FusionServer::new(device))),
@@ -63,7 +67,7 @@ where
         FusionTensor::new(id, shape, self.clone(), StreamId::current())
     }
 
-    fn device(&self) -> &<Self::FusionBackend as FusionBackend>::FusionDevice {
+    fn device(&self) -> &<Self::FusionBackend as Backend>::Device {
         &self.device
     }
     fn register_tensor(
@@ -82,15 +86,15 @@ where
 
     fn read_tensor_float<const D: usize>(
         &self,
-        tensor: crate::TensorDescription,
+        tensor: TensorDescription,
         stream: StreamId,
     ) -> burn_tensor::Reader<burn_tensor::Data<FloatElem<Self::FusionBackend>, D>> {
         self.server.lock().read_float(tensor, stream)
     }
 
     fn read_tensor_int<const D: usize>(
         &self,
-        tensor: crate::TensorDescription,
+        tensor: TensorDescription,
         id: StreamId,
     ) -> burn_tensor::Reader<burn_tensor::Data<burn_tensor::ops::IntElem<Self::FusionBackend>, D>>
     {
@@ -99,25 +103,24 @@ where
 
     fn read_tensor_bool<const D: usize>(
         &self,
-        tensor: crate::TensorDescription,
+        tensor: TensorDescription,
         stream: StreamId,
     ) -> burn_tensor::Reader<burn_tensor::Data<bool, D>> {
         self.server.lock().read_bool(tensor, stream)
     }
 
     fn change_client_float<const D: usize>(
         &self,
-        tensor: crate::TensorDescription,
+        tensor: TensorDescription,
         client: Self,
         stream: StreamId,
     ) -> FusionTensor<Self> {
-        let device = client.device.clone().into();
-
         let mut server_other = client.server.lock();
         let mut server_current = self.server.lock();
         server_current.drain_stream(stream);
 
-        let id = server_current.change_server_float::<D>(&tensor, &device, &mut server_other);
+        let id =
+            server_current.change_server_float::<D>(&tensor, &client.device, &mut server_other);
 
         core::mem::drop(server_other);
         core::mem::drop(server_current);
@@ -127,17 +130,15 @@ where
 
     fn change_client_int<const D: usize>(
         &self,
-        tensor: crate::TensorDescription,
+        tensor: TensorDescription,
         client: Self,
         stream: StreamId,
     ) -> FusionTensor<Self> {
-        let device = client.device.clone().into();
-
         let mut server_other = client.server.lock();
         let mut server_current = self.server.lock();
         server_current.drain_stream(stream);
 
-        let id = server_current.change_server_int::<D>(&tensor, &device, &mut server_other);
+        let id = server_current.change_server_int::<D>(&tensor, &client.device, &mut server_other);
 
         core::mem::drop(server_other);
         core::mem::drop(server_current);
@@ -147,25 +148,23 @@ where
 
     fn change_client_bool<const D: usize>(
         &self,
-        tensor: crate::TensorDescription,
+        tensor: TensorDescription,
         client: Self,
         stream: StreamId,
     ) -> FusionTensor<Self> {
-        let device = client.device.clone().into();
-
         let mut server_other = client.server.lock();
         let mut server_current = self.server.lock();
         server_current.drain_stream(stream);
 
-        let id = server_current.change_server_bool::<D>(&tensor, &device, &mut server_other);
+        let id = server_current.change_server_bool::<D>(&tensor, &client.device, &mut server_other);
 
         core::mem::drop(server_other);
         core::mem::drop(server_current);
 
         FusionTensor::new(id, tensor.shape, client, StreamId::current())
     }
 
-    fn register_orphan(&self, id: &crate::TensorId) {
+    fn register_orphan(&self, id: &TensorId) {
         self.server.lock().drop_tensor_handle(*id);
     }
 }

crates/burn-fusion/src/fusion.rs

@@ -1,8 +1,14 @@
-use crate::{client::FusionClient, DeviceId, FusionBackend, FusionDevice};
+use burn_tensor::{
+    backend::{Backend, DeviceId, DeviceOps},
+    repr::ReprBackend,
+};
+
+use crate::client::FusionClient;
+
 use std::{any::Any, collections::HashMap, ops::DerefMut};
 
-/// Type alias for [fusion backend handle](FusionBackend::Handle).
-pub type Handle<B> = <B as FusionBackend>::Handle;
+/// Type alias for [representation backend handle](burn_tensor::repr::ReprBackend::Handle).
+pub type Handle<B> = <B as ReprBackend>::Handle;
 type Key = (core::any::TypeId, DeviceId);
 
 pub(crate) struct FusionClientLocator {
@@ -22,7 +28,7 @@ impl FusionClientLocator {
     /// Provide the init function to create a new client if it isn't already initialized.
     pub fn client<C: FusionClient + 'static>(
         &self,
-        device: &<C::FusionBackend as FusionBackend>::FusionDevice,
+        device: &<C::FusionBackend as Backend>::Device,
     ) -> C {
         let device_id = device.id();
         let client_id = (core::any::TypeId::of::<C>(), device_id);

crates/burn-fusion/src/lib.rs

@@ -16,7 +16,6 @@ pub mod stream;
 mod backend;
 mod bridge;
 mod fusion;
-mod handle;
 mod ops;
 mod server;
 mod tensor;
@@ -26,5 +25,4 @@ pub(crate) use server::*;
 pub use backend::*;
 pub use bridge::*;
 pub use fusion::*;
-pub use handle::*;
 pub use tensor::*;