thread_aware/
affinity.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Identifiers for threads and NUMA regions.
5
6#![allow(clippy::allow_attributes, reason = "Needed for conditional compilation")]
7
8/// An `MemoryAffinity` can be thought of as a placement in a system.
9///
10/// It is used to represent a specific context or environment where data can be processed.
11/// For example a NUMA node, a thread, a specific CPU core, or a specific memory region.
12#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
13pub enum MemoryAffinity {
14    /// An unknown affinity represents no specific binding, like an unpinned thread.
15    Unknown,
16    /// A pinned affinity represents a specific binding to a processor and memory region.
17    Pinned(PinnedAffinity),
18}
19
20impl From<PinnedAffinity> for MemoryAffinity {
21    fn from(pinned: PinnedAffinity) -> Self {
22        Self::Pinned(pinned)
23    }
24}
25
26impl MemoryAffinity {
27    /// Returns an unknown affinity.
28    #[must_use]
29    pub const fn unknown() -> Self {
30        Self::Unknown
31    }
32
33    /// Returns `true` if the affinity is unknown.
34    #[must_use]
35    pub const fn is_unknown(self) -> bool {
36        matches!(self, Self::Unknown)
37    }
38}
39
40/// A `PinnedAffinity` represents a specific binding to a processor and memory region.
41#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
42pub struct PinnedAffinity {
43    processor_index: u16,
44    memory_region_index: u16,
45
46    processor_count: u16,
47    memory_region_count: u16,
48}
49
50impl PinnedAffinity {
51    // You should probably use the `ThreadRegistry` or `create_manual_pinned_affinities` to create these.
52    #[must_use]
53    pub(crate) const fn new(processor_index: u16, memory_region_index: u16, processor_count: u16, memory_region_count: u16) -> Self {
54        Self {
55            processor_index,
56            memory_region_index,
57            processor_count,
58            memory_region_count,
59        }
60    }
61
62    /// Returns the processor index of this affinity.
63    #[must_use]
64    pub const fn processor_index(self) -> usize {
65        self.processor_index as _
66    }
67
68    /// Returns the memory region index of this affinity.
69    #[must_use]
70    pub const fn memory_region_index(self) -> usize {
71        self.memory_region_index as _
72    }
73
74    /// Returns the processor count of this affinity.
75    #[must_use]
76    pub const fn processor_count(self) -> usize {
77        self.processor_count as _
78    }
79
80    /// Returns the number of memory regions of this affinity.
81    #[must_use]
82    pub const fn memory_region_count(self) -> usize {
83        self.memory_region_count as _
84    }
85}
86
87/// Create pinned affinities manually when not using the `ThreadRegistry`.
88///
89/// # Parameters
90///
91/// * `counts`: A slice of `usize` representing the number of processors in each memory region.
92///
93/// # Panics
94///
95/// If there are more than `u16::MAX` processors or memory regions.
96#[must_use]
97#[expect(clippy::needless_range_loop, reason = "clearer in this case")]
98pub fn pinned_affinities(counts: &[usize]) -> Vec<PinnedAffinity> {
99    let numa_count = counts.len();
100    let core_count = counts.iter().sum();
101    let mut affinities = Vec::with_capacity(core_count);
102    let mut processor_index = 0;
103
104    for numa_index in 0..numa_count {
105        for _ in 0..counts[numa_index] {
106            affinities.push(PinnedAffinity::new(
107                processor_index.try_into().expect("Too many processors"),
108                numa_index.try_into().expect("Too many memory regions"),
109                core_count.try_into().expect("Too many processors"),
110                numa_count.try_into().expect("Too many memory regions"),
111            ));
112            processor_index += 1;
113        }
114    }
115
116    affinities
117}
118
119/// Create memory affinities manually when not using the `ThreadRegistry`.
120///
121/// This is similar to `create_manual_pinned_affinities` but returns `MemoryAffinity` values.
122///
123/// # Parameters
124///
125/// * `counts`: A slice of `usize` representing the number of processors in each memory region.
126///
127/// # Panics
128///
129/// If there are more than `u16::MAX` processors or memory regions.
130#[must_use]
131pub fn memory_affinities(counts: &[usize]) -> Vec<MemoryAffinity> {
132    pinned_affinities(counts).into_iter().map(MemoryAffinity::Pinned).collect()
133}
134
135#[cfg(test)]
136mod tests {
137    use super::{MemoryAffinity, PinnedAffinity};
138
139    #[test]
140    fn test_pinned_affinity() {
141        let affinity = PinnedAffinity::new(2, 1, 4, 2);
142
143        assert_eq!(affinity.processor_index(), 2);
144        assert_eq!(affinity.processor_count(), 4);
145
146        assert_eq!(affinity.memory_region_index(), 1);
147        assert_eq!(affinity.memory_region_count(), 2);
148    }
149
150    #[test]
151    fn test_memory_affinity_unknown() {
152        let affinity = MemoryAffinity::unknown();
153        assert!(affinity.is_unknown());
154    }
155
156    #[test]
157    fn test_memory_affinity_pinned() {
158        let affinity = PinnedAffinity::new(2, 1, 4, 2);
159        let affinity = MemoryAffinity::from(affinity);
160        assert!(!affinity.is_unknown());
161    }
162}
thread_aware/affinity.rs

thread_aware/
affinity.rs
