Skip to main content

moirai_executor/schedule/route/
ids.rs

1//! Route identifier and count newtypes.
2
3/// Scheduler worker-thread identifier.
4#[repr(transparent)]
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub struct ThreadId(usize);
7
8impl ThreadId {
9    /// Construct a thread identifier.
10    #[inline]
11    pub const fn new(id: usize) -> Self {
12        Self(id)
13    }
14
15    /// Return the underlying zero-based identifier.
16    #[inline]
17    pub const fn get(self) -> usize {
18        self.0
19    }
20}
21
22/// Scheduler process identifier.
23#[repr(transparent)]
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
25pub struct ProcessId(usize);
26
27impl ProcessId {
28    /// Construct a process identifier.
29    #[inline]
30    pub const fn new(id: usize) -> Self {
31        Self(id)
32    }
33
34    /// Return the underlying zero-based identifier.
35    #[inline]
36    pub const fn get(self) -> usize {
37        self.0
38    }
39}
40
41/// Scheduler server identifier.
42#[repr(transparent)]
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44pub struct ServerId(usize);
45
46impl ServerId {
47    /// Construct a server identifier.
48    #[inline]
49    pub const fn new(id: usize) -> Self {
50        Self(id)
51    }
52
53    /// Return the underlying zero-based identifier.
54    #[inline]
55    pub const fn get(self) -> usize {
56        self.0
57    }
58}
59
60/// Async lane identifier inside a routed process.
61#[repr(transparent)]
62#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
63pub struct AsyncLaneId(usize);
64
65impl AsyncLaneId {
66    /// Construct an async lane identifier.
67    #[inline]
68    pub const fn new(id: usize) -> Self {
69        Self(id)
70    }
71
72    /// Return the underlying zero-based identifier.
73    #[inline]
74    pub const fn get(self) -> usize {
75        self.0
76    }
77}
78
79/// Accelerator device identifier within one accelerator kind.
80#[repr(transparent)]
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
82pub struct AcceleratorId(usize);
83
84impl AcceleratorId {
85    /// Construct an accelerator identifier.
86    #[inline]
87    pub const fn new(id: usize) -> Self {
88        Self(id)
89    }
90
91    /// Return the underlying zero-based identifier.
92    #[inline]
93    pub const fn get(self) -> usize {
94        self.0
95    }
96}
97
98/// Heterogeneous compute target family.
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
100pub enum AcceleratorKind {
101    /// CPU device placement metadata.
102    Cpu,
103    /// GPU device placement metadata.
104    Gpu,
105    /// TPU device placement metadata.
106    Tpu,
107    /// NPU device placement metadata.
108    Npu,
109}
110
111impl AcceleratorKind {
112    #[inline]
113    pub(crate) const fn checksum_tag(self) -> usize {
114        match self {
115            Self::Cpu => 1,
116            Self::Gpu => 2,
117            Self::Tpu => 3,
118            Self::Npu => 4,
119        }
120    }
121}
122
123/// Number of local scheduler worker threads.
124#[repr(transparent)]
125#[derive(Debug, Clone, Copy, PartialEq, Eq)]
126pub struct WorkerCount(usize);
127
128impl WorkerCount {
129    /// Construct a non-zero worker count.
130    #[inline]
131    pub const fn new(count: usize) -> Self {
132        Self(if count == 0 { 1 } else { count })
133    }
134
135    /// Return the normalized count.
136    #[inline]
137    pub const fn get(self) -> usize {
138        self.0
139    }
140}
141
142/// Number of process route targets.
143#[repr(transparent)]
144#[derive(Debug, Clone, Copy, PartialEq, Eq)]
145pub struct ProcessCount(usize);
146
147impl ProcessCount {
148    /// Construct a non-zero process count.
149    #[inline]
150    pub const fn new(count: usize) -> Self {
151        Self(if count == 0 { 1 } else { count })
152    }
153
154    /// Return the normalized count.
155    #[inline]
156    pub const fn get(self) -> usize {
157        self.0
158    }
159}
160
161/// Number of async lanes in each process target.
162#[repr(transparent)]
163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
164pub struct AsyncLanesPerProcess(usize);
165
166impl AsyncLanesPerProcess {
167    /// Construct a non-zero async-lane count.
168    #[inline]
169    pub const fn new(count: usize) -> Self {
170        Self(if count == 0 { 1 } else { count })
171    }
172
173    /// Return the normalized count.
174    #[inline]
175    pub const fn get(self) -> usize {
176        self.0
177    }
178}
179
180/// Number of server route targets.
181#[repr(transparent)]
182#[derive(Debug, Clone, Copy, PartialEq, Eq)]
183pub struct ServerCount(usize);
184
185impl ServerCount {
186    /// Construct a server count. Zero disables server route targets.
187    #[inline]
188    pub const fn new(count: usize) -> Self {
189        Self(count)
190    }
191
192    /// Return the count.
193    #[inline]
194    pub const fn get(self) -> usize {
195        self.0
196    }
197}
198
199/// Available accelerator route counts by device family.
200#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
201pub struct AcceleratorCounts {
202    cpu: usize,
203    gpu: usize,
204    tpu: usize,
205    npu: usize,
206}
207
208impl AcceleratorCounts {
209    /// Construct accelerator counts by family.
210    #[inline]
211    pub const fn new(cpu: usize, gpu: usize, tpu: usize, npu: usize) -> Self {
212        Self { cpu, gpu, tpu, npu }
213    }
214
215    /// Return the CPU placement count.
216    #[inline]
217    pub const fn cpu(self) -> usize {
218        self.cpu
219    }
220
221    /// Return the GPU placement count.
222    #[inline]
223    pub const fn gpu(self) -> usize {
224        self.gpu
225    }
226
227    /// Return the TPU placement count.
228    #[inline]
229    pub const fn tpu(self) -> usize {
230        self.tpu
231    }
232
233    /// Return the NPU placement count.
234    #[inline]
235    pub const fn npu(self) -> usize {
236        self.npu
237    }
238
239    /// Return the total accelerator target count.
240    #[inline]
241    pub const fn total(self) -> usize {
242        self.cpu + self.gpu + self.tpu + self.npu
243    }
244}