tenferro-cpu 0.4.0

CPU backend, kernels, provider selection, and CPU resource pools for tenferro.
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};

use crate::buffer_pool::BufferPool;
use crate::gemm::GemmAnalysisCache;
use crate::indexed_plan_cache::IndexedPlanCache;
use crate::resource_domain::CpuResourceDomain;
use crate::{
    CpuContext, CpuContextError, CpuDomainExecutor, CpuDomainId, CpuDomainOwnership,
    CpuExecutorAffinity, CpuPlacementGuarantee, ExternalCpuDomain, ResolvedCpuPlacement,
};

#[derive(Debug)]
pub(crate) struct EngineResources {
    pub(crate) buffers: BufferPool,
    pub(crate) gemm_analysis_cache: GemmAnalysisCache,
    pub(crate) indexed_plan_cache: IndexedPlanCache,
}

impl EngineResources {
    pub(crate) fn new(buffer_limit: usize) -> Self {
        Self {
            buffers: BufferPool::with_max_retained_capacity_bytes(buffer_limit),
            gemm_analysis_cache: GemmAnalysisCache::default(),
            indexed_plan_cache: IndexedPlanCache::default(),
        }
    }
}

#[derive(Debug)]
pub(crate) struct CpuEngine {
    domain: CpuResourceDomain,
    pub(crate) resources: Mutex<EngineResources>,
}

impl CpuEngine {
    pub(crate) fn new_managed(
        id: CpuDomainId,
        placement: ResolvedCpuPlacement,
        thread_budget: usize,
        buffer_limit: usize,
    ) -> Result<Self, CpuContextError> {
        let worker_count = thread_budget.min(placement.cpus().len());
        let thread_budget =
            NonZeroUsize::new(worker_count).ok_or(CpuContextError::InvalidThreadCount)?;
        let context = CpuContext::with_pinned_cpus(placement.cpus().clone(), worker_count)?;
        Ok(Self::from_managed_context(
            id,
            placement,
            Arc::new(context),
            thread_budget,
            CpuPlacementGuarantee::ExactDeclared,
            buffer_limit,
        ))
    }

    pub(crate) fn from_context(
        id: CpuDomainId,
        placement: ResolvedCpuPlacement,
        context: Arc<CpuContext>,
        buffer_limit: usize,
    ) -> Self {
        let capabilities = context.capabilities();
        let placement_guarantee =
            if capabilities.affinity == CpuExecutorAffinity::TenferroPinnedVerified {
                CpuPlacementGuarantee::ExactDeclared
            } else {
                CpuPlacementGuarantee::AdvisoryDeclared
            };
        Self::from_managed_context(
            id,
            placement,
            context,
            capabilities.worker_count,
            placement_guarantee,
            buffer_limit,
        )
    }

    fn from_managed_context(
        id: CpuDomainId,
        placement: ResolvedCpuPlacement,
        context: Arc<CpuContext>,
        thread_budget: NonZeroUsize,
        placement_guarantee: CpuPlacementGuarantee,
        buffer_limit: usize,
    ) -> Self {
        let executor: Arc<dyn CpuDomainExecutor> = context.clone();
        Self {
            domain: CpuResourceDomain::new(
                id,
                placement,
                executor,
                thread_budget,
                placement_guarantee,
                CpuDomainOwnership::Managed,
            ),
            resources: Mutex::new(EngineResources::new(buffer_limit)),
        }
    }

    pub(crate) fn from_external(domain: ExternalCpuDomain, buffer_limit: usize) -> Self {
        Self {
            domain: domain.into(),
            resources: Mutex::new(EngineResources::new(buffer_limit)),
        }
    }

    pub(crate) fn domain(&self) -> &CpuResourceDomain {
        &self.domain
    }

    pub(crate) fn placement(&self) -> Option<&ResolvedCpuPlacement> {
        self.domain.placement()
    }
}

#[cfg(test)]
mod tests;