Skip to main content

candela/tensor/backend/cpu_pure/
mod.rs

1mod f32;
2mod f64;
3mod kernels;
4
5use crate::Layout;
6use crate::tensor::backend::{Backend, ComputeFor, Dtype};
7use crate::tensor::ops::def_op::OpKind;
8use crate::tensor::storage::TensorData;
9
10#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
11pub struct CpuPure;
12
13impl Backend for CpuPure {
14    const SUPPORTS_2D_TRANSPOSED_MATMUL: bool = true;
15    const SUPPORTS_NON_CONTIGUOUS_MATMUL: bool = true;
16
17    fn compute<T>(
18        op: &OpKind<T>,
19        output_buffer: Vec<T>,
20        output_layout: &Layout,
21        inputs: &[TensorData<T>],
22    ) -> TensorData<T>
23    where
24        T: Dtype + ComputeFor<CpuPure>,
25    {
26        T::compute(op, output_buffer, output_layout, inputs)
27    }
28
29    fn compute_inplace<T>(
30        op: &OpKind<T>,
31        output_layout: &Layout,
32        inputs: Vec<TensorData<T>>,
33        output_idx: usize,
34    ) -> TensorData<T>
35    where
36        T: Dtype + ComputeFor<Self>,
37    {
38        T::compute_inplace(op, output_layout, inputs, output_idx)
39    }
40}
41
42impl ComputeFor<CpuPure> for f64 {
43    fn compute(
44        op: &OpKind<f64>,
45        output_buffer: Vec<f64>,
46        output_layout: &Layout,
47        inputs: &[TensorData<f64>],
48    ) -> TensorData<f64> {
49        f64::compute_op(op, output_buffer, output_layout, inputs)
50    }
51
52    fn compute_inplace(
53        op: &OpKind<Self>,
54        output_layout: &Layout,
55        inputs: Vec<TensorData<Self>>,
56        output_idx: usize,
57    ) -> TensorData<Self> {
58        f64::compute_op_inplace(op, output_layout, inputs, output_idx)
59    }
60}
61
62impl ComputeFor<CpuPure> for f32 {
63    fn compute(
64        op: &OpKind<f32>,
65        output_buffer: Vec<f32>,
66        output_layout: &Layout,
67        inputs: &[TensorData<f32>],
68    ) -> TensorData<f32> {
69        f32::compute_op(op, output_buffer, output_layout, inputs)
70    }
71
72    fn compute_inplace(
73        op: &OpKind<Self>,
74        output_layout: &Layout,
75        inputs: Vec<TensorData<Self>>,
76        output_idx: usize,
77    ) -> TensorData<Self> {
78        f32::compute_op_inplace(op, output_layout, inputs, output_idx)
79    }
80}