From 9b0e4dc05dd187f61daca4805d5928b402b4fe00 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Thu, 15 Aug 2024 12:22:38 -0700
Subject: [PATCH] Change notifier fixture to prep for other fixtures

[ROCm/rocshmem commit: 1289d50be5e213a10269926cdf5219ab9dca5a82]
---
 .../tests/unit_tests/notifier_gtest.cpp       | 14 ++--
 .../tests/unit_tests/notifier_gtest.hpp       | 64 +++++++++----------
 2 files changed, 38 insertions(+), 40 deletions(-)

diff --git a/projects/rocshmem/tests/unit_tests/notifier_gtest.cpp b/projects/rocshmem/tests/unit_tests/notifier_gtest.cpp
index e1ebea0210..073c37a495 100644
--- a/projects/rocshmem/tests/unit_tests/notifier_gtest.cpp
+++ b/projects/rocshmem/tests/unit_tests/notifier_gtest.cpp
@@ -29,29 +29,29 @@ using namespace rocshmem;
  *****************************************************************************/
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_1_1) {
-  run_all_threads_once_block(1, 1);
+  run_all_threads_once(1, 1);
 }
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_2_1) {
-  run_all_threads_once_block(2, 1);
+  run_all_threads_once(2, 1);
 }
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_64_1) {
-  run_all_threads_once_block(64, 1);
+  run_all_threads_once(64, 1);
 }
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_128_1) {
-  run_all_threads_once_block(128, 1);
+  run_all_threads_once(128, 1);
 }
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_256_1) {
-  run_all_threads_once_block(256, 1);
+  run_all_threads_once(256, 1);
 }
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_512_1) {
-  run_all_threads_once_block(512, 1);
+  run_all_threads_once(512, 1);
 }
 
 TEST_F(NotifierBlockTestFixture, run_all_threads_once_1024_1) {
-  run_all_threads_once_block(1024, 1);
+  run_all_threads_once(1024, 1);
 }
diff --git a/projects/rocshmem/tests/unit_tests/notifier_gtest.hpp b/projects/rocshmem/tests/unit_tests/notifier_gtest.hpp
index 453b93680d..b82f47b1d9 100644
--- a/projects/rocshmem/tests/unit_tests/notifier_gtest.hpp
+++ b/projects/rocshmem/tests/unit_tests/notifier_gtest.hpp
@@ -43,17 +43,18 @@ static const uint64_t NOTIFIER_OFFSET {0x100B00};
 
 inline __device__
 void
-write_to_memory_block(uint8_t* raw_memory) {
-    auto thread_idx {get_flat_block_id()};
+write_to_memory(uint8_t* raw_memory) {
+    auto thread_idx {get_flat_id()};
     raw_memory[thread_idx] = THREAD_VALUE;
     __threadfence();
 }
 
+template <typename NotifierT>
 __global__
 void
-all_threads_once_block(uint8_t* raw_memory,
-                       Notifier<detail::atomic::memory_scope_workgroup> * notifier) {
-    if (!threadIdx.x) {
+all_threads_once(uint8_t* raw_memory,
+                 Notifier<detail::atomic::memory_scope_workgroup> * notifier) {
+    if (!get_flat_id()) {
       notifier->store(NOTIFIER_OFFSET);
       notifier->fence();
     }
@@ -62,49 +63,26 @@ all_threads_once_block(uint8_t* raw_memory,
     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)};
-    write_to_memory_block(address);
+    write_to_memory(address);
 }
 
-class NotifierBlockTestFixture : public ::testing::Test {
-    using NotifierProxyT = NotifierProxy<HIPAllocator, detail::atomic::memory_scope_workgroup>;
-
+class NotifierBase : public ::testing::Test {
   public:
-    NotifierBlockTestFixture() {
+    NotifierBase() {
         assert(raw_memory_ == nullptr);
         hip_allocator_.allocate((void**)&raw_memory_, GIBIBYTE_);
         assert(raw_memory_);
     }
 
-    ~NotifierBlockTestFixture() {
+    ~NotifierBase() {
         if (raw_memory_) {
             hip_allocator_.deallocate(raw_memory_);
         }
     }
 
     void
-    run_all_threads_once_block(uint32_t x_block_dim,
-                               uint32_t x_grid_dim) {
-        const dim3 hip_blocksize(x_block_dim, 1, 1);
-        const dim3 hip_gridsize(x_grid_dim, 1, 1);
-
-        hipLaunchKernelGGL(all_threads_once_block,
-                           hip_gridsize,
-                           hip_blocksize,
-                           0,
-                           nullptr,
-                           raw_memory_,
-                           notifier_.get());
-
-        hipError_t return_code = hipStreamSynchronize(nullptr);
-        if (return_code != hipSuccess) {
-            printf("Failed in stream synchronize\n");
-            assert(return_code == hipSuccess);
-        }
-
-        size_t number_threads {x_block_dim * x_grid_dim};
-
+    verify(size_t number_threads) {
         uint8_t* offset_addr {compute_offset_addr()};
-
         for (size_t i {0}; i < number_threads; i++) {
             ASSERT_EQ(offset_addr[i], THREAD_VALUE);
         }
@@ -137,6 +115,26 @@ class NotifierBlockTestFixture : public ::testing::Test {
      */
     uint8_t *raw_memory_ {nullptr};
 
+};
+
+class NotifierBlockTestFixture : public NotifierBase {
+    using NotifierT = Notifier<detail::atomic::memory_scope_workgroup>;
+    using NotifierProxyT = NotifierProxy<HIPAllocator, detail::atomic::memory_scope_workgroup>;
+
+  public:
+    void
+    run_all_threads_once(uint32_t x_block_dim,
+                         uint32_t x_grid_dim) {
+        const dim3 block(x_block_dim, 1, 1);
+        const dim3 grid(x_grid_dim, 1, 1);
+
+        all_threads_once<NotifierT><<<grid, block>>>(raw_memory_, notifier_.get());
+
+        CHECK_HIP(hipStreamSynchronize(nullptr));
+
+        verify(x_block_dim * x_grid_dim);
+    }
+
     /**
      * @brief Used to broadcast base offset for writing.
      */