LKMM atomics in Rust

[PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations

Posted by Boqun Feng 2 months, 2 weeks ago

One important set of atomic operations is the arithmetic operations,
i.e. add(), sub(), fetch_add(), add_return(), etc. However it may not
make senses for all the types that `AtomicType` to have arithmetic
operations, for example a `Foo(u32)` may not have a reasonable add() or
sub(), plus subword types (`u8` and `u16`) currently don't have
atomic arithmetic operations even on C side and might not have them in
the future in Rust (because they are usually suboptimal on a few
architecures). Therefore the plan is to add a few subtraits of
`AtomicType` describing which types have and can do atomic arithemtic
operations.

One trait `AtomicAdd` is added, and only add() and fetch_add() are
added. The rest will be added in the future.

Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 rust/kernel/sync/atomic.rs           | 93 +++++++++++++++++++++++++++-
 rust/kernel/sync/atomic/predefine.rs | 14 +++++
 2 files changed, 105 insertions(+), 2 deletions(-)

diff --git a/rust/kernel/sync/atomic.rs b/rust/kernel/sync/atomic.rs
index 793134aeaac1..e3a30b6aaee4 100644
--- a/rust/kernel/sync/atomic.rs
+++ b/rust/kernel/sync/atomic.rs
@@ -16,7 +16,6 @@
 //!
 //! [`LKMM`]: srctree/tools/memory-model/
 
-#[allow(dead_code, unreachable_pub)]
 mod internal;
 pub mod ordering;
 mod predefine;
@@ -25,7 +24,7 @@
 pub use ordering::{Acquire, Full, Relaxed, Release};
 
 use crate::build_error;
-use internal::{AtomicBasicOps, AtomicExchangeOps, AtomicRepr};
+use internal::{AtomicArithmeticOps, AtomicBasicOps, AtomicExchangeOps, AtomicRepr};
 use ordering::OrderingType;
 
 /// A memory location which can be safely modified from multiple execution contexts.
@@ -115,6 +114,18 @@ pub unsafe trait AtomicType: Sized + Send + Copy {
     type Repr: AtomicImpl;
 }
 
+/// Types that support atomic add operations.
+///
+/// # Safety
+///
+/// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
+/// any value of type `Self::Repr` obtained through transmuting a value of type `Self` to must
+/// yield a value with a bit pattern also valid for `Self`.
+pub unsafe trait AtomicAdd<Rhs = Self>: AtomicType {
+    /// Converts `Rhs` into the `Delta` type of the atomic implementation.
+    fn rhs_into_delta(rhs: Rhs) -> <Self::Repr as AtomicImpl>::Delta;
+}
+
 #[inline(always)]
 const fn into_repr<T: AtomicType>(v: T) -> T::Repr {
     // SAFETY: Per the safety requirement of `AtomicType`, `T` is round-trip transmutable to
@@ -462,3 +473,81 @@ fn try_cmpxchg<Ordering: ordering::Ordering>(&self, old: &mut T, new: T, _: Orde
         ret
     }
 }
+
+impl<T: AtomicType> Atomic<T>
+where
+    T::Repr: AtomicArithmeticOps,
+{
+    /// Atomic add.
+    ///
+    /// Atomically updates `*self` to `(*self).wrapping_add(v)`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::sync::atomic::{Atomic, Relaxed};
+    ///
+    /// let x = Atomic::new(42);
+    ///
+    /// assert_eq!(42, x.load(Relaxed));
+    ///
+    /// x.add(12, Relaxed);
+    ///
+    /// assert_eq!(54, x.load(Relaxed));
+    /// ```
+    #[inline(always)]
+    pub fn add<Rhs>(&self, v: Rhs, _: ordering::Relaxed)
+    where
+        T: AtomicAdd<Rhs>,
+    {
+        let v = T::rhs_into_delta(v);
+
+        // INVARIANT: `self.0` is a valid `T` after `atomic_add()` due to safety requirement of
+        // `AtomicAdd`.
+        T::Repr::atomic_add(&self.0, v);
+    }
+
+    /// Atomic fetch and add.
+    ///
+    /// Atomically updates `*self` to `(*self).wrapping_add(v)`, and returns the value of `*self`
+    /// before the update.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::sync::atomic::{Atomic, Acquire, Full, Relaxed};
+    ///
+    /// let x = Atomic::new(42);
+    ///
+    /// assert_eq!(42, x.load(Relaxed));
+    ///
+    /// assert_eq!(54, { x.fetch_add(12, Acquire); x.load(Relaxed) });
+    ///
+    /// let x = Atomic::new(42);
+    ///
+    /// assert_eq!(42, x.load(Relaxed));
+    ///
+    /// assert_eq!(54, { x.fetch_add(12, Full); x.load(Relaxed) } );
+    /// ```
+    #[inline(always)]
+    pub fn fetch_add<Rhs, Ordering: ordering::Ordering>(&self, v: Rhs, _: Ordering) -> T
+    where
+        T: AtomicAdd<Rhs>,
+    {
+        let v = T::rhs_into_delta(v);
+
+        // INVARIANT: `self.0` is a valid `T` after `atomic_fetch_add*()` due to safety requirement
+        // of `AtomicAdd`.
+        let ret = {
+            match Ordering::TYPE {
+                OrderingType::Full => T::Repr::atomic_fetch_add(&self.0, v),
+                OrderingType::Acquire => T::Repr::atomic_fetch_add_acquire(&self.0, v),
+                OrderingType::Release => T::Repr::atomic_fetch_add_release(&self.0, v),
+                OrderingType::Relaxed => T::Repr::atomic_fetch_add_relaxed(&self.0, v),
+            }
+        };
+
+        // SAFETY: `ret` comes from reading `self.0`, which is a valid `T` per type invariants.
+        unsafe { from_repr(ret) }
+    }
+}
diff --git a/rust/kernel/sync/atomic/predefine.rs b/rust/kernel/sync/atomic/predefine.rs
index 33356deee952..a6e5883be7cb 100644
--- a/rust/kernel/sync/atomic/predefine.rs
+++ b/rust/kernel/sync/atomic/predefine.rs
@@ -8,8 +8,22 @@ unsafe impl super::AtomicType for i32 {
     type Repr = i32;
 }
 
+// SAFETY: The wrapping add result of two `i32`s is a valid `i32`.
+unsafe impl super::AtomicAdd<i32> for i32 {
+    fn rhs_into_delta(rhs: i32) -> i32 {
+        rhs
+    }
+}
+
 // SAFETY: `i64` has the same size and alignment with itself, and is round-trip transmutable to
 // itself.
 unsafe impl super::AtomicType for i64 {
     type Repr = i64;
 }
+
+// SAFETY: The wrapping add result of two `i64`s is a valid `i64`.
+unsafe impl super::AtomicAdd<i64> for i64 {
+    fn rhs_into_delta(rhs: i64) -> i64 {
+        rhs
+    }
+}
-- 
2.39.5 (Apple Git-154)

Re: [PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations

Posted by Benno Lossin 1 month, 3 weeks ago

On Sat Jul 19, 2025 at 5:08 AM CEST, Boqun Feng wrote:
> One important set of atomic operations is the arithmetic operations,
> i.e. add(), sub(), fetch_add(), add_return(), etc. However it may not
> make senses for all the types that `AtomicType` to have arithmetic
> operations, for example a `Foo(u32)` may not have a reasonable add() or
> sub(), plus subword types (`u8` and `u16`) currently don't have
> atomic arithmetic operations even on C side and might not have them in
> the future in Rust (because they are usually suboptimal on a few
> architecures). Therefore the plan is to add a few subtraits of
> `AtomicType` describing which types have and can do atomic arithemtic
> operations.
>
> One trait `AtomicAdd` is added, and only add() and fetch_add() are
> added. The rest will be added in the future.
>
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Boqun Feng <boqun.feng@gmail.com>

Reviewed-by: Benno Lossin <lossin@kernel.org>

> ---
>  rust/kernel/sync/atomic.rs           | 93 +++++++++++++++++++++++++++-
>  rust/kernel/sync/atomic/predefine.rs | 14 +++++
>  2 files changed, 105 insertions(+), 2 deletions(-)
>
> diff --git a/rust/kernel/sync/atomic.rs b/rust/kernel/sync/atomic.rs
> index 793134aeaac1..e3a30b6aaee4 100644
> --- a/rust/kernel/sync/atomic.rs
> +++ b/rust/kernel/sync/atomic.rs
> @@ -16,7 +16,6 @@
>  //!
>  //! [`LKMM`]: srctree/tools/memory-model/
>  
> -#[allow(dead_code, unreachable_pub)]
>  mod internal;
>  pub mod ordering;
>  mod predefine;
> @@ -25,7 +24,7 @@
>  pub use ordering::{Acquire, Full, Relaxed, Release};
>  
>  use crate::build_error;
> -use internal::{AtomicBasicOps, AtomicExchangeOps, AtomicRepr};
> +use internal::{AtomicArithmeticOps, AtomicBasicOps, AtomicExchangeOps, AtomicRepr};
>  use ordering::OrderingType;
>  
>  /// A memory location which can be safely modified from multiple execution contexts.
> @@ -115,6 +114,18 @@ pub unsafe trait AtomicType: Sized + Send + Copy {
>      type Repr: AtomicImpl;
>  }
>  
> +/// Types that support atomic add operations.
> +///
> +/// # Safety
> +///
> +/// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to

Can you add a normal comment TODO here:

    // TODO: properly define `wrapping_add` in this context.

---
Cheers,
Benno

> +/// any value of type `Self::Repr` obtained through transmuting a value of type `Self` to must
> +/// yield a value with a bit pattern also valid for `Self`.
> +pub unsafe trait AtomicAdd<Rhs = Self>: AtomicType {
> +    /// Converts `Rhs` into the `Delta` type of the atomic implementation.
> +    fn rhs_into_delta(rhs: Rhs) -> <Self::Repr as AtomicImpl>::Delta;
> +}
> +

Re: [PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations

Posted by Boqun Feng 1 month, 2 weeks ago

On Tue, Aug 12, 2025 at 10:04:12AM +0200, Benno Lossin wrote:
> On Sat Jul 19, 2025 at 5:08 AM CEST, Boqun Feng wrote:
> > One important set of atomic operations is the arithmetic operations,
> > i.e. add(), sub(), fetch_add(), add_return(), etc. However it may not
> > make senses for all the types that `AtomicType` to have arithmetic
> > operations, for example a `Foo(u32)` may not have a reasonable add() or
> > sub(), plus subword types (`u8` and `u16`) currently don't have
> > atomic arithmetic operations even on C side and might not have them in
> > the future in Rust (because they are usually suboptimal on a few
> > architecures). Therefore the plan is to add a few subtraits of
> > `AtomicType` describing which types have and can do atomic arithemtic
> > operations.
> >
> > One trait `AtomicAdd` is added, and only add() and fetch_add() are
> > added. The rest will be added in the future.
> >
> > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
> 
> Reviewed-by: Benno Lossin <lossin@kernel.org>
> 

Thank you!

> > ---
> >  rust/kernel/sync/atomic.rs           | 93 +++++++++++++++++++++++++++-
> >  rust/kernel/sync/atomic/predefine.rs | 14 +++++
> >  2 files changed, 105 insertions(+), 2 deletions(-)
> >
> > diff --git a/rust/kernel/sync/atomic.rs b/rust/kernel/sync/atomic.rs
> > index 793134aeaac1..e3a30b6aaee4 100644
> > --- a/rust/kernel/sync/atomic.rs
> > +++ b/rust/kernel/sync/atomic.rs
> > @@ -16,7 +16,6 @@
> >  //!
> >  //! [`LKMM`]: srctree/tools/memory-model/
> >  
> > -#[allow(dead_code, unreachable_pub)]
> >  mod internal;
> >  pub mod ordering;
> >  mod predefine;
> > @@ -25,7 +24,7 @@
> >  pub use ordering::{Acquire, Full, Relaxed, Release};
> >  
> >  use crate::build_error;
> > -use internal::{AtomicBasicOps, AtomicExchangeOps, AtomicRepr};
> > +use internal::{AtomicArithmeticOps, AtomicBasicOps, AtomicExchangeOps, AtomicRepr};
> >  use ordering::OrderingType;
> >  
> >  /// A memory location which can be safely modified from multiple execution contexts.
> > @@ -115,6 +114,18 @@ pub unsafe trait AtomicType: Sized + Send + Copy {
> >      type Repr: AtomicImpl;
> >  }
> >  
> > +/// Types that support atomic add operations.
> > +///
> > +/// # Safety
> > +///
> > +/// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
> 
> Can you add a normal comment TODO here:
> 
>     // TODO: properly define `wrapping_add` in this context.

Yeah, this sounds good to me. How do you propose we arrange the normal
comment with the doc comment, somthing like:

    // TODO: properly define `wrapping_add` in this context.
    
    /// Types that support atomic add operations.
    ///
    /// # Safety
    ///
    /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
    ...
    pub unsafe trait AtomicAdd<...> {
        ...
    }

Regards,
Boqun

> 
> ---
> Cheers,
> Benno
> 
> > +/// any value of type `Self::Repr` obtained through transmuting a value of type `Self` to must
> > +/// yield a value with a bit pattern also valid for `Self`.
> > +pub unsafe trait AtomicAdd<Rhs = Self>: AtomicType {
> > +    /// Converts `Rhs` into the `Delta` type of the atomic implementation.
> > +    fn rhs_into_delta(rhs: Rhs) -> <Self::Repr as AtomicImpl>::Delta;
> > +}
> > +
>

Re: [PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations

Posted by Benno Lossin 1 month, 2 weeks ago

On Sat Aug 16, 2025 at 6:10 PM CEST, Boqun Feng wrote:
> On Tue, Aug 12, 2025 at 10:04:12AM +0200, Benno Lossin wrote:
>> On Sat Jul 19, 2025 at 5:08 AM CEST, Boqun Feng wrote:
>> > +/// Types that support atomic add operations.
>> > +///
>> > +/// # Safety
>> > +///
>> > +/// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
>> 
>> Can you add a normal comment TODO here:
>> 
>>     // TODO: properly define `wrapping_add` in this context.
>
> Yeah, this sounds good to me. How do you propose we arrange the normal
> comment with the doc comment, somthing like:
>
>     // TODO: properly define `wrapping_add` in this context.
>     
>     /// Types that support atomic add operations.
>     ///
>     /// # Safety
>     ///
>     /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
>     ...
>     pub unsafe trait AtomicAdd<...> {
>         ...
>     }


Inline maybe?

    /// Types that support atomic add operations.
    ///
    /// # Safety
    ///
    // TODO: properly define `wrapping_add` in this context:
    /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
    /// any value of type `Self::Repr` obtained through transmuting a value of type `Self` to must
    /// yield a value with a bit pattern also valid for `Self`.

---
Cheers,
Benno

Re: [PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations

Posted by Boqun Feng 1 month, 2 weeks ago

On Sat, Aug 16, 2025 at 09:35:26PM +0200, Benno Lossin wrote:
> On Sat Aug 16, 2025 at 6:10 PM CEST, Boqun Feng wrote:
> > On Tue, Aug 12, 2025 at 10:04:12AM +0200, Benno Lossin wrote:
> >> On Sat Jul 19, 2025 at 5:08 AM CEST, Boqun Feng wrote:
> >> > +/// Types that support atomic add operations.
> >> > +///
> >> > +/// # Safety
> >> > +///
> >> > +/// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
> >> 
> >> Can you add a normal comment TODO here:
> >> 
> >>     // TODO: properly define `wrapping_add` in this context.
> >
> > Yeah, this sounds good to me. How do you propose we arrange the normal
> > comment with the doc comment, somthing like:
> >
> >     // TODO: properly define `wrapping_add` in this context.
> >     
> >     /// Types that support atomic add operations.
> >     ///
> >     /// # Safety
> >     ///
> >     /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
> >     ...
> >     pub unsafe trait AtomicAdd<...> {
> >         ...
> >     }
> 
> 
> Inline maybe?
> 
>     /// Types that support atomic add operations.
>     ///
>     /// # Safety
>     ///
>     // TODO: properly define `wrapping_add` in this context:

The colon looks a bit weird to me. I would replace that with a period,
i.e.

     // TODO: properly define `wrapping_add` in the following comment.
     /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to

Thoughts?

Regards,
Boqun

>     /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
>     /// any value of type `Self::Repr` obtained through transmuting a value of type `Self` to must
>     /// yield a value with a bit pattern also valid for `Self`.
> 
> ---
> Cheers,
> Benno
>

Re: [PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations

Posted by Benno Lossin 1 month, 2 weeks ago

On Sun Aug 17, 2025 at 5:04 AM CEST, Boqun Feng wrote:
> On Sat, Aug 16, 2025 at 09:35:26PM +0200, Benno Lossin wrote:
>> On Sat Aug 16, 2025 at 6:10 PM CEST, Boqun Feng wrote:
>> > On Tue, Aug 12, 2025 at 10:04:12AM +0200, Benno Lossin wrote:
>> >> On Sat Jul 19, 2025 at 5:08 AM CEST, Boqun Feng wrote:
>> >> > +/// Types that support atomic add operations.
>> >> > +///
>> >> > +/// # Safety
>> >> > +///
>> >> > +/// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
>> >> 
>> >> Can you add a normal comment TODO here:
>> >> 
>> >>     // TODO: properly define `wrapping_add` in this context.
>> >
>> > Yeah, this sounds good to me. How do you propose we arrange the normal
>> > comment with the doc comment, somthing like:
>> >
>> >     // TODO: properly define `wrapping_add` in this context.
>> >     
>> >     /// Types that support atomic add operations.
>> >     ///
>> >     /// # Safety
>> >     ///
>> >     /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
>> >     ...
>> >     pub unsafe trait AtomicAdd<...> {
>> >         ...
>> >     }
>> 
>> 
>> Inline maybe?
>> 
>>     /// Types that support atomic add operations.
>>     ///
>>     /// # Safety
>>     ///
>>     // TODO: properly define `wrapping_add` in this context:
>
> The colon looks a bit weird to me. I would replace that with a period,
> i.e.
>
>      // TODO: properly define `wrapping_add` in the following comment.
>      /// `wrapping_add` any value of type `Self::Repr::Delta` obtained by [`Self::rhs_into_delta()`] to
>
> Thoughts?

Also works for me :)

---
Cheers,
Benno

[PATCH v8 1/9] rust: Introduce atomic API helpers
[PATCH v8 2/9] rust: sync: Add basic atomic operation mapping framework
[PATCH v8 3/9] rust: sync: atomic: Add ordering annotation types
[PATCH v8 4/9] rust: sync: atomic: Add generic atomics
[PATCH v8 5/9] rust: sync: atomic: Add atomic {cmp,}xchg operations
[PATCH v8 6/9] rust: sync: atomic: Add the framework of arithmetic operations
[PATCH v8 7/9] rust: sync: atomic: Add Atomic<u{32,64}>
[PATCH v8 8/9] rust: sync: atomic: Add Atomic<{usize,isize}>
[PATCH v8 9/9] rust: sync: Add memory barriers