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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
//! Batching tasks.
//!
//! [`Batch`] represents a prepared batch of tasks. It offers functionality for
//! preparing new batches and extracting tasks from them, but does not consider
//! the collection of tasks from external sources.
use core::{num::NonZeroUsize, ops::Range};
use crate::{Queue, Task, TaskPriority, config::Config, pubqueue::PubQueue};
//----------- Batch ------------------------------------------------------------
/// A worker's batch of tasks.
pub struct Batch<T: Task> {
/// The batch counter.
///
/// This counts down the number of tasks in the batch and decides where the
/// next task will be retrieved from (the local or the public queue).
pub counter: BatchCounter,
/// The local queue of tasks.
pub local: Vec<T>,
/// The public queue of tasks.
pub pub_queue: PubQueue,
/// The priority of the public queue.
///
/// This identifies the priority of the tasks in the public queue, if it is
/// non-empty. Otherwise, it is [`None`].
pub pq_priority: Option<T::Priority>,
}
impl<T: Task> Batch<T> {
/// Construct a new [`Batch`].
///
/// ## Safety
///
/// - `id` is the worker's ID.
/// - The worker's public queue has not been modified since initialization.
pub unsafe fn new(id: usize, config: &Config) -> Self {
debug_assert!(id < config.num_workers.get());
Self {
counter: BatchCounter::empty(),
local: Vec::with_capacity(config.batch_size.get()),
// SAFETY:
// - 'id' is the worker's ID, and is thus a valid ID.
// - 'config.batch_size <= config.batch_size'.
pub_queue: unsafe { PubQueue::new(id, config.batch_size, config) },
pq_priority: None,
}
}
/// Fill the batch.
///
/// ## Safety
///
/// - `id` is the worker's ID.
/// - `global` is the worker's global queue.
#[inline]
pub unsafe fn fill(
&mut self,
reserved: usize,
tasks: &mut Vec<T>,
id: usize,
global: &Queue<T>,
) {
// Initialize the batch.
self.counter = BatchCounter::new(tasks.len(), reserved, &global.config);
// Determine which tasks are part of the batch.
let batch_offset = tasks.len() - self.counter.total();
let priority = tasks.get(batch_offset + 1).map(|t| t.priority());
let mut tasks = tasks.drain(batch_offset..);
debug_assert!(self.pq_priority.is_none());
let pq = self.pub_queue.id(&global.config);
let pq = &global.pq_contents[pq];
if <T::Priority as TaskPriority>::TRIVIAL {
// Put half the tasks in the public queue.
// SAFETY:
// - 'tasks' has 'batch.total() <= pq_size * 2' elements, of which
// half are filtered out and sent to the local queue. Thus, the
// iterator will produce at most 'pq_size' elements.
// - 'pq' is owned by this worker, and it has not published a
// priority.
unsafe { pq.fill(tasks.by_ref().take(self.counter.public_num())) };
// Put half the tasks in the local queue.
self.local.extend(tasks);
} else {
// Divide the tasks between the local and public queues.
let mut local = false;
let tasks = tasks.filter_map(|task| {
local = !local;
if local {
self.local.push(task);
None
} else {
Some(task)
}
});
// Fill the local and public queues.
// SAFETY:
// - 'tasks' has 'batch.total() <= pq_size * 2' elements, of which
// half are filtered out and sent to the local queue. Thus, the
// iterator will produce at most 'pq_size' elements.
// - 'pq' is owned by this worker, and it has not published a
// priority.
unsafe { pq.fill(tasks) };
debug_assert_eq!(self.local.len(), self.counter.local_num());
}
// If the public queue is non-empty, publish it.
if let Some(priority) = priority {
let pq_len = self.counter.public_num();
// SAFETY: As per the batch, the public queue is non-empty.
let pq_len = unsafe { NonZeroUsize::new_unchecked(pq_len) };
let stealer = &global.stealer[id];
// SAFETY:
// - 'pq_len <= pq_size' as per 'Batch::public_num()'.
let new_pubqueue =
unsafe { self.pub_queue.with_len(pq_len, &global.config) };
// SAFETY:
// - This worker owns 'stealer'.
// - The published priority is 'None'.
// - 'last_pubqueue' is the last known value of 'stealer'.
unsafe { stealer.set(self.pub_queue, new_pubqueue) };
self.pub_queue = new_pubqueue;
// SAFETY:
// - This worker owns this priority.
// - The corresponding stealer has been initialized.
unsafe { global.priority[id].set(priority) };
self.pq_priority = Some(priority);
// Wake up a sleeping thread, if any.
global.wake();
}
}
/// Retrieve the next task in the batch.
///
/// As per the batch counter, a task will be retrieved from the local or
/// public queue. If the batch is depleted, [`None`] is returned instead.
///
/// ## Safety
///
/// - `id` is the worker's ID.
/// - `global` is the worker's global queue.
#[inline]
pub unsafe fn next(&mut self, id: usize, global: &Queue<T>) -> Option<T> {
if <T::Priority as TaskPriority>::TRIVIAL {
// NOTE: We don't keep the batch counter updated here.
// All tasks have equal priority, so it doesn't matter whether we
// read from the local queue or the public queue. The local queue
// is easier and more efficient to read from.
return self.local.pop();
}
debug_assert_eq!(self.counter.local_num(), self.local.len());
if self.counter.public_num() != 0 {
debug_assert_eq!(
self.counter.public_num(),
self.pub_queue.len(&global.config).get()
);
debug_assert!(self.pq_priority.is_some());
}
// Tick off the counter and pick the queue to read from.
if self.counter.take()? {
// Read the task from the local queue.
// SAFETY:
// - Before 'take()', 'counter.local_num() == local.len()'.
// - After 'take()', 'local_num()' decreased by 1.
// - Thus 'local_num()' was non-zero.
// - Thus 'local.len()' is still non-zero.
// - Thus 'local' is non-empty.
Some(unsafe { self.local.pop().unwrap_unchecked() })
} else {
// Read the task from the public queue.
let index = self.counter.public_num();
let pq = self.pub_queue.id(&global.config);
let len = self.pub_queue.len(&global.config);
let pq = &global.pq_contents[pq];
// Pre-emptively read this task from the public queue.
debug_assert_eq!(index + 1, len.get());
// SAFETY: 'index < len <= pq_size'
let task = unsafe { pq.read(index) };
// Take ownership of this task.
if index > 0 {
// Try updating the public queue length.
let stealer = &global.stealer[id];
// SAFETY:
// - 'stealer' belongs to this worker.
// - 'last_pubqueue' is the last known stealer value.
// - 'index = len - 1 > 0' so 'len > 0'.
match unsafe { stealer.try_dec(self.pub_queue, &global.config) }
{
Ok(pq) => {
// The public queue was updated successfully.
self.pub_queue = pq;
}
Err(pq) => {
// The public queue has been stolen.
self.pq_priority = None;
self.pub_queue = pq;
return None;
}
}
} else {
// Taking this task would empty the public queue, which
// means ownership of the whole public queue is needed.
// SAFETY: As per the counter, the public queue contains tasks
// and has not yet been stolen.
let last_priority =
unsafe { self.pq_priority.take().unwrap_unchecked() };
let priority = &global.priority[id];
match unsafe { priority.steal_self(last_priority) } {
Ok(()) => {
// Ownership of the public queue has been taken.
}
Err(()) => {
// Another worker is in the middle of stealing from
// us.
let stealer = &global.stealer[id];
// Wait until the public queue is stolen.
// SAFETY:
// - As per 'Priority::steal_self()', the priority
// of this worker was set to 'None' by another
// worker.
// - 'last_pubqueue' is the last value for 'stealer'
// seen by this worker.
self.pub_queue =
unsafe { stealer.wait_for_theft(self.pub_queue) };
return None;
}
}
}
// SAFETY: Taking ownership of the task was successful. This
// means that nobody could have overwritten the task since we
// had filled the public queue, so our read of the task was
// coherent and read initialized data.
Some(unsafe { task.assume_init() })
}
}
/// Drain the batch.
///
/// ## Safety
///
/// - `id` is the worker's ID.
/// - `global` is the worker's global queue.
#[inline]
pub unsafe fn drain(
&mut self,
tasks: &mut Vec<T>,
id: usize,
global: &Queue<T>,
) {
// Move out of the local queue.
tasks.append(&mut self.local);
// Collect tasks from the public queue.
if let Some(priority) = self.pq_priority.take() {
// SAFETY:
// - This 'Priority' belongs to the current worker.
// - 'last_priority' was 'Some(priority)'.
match unsafe { global.priority[id].steal_self(priority) } {
Ok(()) => {
// We have taken ownership of the public queue.
let pq = self.pub_queue.id(&global.config);
let len = self.pub_queue.len(&global.config);
let pq = &global.pq_contents[pq];
// Move the tasks out of the public queue.
// SAFETY:
// - 'len <= pq_size' as per 'PubQueue::len()'.
// - We set the length of the public queue to 'len', and
// nobody stole it, so the length is unchanged.
// - This worker owns 'pq' so nobody will write to it.
unsafe { pq.move_to(len.get(), tasks) };
}
Err(()) => {
// Another worker is in the middle of stealing from us.
let stealer = &global.stealer[id];
// Wait until the public queue is stolen.
// SAFETY:
// - As per 'Priority::steal_self()', the priority of this
// worker was set to 'None' by another worker.
// - 'last_pubqueue' is the last value for 'stealer' seen by
// this worker.
self.pub_queue =
unsafe { stealer.wait_for_theft(self.pub_queue) };
}
}
}
// Reset the batch counter.
self.counter = BatchCounter::empty();
}
}
//----------- BatchCounter -----------------------------------------------------
/// A counter for a [`Batch`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BatchCounter {
/// The range of pending task numbers.
///
/// `range.end - range.start` is the number of pending tasks.
///
/// Invariants:
/// - `range.start <= range.end`.
/// - `range.end <= pq_size * 2`.
range: Range<usize>,
}
impl BatchCounter {
/// Construct an empty [`BatchCounter`].
pub const fn empty() -> Self {
Self { range: 0..0 }
}
/// Construct a new [`BatchCounter`].
///
/// `remaining` is the number of tasks to be divided between the local and
/// public queues. If some tasks have already been removed (for special
/// handling) but should be counted against the maximum batch size, their
/// count should be passed in `reserved`. The batch will contain at most
/// `config.batch_size - reserved` tasks.
pub const fn new(
remaining: usize,
reserved: usize,
config: &Config,
) -> Self {
let batch_size = config.batch_size.get();
debug_assert!(reserved <= batch_size * 2);
// Determine the total number of tasks for the batch.
//let num = remaining.max(pq_size * 2 - reserved);
let total = if remaining < batch_size * 2 - reserved {
remaining
} else {
batch_size - reserved
};
// Separate the tasks into pending and not-pending sets.
let not_pending = total / 2;
Self {
range: not_pending..total,
}
}
/// Whether the batch is empty.
pub const fn is_empty(&self) -> bool {
self.range.start >= self.range.end
}
/// The total number of tasks in the batch.
pub const fn total(&self) -> usize {
self.range.end
}
/// The number of batch tasks in the local queue.
///
/// There may be additional tasks in the local queue before the batch.
pub const fn local_num(&self) -> usize {
// 0, 1, 2, 3, 4, 5, 6, 7, ... -> 0, 1, 1, 2, 2, 3, 3, 4, ...
self.range.end.div_ceil(2)
}
/// The number of batch tasks in the public queue.
///
/// There are no other tasks in the public queue.
pub const fn public_num(&self) -> usize {
// 0, 1, 2, 3, 4, 5, 6, 7, ... -> 0, 0, 1, 1, 2, 2, 3, 3, ...
self.range.end / 2
}
/// Extract a pending task.
///
/// If the task is part of the local queue, `true` is returned.
pub const fn take(&mut self) -> Option<bool> {
if self.range.start >= self.range.end {
return None;
}
// SAFETY: '0 <= range.start < range.end' so 'range.end > 0'.
self.range.end -= 1;
// 0, 1, 2, 3, 4, 5, 6, 7, ... -> L, P, L, P, L, P, L, P, ...
Some(self.range.end % 2 == 0)
}
}