tenferro-cpu 0.8.0

CPU backend, kernels, provider selection, and CPU resource pools for tenferro.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

//! Pool helpers shared by the CPU engine's operations.

use num_complex::{Complex32, Complex64};
use tenferro_tensor::{DType, ElementwiseReadOp};
use threadteam::Exec;

use crate::buffer_pool::{BufferPool, PoolScalar};
use crate::{Tensor, TensorRead, TensorWrite, TypedTensor};

use super::{copy_tensor_read_into, elementwise};

pub(crate) fn elementwise_read_into_fallback_with_pool(
    buffers: &mut BufferPool,
    op: ElementwiseReadOp,
    inputs: &[TensorRead<'_>],
    out: TensorWrite<'_>,
    mut exec: Exec<'_>,
) -> crate::Result<()> {
    let result = match op {
        ElementwiseReadOp::Add => elementwise::add_read_with_pool(
            buffers,
            inputs[0].clone(),
            inputs[1].clone(),
            exec.reborrow(),
        )?,
        ElementwiseReadOp::Subtract => elementwise::sub_read_with_pool(
            buffers,
            inputs[0].clone(),
            inputs[1].clone(),
            exec.reborrow(),
        )?,
        ElementwiseReadOp::Multiply => elementwise::mul_read_with_pool(
            buffers,
            inputs[0].clone(),
            inputs[1].clone(),
            exec.reborrow(),
        )?,
        ElementwiseReadOp::Negate => {
            elementwise::neg_read_with_pool(buffers, inputs[0].clone(), exec.reborrow())?
        }
        ElementwiseReadOp::Conj => {
            elementwise::conj_read_with_pool(buffers, inputs[0].clone(), exec.reborrow())?
        }
        ElementwiseReadOp::Divide => elementwise::div_read_with_pool(
            buffers,
            inputs[0].clone(),
            inputs[1].clone(),
            exec.reborrow(),
        )?,
        _ => {
            return Err(crate::Error::unsupported(
                "CpuEngine::elementwise_read_into",
                format!("CPU backend does not implement {op:?}"),
            ))
        }
    };
    let copied = copy_tensor_read_into(
        "CpuEngine::elementwise_read_into",
        TensorRead::from_tensor(&result),
        out,
        exec,
    );
    // The staged result is scratch: return it to the pool for the next op.
    reclaim_tensor(buffers, result);
    copied
}

pub(crate) fn reclaim_tensor(buffers: &mut BufferPool, tensor: Tensor) {
    match tensor.dtype() {
        DType::F32 => reclaim_tensor_typed::<f32>(buffers, tensor),
        DType::F64 => reclaim_tensor_typed::<f64>(buffers, tensor),
        DType::I32 => reclaim_tensor_typed::<i32>(buffers, tensor),
        DType::I64 => reclaim_tensor_typed::<i64>(buffers, tensor),
        DType::Bool => reclaim_tensor_typed::<bool>(buffers, tensor),
        DType::C32 => reclaim_tensor_typed::<Complex32>(buffers, tensor),
        DType::C64 => reclaim_tensor_typed::<Complex64>(buffers, tensor),
        DType::External(_) => {}
    }
}

/// Hand the typed tensor back to the pool when the tag table reached the matching tag.
fn reclaim_tensor_typed<T: tenferro_cpu_basic::PoolScalar>(
    buffers: &mut BufferPool,
    tensor: Tensor,
) {
    if let Ok(typed) = tensor.into_typed::<T>() {
        reclaim_typed(buffers, typed);
    }
}

pub(crate) fn reclaim_typed<T: PoolScalar>(pool: &mut BufferPool, typed: TypedTensor<T>) {
    if typed.backend_buffer().is_some() {
        return;
    }
    if let Ok(data) = typed.into_host_vec() {
        T::pool_release(pool, data);
    }
}

#[cfg(test)]
mod tests;