Updates to Notifier

[ROCm/rocshmem commit: 51c33b2a66]

projects/rocshmem/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;
 

projects/rocshmem/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).