Merge pull request #5923 from dalg24/drop_kokkos_memory_order

Drop legacy `Impl::memory_order*`

core/src/HIP/Kokkos_HIP_Instance.hpp

@@ -22,6 +22,7 @@
 #include <HIP/Kokkos_HIP_Space.hpp>
 #include <HIP/Kokkos_HIP_Error.hpp>
 
+#include <atomic>
 #include <mutex>
 
 namespace Kokkos {

core/src/Kokkos_Atomics_Desul_Wrapper.hpp

@@ -26,7 +26,6 @@ static_assert(false,
 #include <Kokkos_Atomics_Desul_Config.hpp>
 #include <desul/atomics.hpp>
 
-#include <impl/Kokkos_Atomic_Memory_Order.hpp>
 #include <impl/Kokkos_Volatile_Load.hpp>
 
 // clang-format off
@@ -228,51 +227,23 @@ T atomic_compare_exchange(T* const dest, desul::Impl::dont_deduce_this_parameter
 }
 
 namespace Impl {
-
-  template<class MemoryOrder>
-  struct KokkosToDesulMemoryOrder;
-
-  template<>
-  struct KokkosToDesulMemoryOrder<memory_order_seq_cst_t> {
-    using type = desul::MemoryOrderSeqCst;
-  };
-  template<>
-  struct KokkosToDesulMemoryOrder<memory_order_acquire_t> {
-    using type = desul::MemoryOrderAcquire;
-  };
-  template<>
-  struct KokkosToDesulMemoryOrder<memory_order_release_t> {
-    using type = desul::MemoryOrderRelease;
-  };
-  template<>
-  struct KokkosToDesulMemoryOrder<memory_order_acq_rel_t> {
-    using type = desul::MemoryOrderAcqRel;
-  };
-  template<>
-  struct KokkosToDesulMemoryOrder<memory_order_relaxed_t> {
-    using type = desul::MemoryOrderRelaxed;
-  };
   template<class T, class MemOrderSuccess, class MemOrderFailure> KOKKOS_INLINE_FUNCTION
-  bool atomic_compare_exchange_strong(T* const dest, T& expected, const T desired, MemOrderSuccess, MemOrderFailure) {
-    return desul::atomic_compare_exchange_strong(dest, expected, desired,
-                  typename KokkosToDesulMemoryOrder<MemOrderSuccess>::type(),
-                  typename KokkosToDesulMemoryOrder<MemOrderFailure>::type(),
-                  KOKKOS_DESUL_MEM_SCOPE);
-
+  bool atomic_compare_exchange_strong(T* const dest, T& expected, const T desired, MemOrderSuccess succ, MemOrderFailure fail) {
+    return desul::atomic_compare_exchange_strong(dest, expected, desired, succ, fail, KOKKOS_DESUL_MEM_SCOPE);
   }
   template<class T, class MemoryOrder>
   KOKKOS_INLINE_FUNCTION
-  T atomic_load(const T* const src, MemoryOrder) {
-    return desul::atomic_load(src, typename KokkosToDesulMemoryOrder<MemoryOrder>::type(), KOKKOS_DESUL_MEM_SCOPE);
+  T atomic_load(const T* const src, MemoryOrder order) {
+    return desul::atomic_load(src, order, KOKKOS_DESUL_MEM_SCOPE);
   }
   template<class T, class MemoryOrder>
   KOKKOS_INLINE_FUNCTION
-  void atomic_store(T* const src, const T val, MemoryOrder) {
-    return desul::atomic_store(src, val, typename KokkosToDesulMemoryOrder<MemoryOrder>::type(), KOKKOS_DESUL_MEM_SCOPE);
+  void atomic_store(T* const src, const T val, MemoryOrder order) {
+    return desul::atomic_store(src, val, order, KOKKOS_DESUL_MEM_SCOPE);
   }
-}
+}  // namespace Impl
 
-}
+}  // namespace Kokkos
 
 #undef KOKKOS_DESUL_MEM_SCOPE
 

core/src/impl/Kokkos_Atomic_View.hpp

@@ -39,7 +39,7 @@ class AtomicDataElement {
 
   KOKKOS_INLINE_FUNCTION
   const_value_type operator=(const_value_type& val) const {
-    Kokkos::Impl::atomic_store(ptr, val, Kokkos::Impl::memory_order_relaxed);
+    Kokkos::atomic_store(ptr, val);
     return val;
   }
 
@@ -194,9 +194,7 @@ class AtomicDataElement {
   bool operator>(const_value_type& val) const { return *ptr > val; }
 
   KOKKOS_INLINE_FUNCTION
-  operator value_type() const {
-    return Kokkos::Impl::atomic_load(ptr, Kokkos::Impl::memory_order_relaxed);
-  }
+  operator value_type() const { return Kokkos::atomic_load(ptr); }
 };
 
 template <class ViewTraits>

core/src/impl/Kokkos_ChaseLev.hpp

@@ -172,7 +172,8 @@ struct ChaseLevDeque {
         }
 #else
         if (!Impl::atomic_compare_exchange_strong(
-                &m_top, t, t + 1, memory_order_seq_cst, memory_order_relaxed)) {
+                &m_top, t, t + 1, desul::MemoryOrderSeqCst(),
+                desul::MemoryOrderRelaxed())) {
           /* failed race, someone else stole it */
           return_value = nullptr;
         }
@@ -195,7 +196,7 @@ struct ChaseLevDeque {
   KOKKOS_INLINE_FUNCTION
   bool push(node_type& node) {
     auto b  = m_bottom;  // memory order relaxed
-    auto t  = Impl::atomic_load(&m_top, memory_order_acquire);
+    auto t  = Impl::atomic_load(&m_top, desul::MemoryOrderAcquire());
     auto& a = m_array;
     if (b - t > a.size() - 1) {
       /* queue is full, resize */
@@ -204,7 +205,7 @@ struct ChaseLevDeque {
       return false;
     }
     a[b] = &node;  // relaxed
-    Impl::atomic_store(&m_bottom, b + 1, memory_order_release);
+    Impl::atomic_store(&m_bottom, b + 1, desul::MemoryOrderRelease());
     return true;
   }
 
@@ -213,7 +214,7 @@ struct ChaseLevDeque {
     auto t = m_top;  // TODO @tasking @memory_order DSH: atomic load acquire
     Kokkos::memory_fence();  // seq_cst fence, so why does the above need to be
                              // acquire?
-    auto b = Impl::atomic_load(&m_bottom, memory_order_acquire);
+    auto b = Impl::atomic_load(&m_bottom, desul::MemoryOrderAcquire());
     OptionalRef<T> return_value;
     if (t < b) {
       /* Non-empty queue */
@@ -231,8 +232,9 @@ struct ChaseLevDeque {
         return_value = nullptr;
       }
 #else
-      if (!Impl::atomic_compare_exchange_strong(
-              &m_top, t, t + 1, memory_order_seq_cst, memory_order_relaxed)) {
+      if (!Impl::atomic_compare_exchange_strong(&m_top, t, t + 1,
+                                                desul::MemoryOrderSeqCst(),
+                                                desul::MemoryOrderRelaxed())) {
         return_value = nullptr;
       }
 #endif
@@ -247,7 +249,7 @@ struct ChaseLevDeque {
       // essentially using the memory order in this version as a fence, which
       // may be unnecessary
       auto buffer_ptr = (node_type***)&m_array.buffer;
-      auto a = Impl::atomic_load(buffer_ptr, memory_order_acquire); //
+      auto a = Impl::atomic_load(buffer_ptr, desul::MemoryOrderAcquire()); //
    technically consume ordered, but acquire should be fine return_value =
    *static_cast<T*>(a[t % m_array->size]); // relaxed; we'd have to replace the
    m_array->size if we ever allow growth

core/src/impl/Kokkos_HostThreadTeam.hpp

@@ -481,14 +481,14 @@ class HostThreadTeamMember {
         // with a return value of 'true'
 
         Kokkos::Impl::atomic_store(shared_value, value,
-                                   Kokkos::Impl::memory_order_release);
+                                   desul::MemoryOrderRelease());
 
         m_data.team_rendezvous_release();
         // This thread released all other threads from 'team_rendezvous'
         // with a return value of 'false'
       } else {
         value = Kokkos::Impl::atomic_load(shared_value,
-                                          Kokkos::Impl::memory_order_acquire);
+                                          desul::MemoryOrderAcquire());
       }
     }))
 
@@ -516,15 +516,15 @@ class HostThreadTeamMember {
 
           if (1 < m_data.m_team_size) {
             Kokkos::Impl::atomic_store(shared_value, value,
-                                       Kokkos::Impl::memory_order_release);
+                                       desul::MemoryOrderRelease());
           }
 
           m_data.team_rendezvous_release();
           // This thread released all other threads from 'team_rendezvous'
           // with a return value of 'false'
         } else {
           value = Kokkos::Impl::atomic_load(shared_value,
-                                            Kokkos::Impl::memory_order_acquire);
+                                            desul::MemoryOrderAcquire());
         }))
 
     KOKKOS_IF_ON_DEVICE(