tenferro_ad/
eager_backend.rs

1use tenferro_cpu::CpuBackend;
2#[cfg(feature = "cuda")]
3use tenferro_gpu::CudaBackend;
4#[cfg(feature = "webgpu")]
5use tenferro_gpu::WebGpuBackend;
6use tenferro_tensor::backend::ElementwiseFusionPlan;
7use tenferro_tensor::{
8    BackendCachedDot, BackendRuntimeCache, BackendSession, BackendSessionHost, CompareDir, DType,
9    DotGeneralConfig, GatherConfig, PadConfig, Result as TensorResult, ScatterConfig, SliceConfig,
10    Tensor, TensorAnalytic, TensorBackend, TensorBuffer, TensorDeviceTransfer, TensorDot,
11    TensorElementwise, TensorFusion, TensorIndexing, TensorRead, TensorReduction, TensorStructural,
12    TensorValue,
13};
14
15pub enum EagerBackend {
16    Cpu(CpuBackend),
17    #[cfg(feature = "cuda")]
18    Cuda(CudaBackend),
19    #[cfg(feature = "webgpu")]
20    WebGpu(WebGpuBackend),
21}
22
23impl std::fmt::Debug for EagerBackend {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        match self {
26            Self::Cpu(backend) => f.debug_tuple("Cpu").field(backend).finish(),
27            #[cfg(feature = "cuda")]
28            Self::Cuda(backend) => f.debug_tuple("Cuda").field(backend).finish(),
29            #[cfg(feature = "webgpu")]
30            Self::WebGpu(backend) => f.debug_tuple("WebGpu").field(backend).finish(),
31        }
32    }
33}
34
35impl EagerBackend {
36    pub(crate) fn cpu(backend: CpuBackend) -> Self {
37        Self::Cpu(backend)
38    }
39
40    #[cfg(feature = "cuda")]
41    pub(crate) fn cuda(backend: CudaBackend) -> Self {
42        Self::Cuda(backend)
43    }
44
45    #[cfg(feature = "webgpu")]
46    pub(crate) fn webgpu(backend: WebGpuBackend) -> Self {
47        Self::WebGpu(backend)
48    }
49
50    pub(crate) fn synchronize(&mut self) -> TensorResult<()> {
51        match self {
52            Self::Cpu(_) => Ok(()),
53            #[cfg(feature = "cuda")]
54            Self::Cuda(backend) => backend.runtime().synchronize(),
55            #[cfg(feature = "webgpu")]
56            Self::WebGpu(backend) => backend.synchronize(),
57        }
58    }
59}
60
61macro_rules! dispatch {
62    ($backend:expr, $method:ident($($arg:expr),* $(,)?)) => {
63        match $backend {
64            EagerBackend::Cpu(backend) => backend.$method($($arg),*),
65            #[cfg(feature = "cuda")]
66            EagerBackend::Cuda(backend) => backend.$method($($arg),*),
67            #[cfg(feature = "webgpu")]
68            EagerBackend::WebGpu(backend) => backend.$method($($arg),*),
69        }
70    };
71}
72
73macro_rules! delegate_tensor_backend_methods {
74    ($(fn $method:ident($($arg:ident: $ty:ty),* $(,)?) -> $ret:ty;)*) => {
75        $(
76            fn $method(&mut self, $($arg: $ty),*) -> $ret {
77                dispatch!(self, $method($($arg),*))
78            }
79        )*
80    };
81}
82
83impl BackendRuntimeCache for EagerBackend {
84    type RuntimeCache = ();
85}
86
87impl TensorElementwise for EagerBackend {
88    delegate_tensor_backend_methods! {
89        fn add(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
90        fn add_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
91        fn sub(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
92        fn sub_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
93        fn mul(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
94        fn mul_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
95        fn neg(input: &Tensor) -> TensorResult<Tensor>;
96        fn neg_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
97        fn conj(input: &Tensor) -> TensorResult<Tensor>;
98        fn conj_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
99        fn div(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
100        fn div_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
101        fn rem(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
102        fn rem_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
103        fn abs(input: &Tensor) -> TensorResult<Tensor>;
104        fn abs_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
105        fn sign(input: &Tensor) -> TensorResult<Tensor>;
106        fn sign_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
107        fn maximum(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
108        fn maximum_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
109        fn minimum(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
110        fn minimum_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
111        fn compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> TensorResult<Tensor>;
112        fn compare_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>, dir: &CompareDir) -> TensorResult<Tensor>;
113        fn select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) -> TensorResult<Tensor>;
114        fn select_read(pred: TensorRead<'_>, on_true: TensorRead<'_>, on_false: TensorRead<'_>) -> TensorResult<Tensor>;
115        fn clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) -> TensorResult<Tensor>;
116        fn clamp_read(input: TensorRead<'_>, lower: TensorRead<'_>, upper: TensorRead<'_>) -> TensorResult<Tensor>;
117    }
118}
119
120impl TensorAnalytic for EagerBackend {
121    delegate_tensor_backend_methods! {
122        fn exp(input: &Tensor) -> TensorResult<Tensor>;
123        fn exp_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
124        fn log(input: &Tensor) -> TensorResult<Tensor>;
125        fn log_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
126        fn sin(input: &Tensor) -> TensorResult<Tensor>;
127        fn sin_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
128        fn cos(input: &Tensor) -> TensorResult<Tensor>;
129        fn cos_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
130        fn tanh(input: &Tensor) -> TensorResult<Tensor>;
131        fn tanh_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
132        fn sqrt(input: &Tensor) -> TensorResult<Tensor>;
133        fn sqrt_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
134        fn rsqrt(input: &Tensor) -> TensorResult<Tensor>;
135        fn rsqrt_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
136        fn pow(lhs: &Tensor, rhs: &Tensor) -> TensorResult<Tensor>;
137        fn pow_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> TensorResult<Tensor>;
138        fn expm1(input: &Tensor) -> TensorResult<Tensor>;
139        fn expm1_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
140        fn log1p(input: &Tensor) -> TensorResult<Tensor>;
141        fn log1p_read(input: TensorRead<'_>) -> TensorResult<Tensor>;
142    }
143}
144
145impl TensorStructural for EagerBackend {
146    delegate_tensor_backend_methods! {
147        fn transpose(input: &Tensor, perm: &[usize]) -> TensorResult<Tensor>;
148        fn reshape(input: &Tensor, shape: &[usize]) -> TensorResult<Tensor>;
149        fn reshape_read(input: TensorRead<'_>, shape: &[usize]) -> TensorResult<Tensor>;
150        fn broadcast_in_dim(input: &Tensor, shape: &[usize], dims: &[usize]) -> TensorResult<Tensor>;
151        fn broadcast_in_dim_read(input: TensorRead<'_>, shape: &[usize], dims: &[usize]) -> TensorResult<Tensor>;
152        fn cast(input: &Tensor, to: DType) -> TensorResult<Tensor>;
153        fn extract_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> TensorResult<Tensor>;
154        fn embed_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> TensorResult<Tensor>;
155        fn tril(input: &Tensor, k: i64) -> TensorResult<Tensor>;
156        fn triu(input: &Tensor, k: i64) -> TensorResult<Tensor>;
157    }
158}
159
160impl TensorReduction for EagerBackend {
161    delegate_tensor_backend_methods! {
162        fn reduce_sum(input: &Tensor, axes: &[usize]) -> TensorResult<Tensor>;
163        fn reduce_prod(input: &Tensor, axes: &[usize]) -> TensorResult<Tensor>;
164        fn reduce_max(input: &Tensor, axes: &[usize]) -> TensorResult<Tensor>;
165        fn reduce_min(input: &Tensor, axes: &[usize]) -> TensorResult<Tensor>;
166    }
167}
168
169impl TensorDot for EagerBackend {
170    delegate_tensor_backend_methods! {
171        fn dot_general(lhs: &Tensor, rhs: &Tensor, config: &DotGeneralConfig) -> TensorResult<Tensor>;
172        fn dot_general_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>, config: &DotGeneralConfig) -> TensorResult<Tensor>;
173        fn dot_general_with_conj(lhs: &Tensor, rhs: &Tensor, config: &DotGeneralConfig, lhs_conj: bool, rhs_conj: bool) -> TensorResult<Tensor>;
174    }
175}
176
177impl TensorIndexing for EagerBackend {
178    delegate_tensor_backend_methods! {
179        fn gather(operand: &Tensor, start_indices: &Tensor, config: &GatherConfig) -> TensorResult<Tensor>;
180        fn scatter(operand: &Tensor, scatter_indices: &Tensor, updates: &Tensor, config: &ScatterConfig) -> TensorResult<Tensor>;
181        fn slice(input: &Tensor, config: &SliceConfig) -> TensorResult<Tensor>;
182        fn dynamic_slice(input: &Tensor, starts: &Tensor, slice_sizes: &[usize]) -> TensorResult<Tensor>;
183        fn dynamic_update_slice(operand: &Tensor, update: &Tensor, starts: &Tensor) -> TensorResult<Tensor>;
184        fn pad(input: &Tensor, config: &PadConfig) -> TensorResult<Tensor>;
185        fn concatenate(inputs: &[&Tensor], axis: usize) -> TensorResult<Tensor>;
186        fn reverse(input: &Tensor, axes: &[usize]) -> TensorResult<Tensor>;
187    }
188}
189
190impl BackendSessionHost for EagerBackend {
191    fn with_backend_session<R: Send>(
192        &mut self,
193        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
194    ) -> R {
195        dispatch!(self, with_backend_session(f))
196    }
197}
198
199impl TensorDeviceTransfer for EagerBackend {
200    delegate_tensor_backend_methods! {
201        fn download_to_host(tensor: &Tensor) -> TensorResult<Tensor>;
202        fn upload_host_tensor(tensor: &Tensor) -> TensorResult<Tensor>;
203    }
204}
205
206impl TensorBuffer for EagerBackend {
207    delegate_tensor_backend_methods! {
208        fn reclaim_buffer(tensor: Tensor) -> ();
209    }
210}
211
212impl TensorFusion for EagerBackend {
213    delegate_tensor_backend_methods! {
214        fn execute_elementwise_fusion(inputs: &[&Tensor], plan: &ElementwiseFusionPlan) -> TensorResult<Option<Vec<Tensor>>>;
215        fn execute_broadcast_multiply(lhs: TensorRead<'_>, lhs_shape: &[usize], lhs_dims: &[usize], rhs: TensorRead<'_>, rhs_shape: &[usize], rhs_dims: &[usize]) -> TensorResult<Option<Tensor>>;
216        fn execute_broadcast_multiply_value(lhs: TensorRead<'_>, lhs_shape: &[usize], lhs_dims: &[usize], rhs: TensorRead<'_>, rhs_shape: &[usize], rhs_dims: &[usize]) -> TensorResult<Option<TensorValue>>;
217    }
218}
219
220impl BackendCachedDot for EagerBackend {}
221
222impl TensorBackend for EagerBackend {}
tenferro_ad/eager_backend.rs

tenferro_ad/
eager_backend.rs
