Replace volatile m_ready_count loads by desul::atomic_loads in task q…

…ueue implementation
nliber · Aug 18, 2023 · 7a24e94 · 7a24e94
1 parent 04d5c55
commit 7a24e94
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 7 deletions.
diff --git a/core/src/HPX/Kokkos_HPX_Task.hpp b/core/src/HPX/Kokkos_HPX_Task.hpp
@@ -164,7 +164,9 @@ class TaskQueueSpecializationConstrained<
         team_queue.complete(task);
       }
 
-      if (*((volatile int *)&team_queue.m_ready_count) > 0) {
+      if (desul::atomic_load(&team_queue.m_ready_count,
+                             desul::MemoryOrderAcquire(),
+                             desul::MemoryScopeDevice()) > 0) {
         task = end;
         for (int i = 0; i < queue_type::NumQueue && end == task; ++i) {
           for (int j = 0; j < 2 && end == task; ++j) {

diff --git a/core/src/OpenMP/Kokkos_OpenMP_Task.hpp b/core/src/OpenMP/Kokkos_OpenMP_Task.hpp
@@ -287,7 +287,9 @@ class TaskQueueSpecializationConstrained<
 
               // If 0 == m_ready_count then set task = 0
 
-              if (*((volatile int*)&team_queue.m_ready_count) > 0) {
+              if (desul::atomic_load(&team_queue.m_ready_count,
+                                     desul::MemoryOrderAcquire(),
+                                     desul::MemoryScopeDevice()) > 0) {
                 task = end;
                 // Attempt to acquire a task
                 // Loop by priority and then type

diff --git a/core/src/impl/Kokkos_TaskQueueCommon.hpp b/core/src/impl/Kokkos_TaskQueueCommon.hpp
@@ -143,14 +143,14 @@ class TaskQueueCommonMixin {
  public:
   KOKKOS_INLINE_FUNCTION
   bool is_done() const noexcept {
-    // TODO @tasking @memory_order DSH Memory order, instead of volatile
-    return (*(volatile int*)(&m_ready_count)) == 0;
+    return desul::atomic_load(&m_ready_count, desul::MemoryOrderAcquire(),
+                              desul::MemoryScopeDevice()) == 0;
   }
 
   KOKKOS_INLINE_FUNCTION
   int32_t ready_count() const noexcept {
-    // TODO @tasking @memory_order DSH Memory order, instead of volatile
-    return (*(volatile int*)(&m_ready_count));
+    return desul::atomic_load(&m_ready_count, desul::MemoryOrderAcquire(),
+                              desul::MemoryScopeDevice());
   }
 
   template <class TaskQueueTraits, class TeamSchedulerInfo>

diff --git a/core/src/impl/Kokkos_TaskQueueMultiple.hpp b/core/src/impl/Kokkos_TaskQueueMultiple.hpp
@@ -119,7 +119,9 @@ class TaskQueueMultiple : public TaskQueue<ExecSpace, MemorySpace> {
         for (int iteam = 0; iteam < m_other_queues->size(); ++iteam) {
           if (iteam == m_league_rank) continue;
           auto& steal_from = get_team_queue(iteam);
-          if (*((volatile int*)&steal_from.m_ready_count) > 0) {
+          if (desul::atomic_load(&steal_from.m_ready_count,
+                                 desul::MemoryOrderAcquire(),
+                                 desul::MemoryScopeDevice()) > 0) {
             // we've found at least one queue that's not done, so even if we
             // can't pop something off of it we shouldn't return a nullptr
             // indicating completion.  rv will be end_tag when the pop fails