mirror of https://github.com/tracel-ai/burn.git
WIP
This commit is contained in:
parent
e9d0bf298d
commit
049ceff0e2
|
@ -11,6 +11,9 @@ extern crate derive_new;
|
||||||
|
|
||||||
extern crate alloc;
|
extern crate alloc;
|
||||||
|
|
||||||
|
/// Server module.
|
||||||
|
pub mod server;
|
||||||
|
|
||||||
mod tensor;
|
mod tensor;
|
||||||
|
|
||||||
/// Burn Tensor representaton
|
/// Burn Tensor representaton
|
||||||
|
|
|
@ -0,0 +1,102 @@
|
||||||
|
use core::marker::PhantomData;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
backend::{Backend, BackendBridge, DeviceOps},
|
||||||
|
quantization::QTensorPrimitive,
|
||||||
|
repr::{OperationDescription, TensorDescription},
|
||||||
|
};
|
||||||
|
|
||||||
|
pub struct Server<B: ServerBackend> {
|
||||||
|
r: PhantomData<B>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub trait ServerBackend: Send + Sync + 'static {
|
||||||
|
type Runtime: ServerRuntime;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub trait ServerRuntime {
|
||||||
|
type Client: ServerClient;
|
||||||
|
type Device: DeviceOps;
|
||||||
|
type Bridge;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ServerTensor<R: ServerRuntime> {
|
||||||
|
desc: TensorDescription,
|
||||||
|
client: R::Client,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<B: ServerBackend> core::fmt::Debug for Server<B> {
|
||||||
|
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||||
|
f.write_fmt(format_args!("server"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: ServerRuntime> core::fmt::Debug for ServerTensor<R> {
|
||||||
|
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||||
|
f.write_fmt(format_args!("tensor"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: ServerRuntime> Clone for ServerTensor<R> {
|
||||||
|
fn clone(&self) -> Self {
|
||||||
|
Self {
|
||||||
|
desc: self.desc.clone(),
|
||||||
|
client: self.client.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<B: ServerBackend> Clone for Server<B> {
|
||||||
|
fn clone(&self) -> Self {
|
||||||
|
Self { r: PhantomData }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<B: ServerBackend> Default for Server<B> {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self { r: PhantomData }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: ServerRuntime> QTensorPrimitive for ServerTensor<R> {
|
||||||
|
fn scheme(&self) -> &crate::quantization::QuantizationScheme {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn strategy(&self) -> crate::quantization::QuantizationStrategy {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub trait ServerClient: Clone + Send + Sync {
|
||||||
|
fn execute(&self, op: OperationDescription);
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<B: ServerBackend> Backend for Server<B>
|
||||||
|
where
|
||||||
|
<B::Runtime as ServerRuntime>::Bridge: BackendBridge<Self> + 'static,
|
||||||
|
{
|
||||||
|
type Device = <B::Runtime as ServerRuntime>::Device;
|
||||||
|
|
||||||
|
type FullPrecisionBridge = <B::Runtime as ServerRuntime>::Bridge;
|
||||||
|
|
||||||
|
type FloatTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
|
||||||
|
|
||||||
|
type FloatElem = f32;
|
||||||
|
|
||||||
|
type IntTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
|
||||||
|
|
||||||
|
type IntElem = i32;
|
||||||
|
|
||||||
|
type BoolTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
|
||||||
|
|
||||||
|
type QuantizedTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
|
||||||
|
|
||||||
|
fn name() -> String {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn seed(seed: u64) {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
}
|
|
@ -0,0 +1,4 @@
|
||||||
|
mod backend;
|
||||||
|
mod ops;
|
||||||
|
|
||||||
|
pub use backend::*;
|
|
@ -0,0 +1,6 @@
|
||||||
|
mod op_activation;
|
||||||
|
mod op_bool;
|
||||||
|
mod op_float;
|
||||||
|
mod op_int;
|
||||||
|
mod op_module;
|
||||||
|
mod op_qfloat;
|
|
@ -0,0 +1,4 @@
|
||||||
|
use crate::server::Server;
|
||||||
|
use crate::{ops::ActivationOps, server::ServerBackend};
|
||||||
|
|
||||||
|
impl<B: ServerBackend> ActivationOps<Self> for Server<B> {}
|
|
@ -0,0 +1,119 @@
|
||||||
|
use crate::{
|
||||||
|
ops::BoolTensorOps,
|
||||||
|
server::{Server, ServerBackend},
|
||||||
|
};
|
||||||
|
|
||||||
|
impl<B: ServerBackend> BoolTensorOps<Self> for Server<B> {
|
||||||
|
fn bool_empty<const D: usize>(
|
||||||
|
shape: crate::Shape<D>,
|
||||||
|
device: &crate::Device<Self>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_shape<const D: usize>(tensor: &crate::ops::BoolTensor<Self, D>) -> crate::Shape<D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_into_data<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
) -> impl core::future::Future<Output = crate::TensorData> + Send {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_from_data<const D: usize>(
|
||||||
|
data: crate::TensorData,
|
||||||
|
device: &crate::Device<Self>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_into_int<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
) -> crate::ops::IntTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_into_float<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
) -> crate::ops::FloatTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_device<const D: usize>(
|
||||||
|
tensor: &crate::ops::BoolTensor<Self, D>,
|
||||||
|
) -> crate::Device<Self> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_to_device<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
device: &crate::Device<Self>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_reshape<const D1: usize, const D2: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D1>,
|
||||||
|
shape: crate::Shape<D2>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D2> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_slice<const D1: usize, const D2: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D1>,
|
||||||
|
ranges: [core::ops::Range<usize>; D2],
|
||||||
|
) -> crate::ops::BoolTensor<Self, D1> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_slice_assign<const D1: usize, const D2: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D1>,
|
||||||
|
ranges: [core::ops::Range<usize>; D2],
|
||||||
|
value: crate::ops::BoolTensor<Self, D1>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D1> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_equal<const D: usize>(
|
||||||
|
lhs: crate::ops::BoolTensor<Self, D>,
|
||||||
|
rhs: crate::ops::BoolTensor<Self, D>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_not<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_swap_dims<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
dim1: usize,
|
||||||
|
dim2: usize,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_permute<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
axes: [usize; D],
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_flip<const D: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D>,
|
||||||
|
axes: &[usize],
|
||||||
|
) -> crate::ops::BoolTensor<Self, D> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bool_expand<const D1: usize, const D2: usize>(
|
||||||
|
tensor: crate::ops::BoolTensor<Self, D1>,
|
||||||
|
shape: crate::Shape<D2>,
|
||||||
|
) -> crate::ops::BoolTensor<Self, D2> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
}
|
|
@ -0,0 +1,4 @@
|
||||||
|
use crate::server::Server;
|
||||||
|
use crate::{ops::FloatTensorOps, server::ServerBackend};
|
||||||
|
|
||||||
|
impl<B: ServerBackend> FloatTensorOps<Self> for Server<B> {}
|
|
@ -0,0 +1,4 @@
|
||||||
|
use crate::server::Server;
|
||||||
|
use crate::{ops::IntTensorOps, server::ServerBackend};
|
||||||
|
|
||||||
|
impl<B: ServerBackend> IntTensorOps<Self> for Server<B> {}
|
|
@ -0,0 +1,7 @@
|
||||||
|
use crate::server::Server;
|
||||||
|
use crate::{
|
||||||
|
ops::{ConvOptions, ConvTransposeOptions, FloatTensor, ModuleOps},
|
||||||
|
server::ServerBackend,
|
||||||
|
};
|
||||||
|
|
||||||
|
impl<B: ServerBackend> ModuleOps<Self> for Server<B> {}
|
|
@ -0,0 +1,10 @@
|
||||||
|
use core::future::Future;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
ops::{FloatTensor, QTensorOps, QuantizedTensor},
|
||||||
|
quantization::{QuantizationParametersPrimitive, QuantizationScheme},
|
||||||
|
server::{Server, ServerBackend},
|
||||||
|
Device, TensorData,
|
||||||
|
};
|
||||||
|
|
||||||
|
impl<B: ServerBackend> QTensorOps<Self> for Server<B> {}
|
|
@ -60,8 +60,8 @@ pub trait Backend:
|
||||||
+ ActivationOps<Self>
|
+ ActivationOps<Self>
|
||||||
+ QTensorOps<Self>
|
+ QTensorOps<Self>
|
||||||
+ Clone
|
+ Clone
|
||||||
+ Sized
|
|
||||||
+ Default
|
+ Default
|
||||||
|
+ Sized
|
||||||
+ Send
|
+ Send
|
||||||
+ Sync
|
+ Sync
|
||||||
+ core::fmt::Debug
|
+ core::fmt::Debug
|
||||||
|
|
Loading…
Reference in New Issue