Updates to Notifier

src/atomic.hpp

@@ -47,7 +47,6 @@ typedef enum rocshmem_memory_order {
 struct rocshmem_memory_orders {
   rocshmem_memory_order load {memory_order_acquire};
   rocshmem_memory_order store {memory_order_release};
-  rocshmem_memory_order fence {memory_order_acq_rel};
   rocshmem_memory_order atomic {memory_order_acq_rel};
   rocshmem_memory_order weak_cas_success {memory_order_acq_rel};
   rocshmem_memory_order weak_cas_failure {memory_order_acq_rel};
@@ -123,7 +122,7 @@ T fetch_min(T* obj, U arg, rocshmem_memory_orders o) {
 
 template <rocshmem_memory_scope s>
 __device__
-void thread_fence([[maybe_unused]] rocshmem_memory_order order) {
+void thread_fence() {
   if constexpr (s == memory_scope_system) {
     __threadfence_system();
   } else if constexpr (s == memory_scope_agent) {

src/memory/notifier.hpp

@@ -40,15 +40,11 @@ class Notifier {
       detail::atomic::store<uint64_t, scope>(&value_, val, orders);
   }
 
-  __device__ void done() { __syncthreads(); }
+  __device__ void fence() {
+      detail::atomic::thread_fence<scope>();
+  }
 
  private:
-  __device__ void publish() {
-    if (is_thread_zero_in_block()) {
-      __threadfence();
-    }
-    __syncthreads();
-  }
 
   detail::atomic::rocshmem_memory_orders orders;
 

src/memory/slab_heap.cpp

@@ -75,9 +75,12 @@ __device__ void SlabHeap::malloc(void** ptr, size_t size) {
    * Notify other threads in block about the allocation result.
    */
   auto notifier{notifier_.get()};
-  notifier->store(ptr_deref_u64);
+  if (!threadIdx.x) {
+    notifier->store(ptr_deref_u64);
+    notifier->fence();
+  }
+  __syncthreads();
   uint64_t notification_u64{notifier->load()};
-  notifier->done();
 
   /*
    * Write to the ptr parameter (to return it back up the call stack).

tests/unit_tests/notifier_gtest.hpp

@@ -53,10 +53,12 @@ __global__
 void
 all_threads_once(uint8_t* raw_memory,
                  Notifier<detail::atomic::memory_scope_workgroup> * notifier) {
-    notifier->store(NOTIFIER_OFFSET);
+    if (!threadIdx.x) {
+      notifier->store(NOTIFIER_OFFSET);
+      notifier->fence();
+    }
+    __syncthreads();
     uint64_t offset_u64 {notifier->load()};
-    notifier->done();
-
     uint64_t raw_memory_u64 {reinterpret_cast<uint64_t>(raw_memory)};
     uint64_t address_u64 {raw_memory_u64 + offset_u64};
     uint8_t* address {reinterpret_cast<uint8_t*>(address_u64)};