Skip to content

ARM64 threading optimization - atomic ordering and cache line alignment #95

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
glorv merged 3 commits into from
Aug 25, 2026
Merged

ARM64 threading optimization - atomic ordering and cache line alignment #95

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16bf2bd
minimum correct ordering changes
mohakgoel1 Jun 7, 2026
cffc852
Cover 256-byte AArch64 cache lines
mohakgoel1 Jul 28, 2026
9d513e5
Replace custom CacheAligned with crossbeam_utils::CachePadded for act…
mohakgoel1 Aug 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions src/pool/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ impl Clone for SchedConfig {
SchedConfig {
max_thread_count: self.max_thread_count,
min_thread_count: self.min_thread_count,
core_thread_count: AtomicUsize::new(self.core_thread_count.load(Ordering::SeqCst)),
core_thread_count: AtomicUsize::new(self.core_thread_count.load(Ordering::Acquire)),
max_inplace_spin: self.max_inplace_spin,
max_idle_time: self.max_idle_time,
max_wait_time: self.max_wait_time,
Expand Down Expand Up @@ -166,7 +166,7 @@ impl Builder {
if count > 0 {
self.sched_config
.core_thread_count
.store(count, Ordering::SeqCst);
.store(count, Ordering::Release);
}
self
}
Expand Down Expand Up @@ -241,15 +241,15 @@ impl Builder {
T: TaskCell + Send,
{
assert!(self.sched_config.min_thread_count <= self.sched_config.max_thread_count);
let core_thread_count = self.sched_config.core_thread_count.load(Ordering::SeqCst);
let core_thread_count = self.sched_config.core_thread_count.load(Ordering::Acquire);
if core_thread_count == 0 || core_thread_count > self.sched_config.max_thread_count {
self.sched_config
.core_thread_count
.store(self.sched_config.max_thread_count, Ordering::SeqCst);
.store(self.sched_config.max_thread_count, Ordering::Release);
} else if core_thread_count < self.sched_config.min_thread_count {
self.sched_config
.core_thread_count
.store(self.sched_config.min_thread_count, Ordering::SeqCst);
.store(self.sched_config.min_thread_count, Ordering::Release);
}
let (injector, local_queues) = queue::build(queue_type, self.sched_config.max_thread_count);
let core = Arc::new(QueueCore::new(injector, self.sched_config.clone()));
Expand Down
37 changes: 21 additions & 16 deletions src/pool/spawn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

use crate::pool::SchedConfig;
use crate::queue::{Extras, LocalQueue, Pop, PopResult, TaskCell, TaskInjector, WithExtras};
use crossbeam_utils::CachePadded;
use fail::fail_point;
use parking_lot_core::{FilterOp, ParkResult, ParkToken, UnparkToken};
use std::sync::{
Expand All @@ -14,6 +15,9 @@ use std::sync::{
};
use std::time::Instant;

// crossbeam_utils::CachePadded is used to isolate active_workers to its
// own cache line, preventing false sharing with global_queue in QueueCore.

/// An usize is used to trace the threads that are working actively.
/// To save additional memory and atomic operation, the number and
/// shutdown hint are merged into one number in the following format
Expand All @@ -36,17 +40,18 @@ pub fn is_shutdown(cnt: usize) -> bool {
///
/// Every thread pool instance should have one and only `QueueCore`. It's
/// saved in an `Arc` and shared between all worker threads and remote handles.

pub(crate) struct QueueCore<T> {
global_queue: TaskInjector<T>,
active_workers: AtomicUsize,
active_workers: CachePadded<AtomicUsize>,
config: SchedConfig,
}

impl<T> QueueCore<T> {
pub fn new(global_queue: TaskInjector<T>, config: SchedConfig) -> QueueCore<T> {
QueueCore {
global_queue,
active_workers: AtomicUsize::new(config.max_thread_count << WORKER_COUNT_SHIFT),
active_workers: CachePadded::new(AtomicUsize::new(config.max_thread_count << WORKER_COUNT_SHIFT)),
config,
}
}
Expand All @@ -56,8 +61,8 @@ impl<T> QueueCore<T> {
/// If the method is going to wake up any threads, source is used to trace who triggers
/// the action.
pub fn ensure_workers(&self, source: usize) {
let cnt = self.active_workers.load(Ordering::SeqCst);
if (cnt >> WORKER_COUNT_SHIFT) >= self.config.core_thread_count.load(Ordering::SeqCst)
let cnt = self.active_workers.load(Ordering::Acquire);
if (cnt >> WORKER_COUNT_SHIFT) >= self.config.core_thread_count.load(Ordering::Acquire)
|| is_shutdown(cnt)
{
return;
Expand All @@ -74,7 +79,7 @@ impl<T> QueueCore<T> {
parking_lot_core::unpark_filter(
addr,
|p: ParkToken| {
if !unparked_once && p.0 <= self.config.core_thread_count.load(Ordering::SeqCst)
if !unparked_once && p.0 <= self.config.core_thread_count.load(Ordering::Acquire)
{
unparked_once = true;
FilterOp::Unpark
Expand All @@ -91,7 +96,7 @@ impl<T> QueueCore<T> {
///
/// `source` is used to trace who triggers the action.
pub fn mark_shutdown(&self, source: usize) {
self.active_workers.fetch_or(SHUTDOWN_BIT, Ordering::SeqCst);
self.active_workers.fetch_or(SHUTDOWN_BIT, Ordering::AcqRel);
let addr = self as *const QueueCore<T> as usize;
unsafe {
parking_lot_core::unpark_all(addr, UnparkToken(source));
Expand All @@ -100,15 +105,15 @@ impl<T> QueueCore<T> {

/// Checks if the thread pool is shutting down.
pub fn is_shutdown(&self) -> bool {
let cnt = self.active_workers.load(Ordering::SeqCst);
let cnt = self.active_workers.load(Ordering::Acquire);
is_shutdown(cnt)
}

/// Marks the current thread in sleep state.
///
/// It can be marked as sleep only when the pool is not shutting down.
pub fn mark_sleep(&self) -> bool {
let mut cnt = self.active_workers.load(Ordering::SeqCst);
let mut cnt = self.active_workers.load(Ordering::Acquire);
loop {
if is_shutdown(cnt) {
return false;
Expand All @@ -117,8 +122,8 @@ impl<T> QueueCore<T> {
match self.active_workers.compare_exchange_weak(
cnt,
cnt - WORKER_COUNT_BASE,
Ordering::SeqCst,
Ordering::SeqCst,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return true,
Err(n) => cnt = n,
Expand All @@ -128,13 +133,13 @@ impl<T> QueueCore<T> {

/// Marks current thread as woken up states.
pub fn mark_woken(&self) {
let mut cnt = self.active_workers.load(Ordering::SeqCst);
let mut cnt = self.active_workers.load(Ordering::Acquire);
loop {
match self.active_workers.compare_exchange_weak(
cnt,
cnt + WORKER_COUNT_BASE,
Ordering::SeqCst,
Ordering::SeqCst,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return,
Err(n) => cnt = n,
Expand All @@ -151,7 +156,7 @@ impl<T> QueueCore<T> {
}
self.config
.core_thread_count
.store(new_thread_count, Ordering::SeqCst);
.store(new_thread_count, Ordering::Release);
}

pub fn config(&self) -> &SchedConfig {
Expand Down Expand Up @@ -319,7 +324,7 @@ impl<T: TaskCell + Send> Local<T> {
}

pub(crate) fn is_scaled_down_worker(&self) -> bool {
self.id > self.core.config.core_thread_count.load(Ordering::SeqCst)
self.id > self.core.config.core_thread_count.load(Ordering::Acquire)
}

pub(crate) fn drain(&mut self) {
Expand Down Expand Up @@ -355,7 +360,7 @@ impl<T: TaskCell + Send> Local<T> {
marked_sleep = true;
// If this thread is above core_thread_count, go to sleep
// without popping so scaled-down threads don't keep working.
if id > self.core.config.core_thread_count.load(Ordering::SeqCst) {
if id > self.core.config.core_thread_count.load(Ordering::Acquire) {
return true;
}
fail_point!("worker-pop-or-sleep-before-validate-pop");
Expand Down
18 changes: 9 additions & 9 deletions src/queue/multilevel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -496,7 +496,7 @@ impl LevelManager {
fn maybe_adjust_chance(&self) {
if self
.adjusting
.compare_exchange(false, true, SeqCst, SeqCst)
.compare_exchange(false, true, AcqRel, Acquire)
.is_err()
{
return;
Expand All @@ -507,7 +507,7 @@ impl LevelManager {
let total_diff = total - self.last_total_elapsed_us.get();
if total_diff < ADJUST_CHANCE_INTERVAL_US {
// Needn't change it now.
self.adjusting.store(false, SeqCst);
self.adjusting.store(false, Release);
return;
}
let level0 = self.level0_elapsed_us.get();
Expand Down Expand Up @@ -549,13 +549,13 @@ impl LevelManager {
},
);
self.max_level_queue_steal_size
.store(new_steal_count, SeqCst);
.store(new_steal_count, Release);
for (i, c) in self.last_exec_tasks_per_level.iter().enumerate() {
c.set(cur_total_tasks_per_level[i]);
}
}

self.adjusting.store(false, SeqCst);
self.adjusting.store(false, Release);
}
}

Expand Down Expand Up @@ -686,7 +686,7 @@ impl TaskElapsedMap {
}

fn get_elapsed(&self, key: u64) -> Arc<ElapsedTime> {
let new_index = self.new_index.load(SeqCst);
let new_index = self.new_index.load(Acquire);
let new_map = &self.maps[new_index];
let old_map = &self.maps[new_index ^ 1];
if let Some(v) = new_map.get(&key) {
Expand Down Expand Up @@ -714,14 +714,14 @@ impl TaskElapsedMap {

fn try_cleanup(&self) -> Option<Instant> {
self.cleaning_up
.compare_exchange(false, true, SeqCst, SeqCst)
.compare_exchange(false, true, AcqRel, Acquire)
.map(|_| {
let old_index = self.new_index.load(SeqCst) ^ 1;
let old_index = self.new_index.load(Acquire) ^ 1;
self.maps[old_index].clear();
self.new_index.store(old_index, SeqCst);
self.new_index.store(old_index, Release);
let now = now();
*self.last_cleanup_time.lock().unwrap() = now;
self.cleaning_up.store(false, SeqCst);
self.cleaning_up.store(false, Release);
now
})
.ok()
Expand Down
14 changes: 7 additions & 7 deletions src/task/future.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use std::pin::Pin;
use std::ptr::NonNull;
use std::sync::atomic::{
AtomicU8, AtomicUsize,
Ordering::{Acquire, Relaxed, Release, SeqCst},
Ordering::{AcqRel, Acquire, Relaxed, Release},
};
use std::sync::{atomic, Arc};
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
Expand Down Expand Up @@ -243,13 +243,13 @@ unsafe fn drop_raw(this: *const ()) {
}

unsafe fn wake_impl(task: Cow<'_, TaskCell>) {
let mut status = task.status().load(SeqCst);
let mut status = task.status().load(Acquire);
loop {
match status {
IDLE => {
match task
.status()
.compare_exchange_weak(IDLE, NOTIFIED, SeqCst, SeqCst)
.compare_exchange_weak(IDLE, NOTIFIED, AcqRel, Acquire)
{
Ok(_) => {
wake_task(task, false);
Expand All @@ -261,7 +261,7 @@ unsafe fn wake_impl(task: Cow<'_, TaskCell>) {
POLLING => {
match task
.status()
.compare_exchange_weak(POLLING, NOTIFIED, SeqCst, SeqCst)
.compare_exchange_weak(POLLING, NOTIFIED, AcqRel, Acquire)
{
Ok(_) => break,
Err(cur) => status = cur,
Expand Down Expand Up @@ -390,9 +390,9 @@ impl crate::pool::Runner for Runner {
let mut cx = waker_ref.to_context();
let mut repoll_times = 0;
loop {
task_cell.status().store(POLLING, SeqCst);
task_cell.status().store(POLLING, Release);
if task_cell.poll(&mut cx).is_ready() {
task_cell.status().store(COMPLETED, SeqCst);
task_cell.status().store(COMPLETED, Release);
return true;
}
let extras = { &mut *task_cell.extras().get() };
Expand All @@ -406,7 +406,7 @@ impl crate::pool::Runner for Runner {
}
match task_cell
.status()
.compare_exchange(POLLING, IDLE, SeqCst, SeqCst)
.compare_exchange(POLLING, IDLE, AcqRel, Acquire)
{
Ok(_) => return false,
Err(NOTIFIED) => {
Expand Down
Loading