From 16941b8d78650fcb5cf2b7053a2422d9bf7ad4c5 Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Sat, 21 Jan 2023 18:05:10 +0900
Subject: [PATCH] Optimize Atomic{I,U}*::{fetch_not,not}

---
 CHANGELOG.md                    |  2 ++
 src/imp/atomic128/aarch64.rs    | 32 ++++++++++++++++++++++++
 src/imp/atomic128/intrinsics.rs | 44 ++++++++++++++++++++-------------
 src/imp/atomic128/macros.rs     | 34 +++++++++++++------------
 src/imp/atomic128/powerpc64.rs  |  6 +++++
 src/imp/core_atomic.rs          | 27 +++++++++++++++-----
 6 files changed, 106 insertions(+), 39 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a554f155..98f5345e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@ Note: In this file, do not use the hard wrap in the middle of a sentence for com
 
 ## [Unreleased]
 
+- Optimize `Atomic{I,U}*::{fetch_not,not}` methods.
+
 ## [1.0.0] - 2023-01-15
 
 - Add `critical-section` feature to use [critical-section](https://github.com/rust-embedded/critical-section) on targets where atomic CAS is not natively available. ([#51](https://github.com/taiki-e/portable-atomic/pull/51), thanks @Dirbaio)
diff --git a/src/imp/atomic128/aarch64.rs b/src/imp/atomic128/aarch64.rs
index 39a0a4d0..9ee0f613 100644
--- a/src/imp/atomic128/aarch64.rs
+++ b/src/imp/atomic128/aarch64.rs
@@ -662,6 +662,38 @@ unsafe fn atomic_xor(dst: *mut u128, val: u128, order: Ordering) -> u128 {
     }
 }
 
+#[inline]
+unsafe fn atomic_not(dst: *mut u128, order: Ordering) -> u128 {
+    debug_assert!(dst as usize % 16 == 0);
+
+    // SAFETY: the caller must uphold the safety contract for `atomic_not`.
+    unsafe {
+        let (mut prev_lo, mut prev_hi);
+        macro_rules! not {
+            ($acquire:tt, $release:tt) => {
+                asm!(
+                    "2:",
+                        concat!("ld", $acquire, "xp {prev_lo}, {prev_hi}, [{dst", ptr_modifier!(), "}]"),
+                        "mvn {tmp_lo}, {prev_lo}",
+                        "mvn {tmp_hi}, {prev_hi}",
+                        concat!("st", $release, "xp {r:w}, {tmp_lo}, {tmp_hi}, [{dst", ptr_modifier!(), "}]"),
+                        // 0 if the store was successful, 1 if no store was performed
+                        "cbnz {r:w}, 2b",
+                    dst = in(reg) dst,
+                    prev_lo = out(reg) prev_lo,
+                    prev_hi = out(reg) prev_hi,
+                    tmp_lo = out(reg) _,
+                    tmp_hi = out(reg) _,
+                    r = out(reg) _,
+                    options(nostack, preserves_flags),
+                )
+            };
+        }
+        atomic_rmw!(not, order);
+        U128 { pair: Pair { lo: prev_lo, hi: prev_hi } }.whole
+    }
+}
+
 #[inline]
 unsafe fn atomic_max(dst: *mut i128, val: i128, order: Ordering) -> i128 {
     debug_assert!(dst as usize % 16 == 0);
diff --git a/src/imp/atomic128/intrinsics.rs b/src/imp/atomic128/intrinsics.rs
index ba9b75ea..f28fb1a9 100644
--- a/src/imp/atomic128/intrinsics.rs
+++ b/src/imp/atomic128/intrinsics.rs
@@ -405,7 +405,7 @@ unsafe fn atomic_umin(dst: *mut u128, val: u128, order: Ordering) -> u128 {
 }
 
 macro_rules! atomic128 {
-    (uint, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
+    (int_general, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
         #[repr(C, align(16))]
         pub(crate) struct $atomic_type {
             v: UnsafeCell<$int_type>,
@@ -583,13 +583,36 @@ macro_rules! atomic128 {
                 unsafe { $atomic_min(self.v.get(), val, order) }
             }
 
+            #[inline]
+            pub(crate) fn not(&self, order: Ordering) {
+                self.fetch_not(order);
+            }
+        }
+    };
+    (uint, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
+        atomic128!(int_general, $atomic_type, $int_type, $atomic_max, $atomic_min);
+        impl $atomic_type {
             #[inline]
             pub(crate) fn fetch_not(&self, order: Ordering) -> $int_type {
-                self.fetch_update_(order, |x| !x)
+                self.fetch_xor(core::$int_type::MAX, order)
             }
+        }
+    };
+    (int, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
+        atomic128!(int_general, $atomic_type, $int_type, $atomic_max, $atomic_min);
+        impl $atomic_type {
             #[inline]
-            pub(crate) fn not(&self, order: Ordering) {
-                self.fetch_not(order);
+            pub(crate) fn fetch_not(&self, order: Ordering) -> $int_type {
+                self.fetch_xor(-1, order)
+            }
+
+            #[inline]
+            pub(crate) fn fetch_neg(&self, order: Ordering) -> $int_type {
+                self.fetch_update_(order, |x| x.wrapping_neg())
+            }
+            #[inline]
+            pub(crate) fn neg(&self, order: Ordering) {
+                self.fetch_neg(order);
             }
 
             #[inline]
@@ -609,19 +632,6 @@ macro_rules! atomic128 {
             }
         }
     };
-    (int, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
-        atomic128!(uint, $atomic_type, $int_type, $atomic_max, $atomic_min);
-        impl $atomic_type {
-            #[inline]
-            pub(crate) fn fetch_neg(&self, order: Ordering) -> $int_type {
-                self.fetch_update_(order, |x| x.wrapping_neg())
-            }
-            #[inline]
-            pub(crate) fn neg(&self, order: Ordering) {
-                self.fetch_neg(order);
-            }
-        }
-    };
 }
 
 atomic128!(int, AtomicI128, i128, atomic_max, atomic_min);
diff --git a/src/imp/atomic128/macros.rs b/src/imp/atomic128/macros.rs
index 79b668f8..c4408976 100644
--- a/src/imp/atomic128/macros.rs
+++ b/src/imp/atomic128/macros.rs
@@ -424,13 +424,29 @@ macro_rules! atomic128 {
 
             #[inline]
             pub(crate) fn fetch_not(&self, order: Ordering) -> $int_type {
-                // TODO: define atomic_not function and use it
-                self.fetch_update_(order, |x| !x)
+                crate::utils::assert_swap_ordering(order);
+                // SAFETY: any data races are prevented by atomic intrinsics and the raw
+                // pointer passed in is valid because we got it from a reference.
+                unsafe { atomic_not(self.v.get().cast(), order) as $int_type }
             }
             #[inline]
             pub(crate) fn not(&self, order: Ordering) {
                 self.fetch_not(order);
             }
+        }
+    };
+    (int, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
+        atomic128!(uint, $atomic_type, $int_type, $atomic_max, $atomic_min);
+        impl $atomic_type {
+            #[inline]
+            pub(crate) fn fetch_neg(&self, order: Ordering) -> $int_type {
+                // TODO: define atomic_neg function and use it
+                self.fetch_update_(order, |x| x.wrapping_neg())
+            }
+            #[inline]
+            pub(crate) fn neg(&self, order: Ordering) {
+                self.fetch_neg(order);
+            }
 
             #[inline]
             fn fetch_update_<F>(&self, set_order: Ordering, mut f: F) -> $int_type
@@ -449,18 +465,4 @@ macro_rules! atomic128 {
             }
         }
     };
-    (int, $atomic_type:ident, $int_type:ident, $atomic_max:ident, $atomic_min:ident) => {
-        atomic128!(uint, $atomic_type, $int_type, $atomic_max, $atomic_min);
-        impl $atomic_type {
-            #[inline]
-            pub(crate) fn fetch_neg(&self, order: Ordering) -> $int_type {
-                // TODO: define atomic_neg function and use it
-                self.fetch_update_(order, |x| x.wrapping_neg())
-            }
-            #[inline]
-            pub(crate) fn neg(&self, order: Ordering) {
-                self.fetch_neg(order);
-            }
-        }
-    };
 }
diff --git a/src/imp/atomic128/powerpc64.rs b/src/imp/atomic128/powerpc64.rs
index 5a358005..c3faf776 100644
--- a/src/imp/atomic128/powerpc64.rs
+++ b/src/imp/atomic128/powerpc64.rs
@@ -507,6 +507,12 @@ unsafe fn atomic_xor(dst: *mut u128, val: u128, order: Ordering) -> u128 {
     }
 }
 
+#[inline]
+unsafe fn atomic_not(dst: *mut u128, order: Ordering) -> u128 {
+    // SAFETY: the caller must uphold the safety contract for `atomic_not`.
+    unsafe { atomic_xor(dst, core::u128::MAX, order) }
+}
+
 #[inline]
 unsafe fn atomic_max(dst: *mut i128, val: i128, order: Ordering) -> i128 {
     debug_assert!(dst as usize % 16 == 0);
diff --git a/src/imp/core_atomic.rs b/src/imp/core_atomic.rs
index 7d336670..d82746eb 100644
--- a/src/imp/core_atomic.rs
+++ b/src/imp/core_atomic.rs
@@ -169,7 +169,7 @@ impl<T> core::ops::DerefMut for AtomicPtr<T> {
 }
 
 macro_rules! atomic_int {
-    (uint, $atomic_type:ident, $int_type:ident) => {
+    (int_general, $atomic_type:ident, $int_type:ident) => {
         #[repr(transparent)]
         pub(crate) struct $atomic_type {
             inner: core::sync::atomic::$atomic_type,
@@ -244,6 +244,7 @@ macro_rules! atomic_int {
                 let success = crate::utils::upgrade_success_ordering(success, failure);
                 self.inner.compare_exchange_weak(current, new, success, failure)
             }
+            #[allow(dead_code)]
             #[inline]
             fn fetch_update_<F>(&self, set_order: Ordering, mut f: F) -> $int_type
             where
@@ -347,10 +348,6 @@ macro_rules! atomic_int {
                     self.fetch_update_(order, |x| core::cmp::min(x, val))
                 }
             }
-            #[inline]
-            pub(crate) fn fetch_not(&self, order: Ordering) -> $int_type {
-                self.fetch_update_(order, |x| !x)
-            }
             #[cfg(not(all(
                 not(any(miri, portable_atomic_sanitize_thread)),
                 any(not(portable_atomic_no_asm), portable_atomic_unstable_asm),
@@ -375,14 +372,32 @@ macro_rules! atomic_int {
             }
         }
     };
+    (uint, $atomic_type:ident, $int_type:ident) => {
+        atomic_int!(int_general, $atomic_type, $int_type);
+        #[cfg_attr(
+            portable_atomic_no_cfg_target_has_atomic,
+            cfg(not(portable_atomic_no_atomic_cas))
+        )]
+        #[cfg_attr(not(portable_atomic_no_cfg_target_has_atomic), cfg(target_has_atomic = "ptr"))]
+        impl $atomic_type {
+            #[inline]
+            pub(crate) fn fetch_not(&self, order: Ordering) -> $int_type {
+                self.fetch_xor(core::$int_type::MAX, order)
+            }
+        }
+    };
     (int, $atomic_type:ident, $int_type:ident) => {
-        atomic_int!(uint, $atomic_type, $int_type);
+        atomic_int!(int_general, $atomic_type, $int_type);
         #[cfg_attr(
             portable_atomic_no_cfg_target_has_atomic,
             cfg(not(portable_atomic_no_atomic_cas))
         )]
         #[cfg_attr(not(portable_atomic_no_cfg_target_has_atomic), cfg(target_has_atomic = "ptr"))]
         impl $atomic_type {
+            #[inline]
+            pub(crate) fn fetch_not(&self, order: Ordering) -> $int_type {
+                self.fetch_xor(-1, order)
+            }
             #[inline]
             pub(crate) fn fetch_neg(&self, order: Ordering) -> $int_type {
                 self.fetch_update_(order, |x| x.wrapping_neg())