cubecl_environment/sync/

base.rs

#[cfg(not(feature = "std"))]
use spin::{Mutex as MutexImported, MutexGuard, Once as OnceImported, RwLock as RwLockImported};
#[cfg(feature = "std")]
use std::sync::{
    Mutex as MutexImported, MutexGuard, OnceLock as OnceImported, RwLock as RwLockImported,
};

#[cfg(not(feature = "std"))]
pub use spin::{LazyLock, RwLockReadGuard, RwLockWriteGuard};
#[cfg(feature = "std")]
pub use std::sync::{LazyLock, RwLockReadGuard, RwLockWriteGuard};

/// A spin-based one-time initialization cell, identical on every target.
///
/// Prefer [`SyncOnceCell`] for plain lazy initialization; use this when the
/// fallible [`spin::Once::try_call_once`] API is needed.
pub use spin::Once;

#[cfg(target_has_atomic = "ptr")]
pub use alloc::sync::Arc;
#[cfg(not(target_has_atomic = "ptr"))]
pub use portable_atomic_util::Arc;

#[cfg(target_has_atomic = "ptr")]
pub use core::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicUsize, Ordering};
#[cfg(not(target_has_atomic = "ptr"))]
pub use portable_atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicUsize, Ordering};

// No `AtomicU64`: `target_has_atomic = "ptr"` says nothing about 64-bit
// atomics, and a 32-bit target can have pointer-width CAS without them
// (thumbv7m). Its one user is behind `stream_local`, which is `std`, and every
// std target has them — so it takes `core`'s directly rather than making this
// shim carry a width it cannot provide on both arms.

/// A mutual exclusion primitive useful for protecting shared data
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can also be statically initialized or created via a [`Mutex::new`]
///
/// [Mutex] wrapper to make `spin::Mutex` API compatible with `std::sync::Mutex` to swap
#[derive(Debug, Default)]
pub struct Mutex<T: ?Sized> {
    inner: MutexImported<T>,
}

impl<T> Mutex<T> {
    /// Creates a new mutex in an unlocked state ready for use.
    #[inline(always)]
    pub const fn new(value: T) -> Self {
        Self {
            inner: MutexImported::new(value),
        }
    }

    /// Consumes the mutex, returning the value it protected. A poisoned lock
    /// is recovered, as in [`lock`](Self::lock).
    #[inline(always)]
    pub fn into_inner(self) -> T {
        #[cfg(not(feature = "std"))]
        {
            self.inner.into_inner()
        }

        #[cfg(feature = "std")]
        {
            self.inner
                .into_inner()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
        }
    }
}

impl<T: ?Sized> Mutex<T> {
    /// Locks the mutex blocking the current thread until it is able to do so.
    ///
    /// Locking cannot fail, so no `Result` is returned. A poisoned lock is
    /// recovered rather than reported: a panic under a guard must not turn
    /// every later lock into a panic, matching the unwind behavior of the spin
    /// implementation used off-std.
    #[inline(always)]
    pub fn lock(&self) -> MutexGuard<'_, T> {
        #[cfg(not(feature = "std"))]
        {
            self.inner.lock()
        }

        #[cfg(feature = "std")]
        {
            self.inner
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
        }
    }
}

/// A reader-writer lock which is exclusively locked for writing or shared for reading.
/// This reader-writer lock will block threads waiting for the lock to become available.
/// The lock can also be statically initialized or created via a [`RwLock::new`]
/// [`RwLock`] wrapper to make `spin::RwLock` API compatible with `std::sync::RwLock` to swap
#[derive(Debug)]
pub struct RwLock<T: ?Sized> {
    inner: RwLockImported<T>,
}

impl<T> RwLock<T> {
    /// Creates a new reader-writer lock in an unlocked state ready for use.
    #[inline(always)]
    pub const fn new(value: T) -> Self {
        Self {
            inner: RwLockImported::new(value),
        }
    }
}

impl<T: ?Sized> RwLock<T> {
    /// Locks this rwlock with shared read access, blocking the current thread
    /// until it can be acquired.
    ///
    /// Poisoning is recovered, never reported; see [`Mutex::lock`].
    #[inline(always)]
    pub fn read(&self) -> RwLockReadGuard<'_, T> {
        #[cfg(not(feature = "std"))]
        {
            self.inner.read()
        }
        #[cfg(feature = "std")]
        {
            self.inner
                .read()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
        }
    }

    /// Locks this rwlock with exclusive write access, blocking the current thread
    /// until it can be acquired.
    ///
    /// Poisoning is recovered, never reported; see [`Mutex::lock`].
    #[inline(always)]
    pub fn write(&self) -> RwLockWriteGuard<'_, T> {
        #[cfg(not(feature = "std"))]
        {
            self.inner.write()
        }

        #[cfg(feature = "std")]
        {
            self.inner
                .write()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
        }
    }
}

/// An opaque thread identifier.
///
/// This is a stub when no std is available to swap with `std::thread::ThreadId`.
/// There is no way to obtain one on such targets, so it exists only to keep
/// thread-keyed types nameable.
#[allow(dead_code)]
#[derive(Eq, PartialEq, Clone, Copy, Hash, Debug)]
pub struct ThreadId(core::num::NonZeroU64);

/// A cell that provides lazy one-time initialization that implements [Sync] and [Send].
///
/// This module is a stub when no std is available to swap with [`std::sync::OnceLock`].
pub struct SyncOnceCell<T>(OnceImported<T>);

impl<T> Default for SyncOnceCell<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> SyncOnceCell<T> {
    /// Create a new once.
    #[inline(always)]
    pub fn new() -> Self {
        Self(OnceImported::new())
    }

    /// Initialize the cell with a value.
    #[inline(always)]
    pub fn initialized(value: T) -> Self {
        #[cfg(not(feature = "std"))]
        {
            let cell = OnceImported::initialized(value);
            Self(cell)
        }

        #[cfg(feature = "std")]
        {
            let cell = OnceImported::new();
            // Infallible: the cell was just created, so it is empty. Ignoring
            // the `Err` is what keeps `T: Debug` off this whole impl.
            let _ = cell.set(value);

            Self(cell)
        }
    }

    /// Gets the contents of the cell, initializing it with `f` if the cell
    /// was empty.
    #[inline(always)]
    pub fn get_or_init<F>(&self, f: F) -> &T
    where
        F: FnOnce() -> T,
    {
        #[cfg(not(feature = "std"))]
        {
            self.0.call_once(f)
        }

        #[cfg(feature = "std")]
        {
            self.0.get_or_init(f)
        }
    }
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;

    /// Regression: a panic under a guard must not poison the lock — one
    /// failed autotune must not take down every later kernel launch.
    #[test]
    fn poisoned_mutex_recovers() {
        let mutex = Mutex::new(0u32);

        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let _guard = mutex.lock();
            panic!("poison the lock");
        }))
        .unwrap_err();

        *mutex.lock() += 1;
        assert_eq!(*mutex.lock(), 1);
    }

    #[test]
    fn poisoned_rwlock_recovers() {
        let lock = RwLock::new(0u32);

        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let _guard = lock.write();
            panic!("poison the lock");
        }))
        .unwrap_err();

        assert_eq!(*lock.read(), 0);
    }

    #[test]
    fn a_mutex_locks_a_trait_object() {
        let mutex: Arc<Mutex<dyn Iterator<Item = u32> + Send>> = Arc::new(Mutex::new(0..3));

        assert_eq!(mutex.lock().next(), Some(0));
        assert_eq!(mutex.lock().next(), Some(1));
    }

    #[test]
    fn an_rwlock_locks_a_trait_object() {
        let lock: Arc<RwLock<dyn Iterator<Item = u32> + Send + Sync>> = Arc::new(RwLock::new(0..3));

        assert_eq!(lock.write().next(), Some(0));
        assert_eq!(lock.read().size_hint(), (2, Some(2)));
    }
}