tenferro_cpu/
affinity_policy.rs

1use std::collections::BTreeMap;
2
3use smallvec::SmallVec;
4use tenferro_tensor::{CpuDomainId, DType, Tensor};
5
6const INLINE_DOMAIN_CAPACITY: usize = 8;
7
8/// Policy used to select a CPU execution domain from input affinity metadata.
9///
10/// # Examples
11///
12/// ```rust
13/// use tenferro_cpu::CpuAffinityPolicy;
14///
15/// let policy = CpuAffinityPolicy::DominantInputBytes;
16/// assert_ne!(policy, CpuAffinityPolicy::RequireSingleDomain);
17/// ```
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
19pub enum CpuAffinityPolicy {
20    /// Select the domain with the largest total of positive logical input bytes.
21    DominantInputBytes,
22    /// Accept zero or one known input domain and reject mixed known domains.
23    RequireSingleDomain,
24}
25
26/// CPU affinity metadata for one logical operation input.
27///
28/// The resolver reads this metadata only. It never changes, copies, or rehomes
29/// tensor payloads.
30///
31/// # Examples
32///
33/// ```rust
34/// use tenferro_cpu::CpuAffinityInput;
35/// use tenferro_tensor::CpuDomainId;
36///
37/// let input = CpuAffinityInput {
38///     domain: Some(CpuDomainId::new(3)),
39///     logical_bytes: 64,
40/// };
41/// assert_eq!(input.logical_bytes, 64);
42/// ```
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub struct CpuAffinityInput {
45    /// Known CPU execution domain, or `None` when affinity is unknown.
46    pub domain: Option<CpuDomainId>,
47    /// Logical input size used by [`CpuAffinityPolicy::DominantInputBytes`].
48    pub logical_bytes: usize,
49}
50
51impl CpuAffinityInput {
52    /// Construct resolver input metadata from a tensor.
53    ///
54    /// The logical byte count is the checked shape product times the tensor's
55    /// scalar byte width. CPU affinity is copied from placement metadata; the
56    /// tensor and its storage are otherwise untouched.
57    ///
58    /// # Examples
59    ///
60    /// ```rust
61    /// use tenferro_cpu::CpuAffinityInput;
62    /// use tenferro_tensor::Tensor;
63    ///
64    /// let tensor = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
65    /// let input = CpuAffinityInput::from_tensor(&tensor)?;
66    /// assert_eq!(input.logical_bytes, 2 * std::mem::size_of::<f64>());
67    /// # Ok::<(), Box<dyn std::error::Error>>(())
68    /// ```
69    ///
70    /// # Errors
71    ///
72    /// Returns [`CpuAffinityInputError`] when the logical element or byte
73    /// count cannot be represented by `usize`.
74    pub fn from_tensor(tensor: &Tensor) -> Result<Self, CpuAffinityInputError> {
75        Self::from_parts(
76            tensor.placement().cpu_affinity,
77            tensor.shape(),
78            tensor.dtype(),
79        )
80    }
81
82    /// Construct resolver input metadata from placement, shape, and dtype.
83    ///
84    /// Scalar shapes have one element. Any zero extent yields zero logical
85    /// bytes without multiplying the other extents.
86    ///
87    /// # Examples
88    ///
89    /// ```rust
90    /// use tenferro_cpu::CpuAffinityInput;
91    /// use tenferro_tensor::DType;
92    ///
93    /// let scalar = CpuAffinityInput::from_parts(None, &[], DType::F32)?;
94    /// let empty = CpuAffinityInput::from_parts(None, &[usize::MAX, 0], DType::F64)?;
95    /// assert_eq!(scalar.logical_bytes, 4);
96    /// assert_eq!(empty.logical_bytes, 0);
97    /// # Ok::<(), tenferro_cpu::CpuAffinityInputError>(())
98    /// ```
99    ///
100    /// # Errors
101    ///
102    /// Returns [`CpuAffinityInputError::ShapeProductOverflow`] when non-zero
103    /// extents overflow, or
104    /// [`CpuAffinityInputError::LogicalByteCountOverflow`] when multiplying by
105    /// the dtype width overflows.
106    pub fn from_parts(
107        domain: Option<CpuDomainId>,
108        shape: &[usize],
109        dtype: DType,
110    ) -> Result<Self, CpuAffinityInputError> {
111        let element_count = if shape.contains(&0) {
112            0
113        } else {
114            shape.iter().try_fold(1_usize, |count, &extent| {
115                count
116                    .checked_mul(extent)
117                    .ok_or(CpuAffinityInputError::ShapeProductOverflow)
118            })?
119        };
120        let byte_width = dtype_byte_width(dtype);
121        let logical_bytes = element_count.checked_mul(byte_width).ok_or(
122            CpuAffinityInputError::LogicalByteCountOverflow {
123                element_count,
124                byte_width,
125            },
126        )?;
127        Ok(Self {
128            domain,
129            logical_bytes,
130        })
131    }
132}
133
134const fn dtype_byte_width(dtype: DType) -> usize {
135    match dtype {
136        DType::F32 | DType::I32 => std::mem::size_of::<u32>(),
137        DType::F64 | DType::I64 => std::mem::size_of::<u64>(),
138        DType::Bool => std::mem::size_of::<bool>(),
139        DType::C32 => std::mem::size_of::<num_complex::Complex32>(),
140        DType::C64 => std::mem::size_of::<num_complex::Complex64>(),
141    }
142}
143
144/// Failure to derive logical input bytes for CPU affinity resolution.
145///
146/// # Examples
147///
148/// ```rust
149/// use tenferro_cpu::{CpuAffinityInput, CpuAffinityInputError};
150/// use tenferro_tensor::DType;
151///
152/// let error = CpuAffinityInput::from_parts(None, &[usize::MAX, 2], DType::F32)
153///     .unwrap_err();
154/// assert_eq!(error, CpuAffinityInputError::ShapeProductOverflow);
155/// ```
156#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
157pub enum CpuAffinityInputError {
158    /// Multiplying non-zero shape extents overflowed `usize`.
159    #[error("logical tensor element count overflowed usize")]
160    ShapeProductOverflow,
161    /// Multiplying element count by dtype width overflowed `usize`.
162    #[error(
163        "logical tensor byte count overflowed: element_count={element_count}, byte_width={byte_width}"
164    )]
165    LogicalByteCountOverflow {
166        /// Checked logical element count.
167        element_count: usize,
168        /// Scalar dtype width in bytes.
169        byte_width: usize,
170    },
171}
172
173/// Why the CPU affinity resolver selected a domain.
174///
175/// # Examples
176///
177/// ```rust
178/// use tenferro_cpu::CpuAffinitySelectionReason;
179///
180/// let reason = CpuAffinitySelectionReason::DefaultDomain;
181/// assert_eq!(reason, CpuAffinitySelectionReason::DefaultDomain);
182/// ```
183#[derive(Clone, Copy, Debug, Eq, PartialEq)]
184pub enum CpuAffinitySelectionReason {
185    /// An operation-local explicit domain override took precedence.
186    ExplicitOverride,
187    /// Positive logical bytes made this domain dominant.
188    DominantInputBytes,
189    /// Strict policy observed exactly one known input domain.
190    SingleInputDomain,
191    /// No relevant input affinity was available.
192    DefaultDomain,
193}
194
195/// Deterministic CPU affinity selection returned by the pure resolver.
196///
197/// # Examples
198///
199/// ```rust
200/// use tenferro_cpu::{CpuAffinitySelection, CpuAffinitySelectionReason};
201/// use tenferro_tensor::CpuDomainId;
202///
203/// let selection = CpuAffinitySelection {
204///     domain: CpuDomainId::new(2),
205///     reason: CpuAffinitySelectionReason::DominantInputBytes,
206/// };
207/// assert_eq!(selection.domain.as_u64(), 2);
208/// ```
209#[derive(Clone, Copy, Debug, Eq, PartialEq)]
210pub struct CpuAffinitySelection {
211    /// Selected CPU execution domain.
212    pub domain: CpuDomainId,
213    /// Deterministic reason for the selection.
214    pub reason: CpuAffinitySelectionReason,
215}
216
217/// Failure to resolve CPU affinity from input metadata.
218///
219/// # Examples
220///
221/// ```rust
222/// use tenferro_cpu::CpuAffinityResolutionError;
223/// use tenferro_tensor::CpuDomainId;
224///
225/// let error = CpuAffinityResolutionError::LogicalByteCountOverflow {
226///     domain: CpuDomainId::new(4),
227/// };
228/// assert!(error.to_string().contains("4"));
229/// ```
230#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
231pub enum CpuAffinityResolutionError {
232    /// Adding logical byte counts overflowed `usize` for one domain.
233    #[error("logical input-byte total overflowed for CPU domain {domain:?}")]
234    LogicalByteCountOverflow {
235        /// Smallest CPU domain whose logical byte total overflowed.
236        domain: CpuDomainId,
237    },
238    /// Strict policy observed at least two different known domains.
239    #[error("CPU affinity policy requires one input domain, found {first:?} and {second:?}")]
240    MultipleKnownDomains {
241        /// Smallest known input domain.
242        first: CpuDomainId,
243        /// Second-smallest known input domain.
244        second: CpuDomainId,
245    },
246}
247
248/// Resolve a CPU execution domain from input affinity metadata.
249///
250/// Unknown affinities and zero-byte inputs do not contribute to dominant-byte
251/// scoring. When no input contributes, `default_domain` is selected. Equal
252/// positive totals are resolved in favor of the smallest [`CpuDomainId`]. The
253/// input slice is only read; the resolver never retags or rehomes an input.
254///
255/// Use [`resolve_cpu_affinity_with_override`] when an operation-local explicit
256/// placement has already been selected.
257///
258/// # Examples
259///
260/// ```rust
261/// use tenferro_cpu::{resolve_cpu_affinity, CpuAffinityInput, CpuAffinityPolicy};
262/// use tenferro_tensor::CpuDomainId;
263///
264/// let inputs = [
265///     CpuAffinityInput { domain: Some(CpuDomainId::new(8)), logical_bytes: 6 },
266///     CpuAffinityInput { domain: Some(CpuDomainId::new(3)), logical_bytes: 2 },
267/// ];
268/// let selected = resolve_cpu_affinity(
269///     CpuAffinityPolicy::DominantInputBytes,
270///     &inputs,
271///     CpuDomainId::new(1),
272/// )?;
273/// assert_eq!(selected.domain, CpuDomainId::new(8));
274/// # Ok::<(), tenferro_cpu::CpuAffinityResolutionError>(())
275/// ```
276///
277/// # Errors
278///
279/// Returns [`CpuAffinityResolutionError::LogicalByteCountOverflow`] when one
280/// domain's logical byte total cannot be represented by `usize`, or
281/// [`CpuAffinityResolutionError::MultipleKnownDomains`] when strict policy sees
282/// more than one known input domain.
283pub fn resolve_cpu_affinity(
284    policy: CpuAffinityPolicy,
285    inputs: &[CpuAffinityInput],
286    default_domain: CpuDomainId,
287) -> Result<CpuAffinitySelection, CpuAffinityResolutionError> {
288    resolve_cpu_affinity_with_override(policy, inputs, default_domain, None)
289}
290
291/// Resolve CPU affinity with an optional operation-local explicit override.
292///
293/// Explicit placement takes precedence before input-byte accounting or strict
294/// mixed-domain validation. Passing `None` applies the same policy resolution
295/// as [`resolve_cpu_affinity`].
296///
297/// # Examples
298///
299/// ```rust
300/// use tenferro_cpu::{
301///     resolve_cpu_affinity_with_override, CpuAffinityInput, CpuAffinityPolicy,
302///     CpuAffinitySelectionReason,
303/// };
304/// use tenferro_tensor::CpuDomainId;
305///
306/// let mixed = [
307///     CpuAffinityInput { domain: Some(CpuDomainId::new(1)), logical_bytes: 1 },
308///     CpuAffinityInput { domain: Some(CpuDomainId::new(2)), logical_bytes: 1 },
309/// ];
310/// let selected = resolve_cpu_affinity_with_override(
311///     CpuAffinityPolicy::RequireSingleDomain,
312///     &mixed,
313///     CpuDomainId::new(1),
314///     Some(CpuDomainId::new(9)),
315/// )?;
316/// assert_eq!(selected.domain, CpuDomainId::new(9));
317/// assert_eq!(selected.reason, CpuAffinitySelectionReason::ExplicitOverride);
318/// # Ok::<(), tenferro_cpu::CpuAffinityResolutionError>(())
319/// ```
320///
321/// # Errors
322///
323/// When `explicit_domain` is `None`, returns
324/// [`CpuAffinityResolutionError::LogicalByteCountOverflow`] for an unrepresentable
325/// domain byte total or [`CpuAffinityResolutionError::MultipleKnownDomains`]
326/// when strict policy sees more than one known input domain. A present explicit
327/// override bypasses both policy errors.
328pub fn resolve_cpu_affinity_with_override(
329    policy: CpuAffinityPolicy,
330    inputs: &[CpuAffinityInput],
331    default_domain: CpuDomainId,
332    explicit_domain: Option<CpuDomainId>,
333) -> Result<CpuAffinitySelection, CpuAffinityResolutionError> {
334    if let Some(domain) = explicit_domain {
335        return Ok(CpuAffinitySelection {
336            domain,
337            reason: CpuAffinitySelectionReason::ExplicitOverride,
338        });
339    }
340    match policy {
341        CpuAffinityPolicy::DominantInputBytes => resolve_dominant(inputs, default_domain),
342        CpuAffinityPolicy::RequireSingleDomain => resolve_single(inputs, default_domain),
343    }
344}
345
346fn resolve_dominant(
347    inputs: &[CpuAffinityInput],
348    default_domain: CpuDomainId,
349) -> Result<CpuAffinitySelection, CpuAffinityResolutionError> {
350    let mut totals = DomainTotals::default();
351    for input in inputs {
352        let Some(domain) = input.domain else {
353            continue;
354        };
355        if input.logical_bytes == 0 {
356            continue;
357        }
358        totals.add(domain, input.logical_bytes);
359    }
360
361    if let Some(domain) = totals.smallest_overflowing_domain() {
362        return Err(CpuAffinityResolutionError::LogicalByteCountOverflow { domain });
363    }
364
365    match totals.dominant_domain() {
366        Some(domain) => Ok(CpuAffinitySelection {
367            domain,
368            reason: CpuAffinitySelectionReason::DominantInputBytes,
369        }),
370        None => Ok(default_selection(default_domain)),
371    }
372}
373
374fn resolve_single(
375    inputs: &[CpuAffinityInput],
376    default_domain: CpuDomainId,
377) -> Result<CpuAffinitySelection, CpuAffinityResolutionError> {
378    let mut first = None;
379    let mut second = None;
380    for domain in inputs.iter().filter_map(|input| input.domain) {
381        observe_smallest_two_distinct(domain, &mut first, &mut second);
382    }
383
384    if let (Some(first), Some(second)) = (first, second) {
385        return Err(CpuAffinityResolutionError::MultipleKnownDomains { first, second });
386    }
387
388    Ok(match first {
389        Some(domain) => CpuAffinitySelection {
390            domain,
391            reason: CpuAffinitySelectionReason::SingleInputDomain,
392        },
393        None => default_selection(default_domain),
394    })
395}
396
397fn observe_smallest_two_distinct(
398    domain: CpuDomainId,
399    first: &mut Option<CpuDomainId>,
400    second: &mut Option<CpuDomainId>,
401) {
402    if *first == Some(domain) || *second == Some(domain) {
403        return;
404    }
405    match *first {
406        None => *first = Some(domain),
407        Some(current_first) if domain < current_first => {
408            *second = *first;
409            *first = Some(domain);
410        }
411        Some(_) if second.is_none_or(|current_second| domain < current_second) => {
412            *second = Some(domain);
413        }
414        Some(_) => {}
415    }
416}
417
418fn default_selection(domain: CpuDomainId) -> CpuAffinitySelection {
419    CpuAffinitySelection {
420        domain,
421        reason: CpuAffinitySelectionReason::DefaultDomain,
422    }
423}
424
425#[derive(Clone, Copy, Debug)]
426struct DomainTotal {
427    domain: CpuDomainId,
428    logical_bytes: Option<usize>,
429}
430
431impl DomainTotal {
432    fn new(domain: CpuDomainId, logical_bytes: usize) -> Self {
433        Self {
434            domain,
435            logical_bytes: Some(logical_bytes),
436        }
437    }
438
439    fn add(&mut self, logical_bytes: usize) {
440        self.logical_bytes = self
441            .logical_bytes
442            .and_then(|total| total.checked_add(logical_bytes));
443    }
444}
445
446enum DomainTotals {
447    Inline(SmallVec<[DomainTotal; INLINE_DOMAIN_CAPACITY]>),
448    Heap(BTreeMap<CpuDomainId, Option<usize>>),
449}
450
451impl Default for DomainTotals {
452    fn default() -> Self {
453        Self::Inline(SmallVec::new())
454    }
455}
456
457impl DomainTotals {
458    fn add(&mut self, domain: CpuDomainId, logical_bytes: usize) {
459        let promoted = match self {
460            Self::Inline(entries) => {
461                // INVARIANT: the linear lookup is bounded by the inline capacity;
462                // larger distinct-domain sets are promoted to `BTreeMap` below.
463                if let Some(entry) = entries.iter_mut().find(|entry| entry.domain == domain) {
464                    entry.add(logical_bytes);
465                    return;
466                }
467                if entries.len() < INLINE_DOMAIN_CAPACITY {
468                    entries.push(DomainTotal::new(domain, logical_bytes));
469                    return;
470                }
471                let mut heap = BTreeMap::new();
472                for entry in entries.drain(..) {
473                    heap.insert(entry.domain, entry.logical_bytes);
474                }
475                heap.insert(domain, Some(logical_bytes));
476                Some(heap)
477            }
478            Self::Heap(entries) => {
479                let total = entries.entry(domain).or_insert(Some(0));
480                *total = total.and_then(|current| current.checked_add(logical_bytes));
481                None
482            }
483        };
484        if let Some(heap) = promoted {
485            *self = Self::Heap(heap);
486        }
487    }
488
489    fn smallest_overflowing_domain(&self) -> Option<CpuDomainId> {
490        match self {
491            Self::Inline(entries) => entries
492                .iter()
493                .filter(|entry| entry.logical_bytes.is_none())
494                .map(|entry| entry.domain)
495                .min(),
496            Self::Heap(entries) => entries
497                .iter()
498                .find_map(|(domain, total)| total.is_none().then_some(*domain)),
499        }
500    }
501
502    fn dominant_domain(&self) -> Option<CpuDomainId> {
503        let mut best = None;
504        match self {
505            Self::Inline(entries) => {
506                for entry in entries {
507                    if let Some(logical_bytes) = entry.logical_bytes {
508                        consider_dominant(&mut best, entry.domain, logical_bytes);
509                    }
510                }
511            }
512            Self::Heap(entries) => {
513                for (&domain, &logical_bytes) in entries {
514                    if let Some(logical_bytes) = logical_bytes {
515                        consider_dominant(&mut best, domain, logical_bytes);
516                    }
517                }
518            }
519        }
520        best.map(|(domain, _)| domain)
521    }
522}
523
524fn consider_dominant(
525    best: &mut Option<(CpuDomainId, usize)>,
526    domain: CpuDomainId,
527    logical_bytes: usize,
528) {
529    let replace = match *best {
530        None => true,
531        Some((best_domain, best_bytes)) => {
532            logical_bytes > best_bytes || (logical_bytes == best_bytes && domain < best_domain)
533        }
534    };
535    if replace {
536        *best = Some((domain, logical_bytes));
537    }
538}
539
540#[cfg(test)]
541mod tests;
tenferro_cpu/affinity_policy.rs

tenferro_cpu/
affinity_policy.rs
