tenferro_cpu/
backend.rs

1//! Pool helpers shared by the CPU engine's operations.
2
3use num_complex::{Complex32, Complex64};
4use tenferro_tensor::{DType, ElementwiseReadOp};
5use threadteam::Exec;
6
7use crate::buffer_pool::{BufferPool, PoolScalar};
8use crate::{Tensor, TensorRead, TensorWrite, TypedTensor};
9
10use super::{copy_tensor_read_into, elementwise};
11
12pub(crate) fn elementwise_read_into_fallback_with_pool(
13    buffers: &mut BufferPool,
14    op: ElementwiseReadOp,
15    inputs: &[TensorRead<'_>],
16    out: TensorWrite<'_>,
17    mut exec: Exec<'_>,
18) -> crate::Result<()> {
19    let result = match op {
20        ElementwiseReadOp::Add => elementwise::add_read_with_pool(
21            buffers,
22            inputs[0].clone(),
23            inputs[1].clone(),
24            exec.reborrow(),
25        )?,
26        ElementwiseReadOp::Subtract => elementwise::sub_read_with_pool(
27            buffers,
28            inputs[0].clone(),
29            inputs[1].clone(),
30            exec.reborrow(),
31        )?,
32        ElementwiseReadOp::Multiply => elementwise::mul_read_with_pool(
33            buffers,
34            inputs[0].clone(),
35            inputs[1].clone(),
36            exec.reborrow(),
37        )?,
38        ElementwiseReadOp::Negate => {
39            elementwise::neg_read_with_pool(buffers, inputs[0].clone(), exec.reborrow())?
40        }
41        ElementwiseReadOp::Conj => {
42            elementwise::conj_read_with_pool(buffers, inputs[0].clone(), exec.reborrow())?
43        }
44        ElementwiseReadOp::Divide => elementwise::div_read_with_pool(
45            buffers,
46            inputs[0].clone(),
47            inputs[1].clone(),
48            exec.reborrow(),
49        )?,
50        _ => {
51            return Err(crate::Error::unsupported(
52                "CpuEngine::elementwise_read_into",
53                format!("CPU backend does not implement {op:?}"),
54            ))
55        }
56    };
57    let copied = copy_tensor_read_into(
58        "CpuEngine::elementwise_read_into",
59        TensorRead::from_tensor(&result),
60        out,
61        exec,
62    );
63    // The staged result is scratch: return it to the pool for the next op.
64    reclaim_tensor(buffers, result);
65    copied
66}
67
68pub(crate) fn reclaim_tensor(buffers: &mut BufferPool, tensor: Tensor) {
69    match tensor.dtype() {
70        DType::F32 => reclaim_tensor_typed::<f32>(buffers, tensor),
71        DType::F64 => reclaim_tensor_typed::<f64>(buffers, tensor),
72        DType::I32 => reclaim_tensor_typed::<i32>(buffers, tensor),
73        DType::I64 => reclaim_tensor_typed::<i64>(buffers, tensor),
74        DType::Bool => reclaim_tensor_typed::<bool>(buffers, tensor),
75        DType::C32 => reclaim_tensor_typed::<Complex32>(buffers, tensor),
76        DType::C64 => reclaim_tensor_typed::<Complex64>(buffers, tensor),
77        DType::External(_) => {}
78    }
79}
80
81/// Hand the typed tensor back to the pool when the tag table reached the matching tag.
82fn reclaim_tensor_typed<T: tenferro_cpu_basic::PoolScalar>(
83    buffers: &mut BufferPool,
84    tensor: Tensor,
85) {
86    if let Ok(typed) = tensor.into_typed::<T>() {
87        reclaim_typed(buffers, typed);
88    }
89}
90
91pub(crate) fn reclaim_typed<T: PoolScalar>(pool: &mut BufferPool, typed: TypedTensor<T>) {
92    if typed.backend_buffer().is_some() {
93        return;
94    }
95    if let Ok(data) = typed.into_host_vec() {
96        T::pool_release(pool, data);
97    }
98}
99
100#[cfg(test)]
101mod tests;
tenferro_cpu/backend.rs

tenferro_cpu/
backend.rs
