rocSHMEM: Add signal operations

projects/rocshmem/CHANGELOG.md

@@ -13,6 +13,14 @@
     improving application performance when such relaxations are known to be safe by the caller.
     These relaxations can be applied independently at each call site.
     This is currently an internal mechanism and is not exposed by the public rocSHMEM APIs.
+* Added standalone signal APIs:
+    * `rocshmem_signal_add` and `rocshmem_ctx_signal_add`
+    * `rocshmem_signal_set` and `rocshmem_ctx_signal_set`
+    * `rocshmem_signal_wait_until`
+
+### Deprecated
+* Deprecated `rocshmem_signal_fetch_wg` and `rocshmem_signal_fetch_wave`.
+  Use `rocshmem_signal_fetch` instead.
 
 ## rocSHMEM 3.7.0 for ROCm 10.1
 

projects/rocshmem/docs/api/sigops.rst

@@ -95,6 +95,22 @@ to ensure completion.
 
 Valid ``sig_op`` values are listed in SIGNAL_OPERATORS_.
 
+ROCSHMEM_SIGNAL_ADD
+-------------------
+
+.. cpp:function:: __device__ void rocshmem_signal_add(uint64_t *sig_addr, uint64_t signal, int pe)
+.. cpp:function:: __device__ void rocshmem_ctx_signal_add(rocshmem_ctx_t ctx, uint64_t *sig_addr, uint64_t signal, int pe)
+
+  :param ctx:      Context with which to perform this operation.
+  :param sig_addr: Symmetric address of the signal data object on the remote PE.
+  :param signal:   Value to add to the signal data object.
+  :param pe:       PE containing the signal data object.
+  :returns:        None.
+
+**Description:**
+This function atomically adds ``signal`` to the signal data object at ``sig_addr``
+on PE ``pe``.
+
 ROCSHMEM_SIGNAL_FETCH
 ---------------------
 
@@ -108,6 +124,42 @@ ROCSHMEM_SIGNAL_FETCH
 **Description:**
 This function atomically fetches the value stored at ``sig_addr``.
 
+.. deprecated:: 3.8.0
+   ``rocshmem_signal_fetch_wg`` and ``rocshmem_signal_fetch_wave`` are
+   deprecated. Use ``rocshmem_signal_fetch`` instead.
+
+ROCSHMEM_SIGNAL_SET
+-------------------
+
+.. cpp:function:: __device__ void rocshmem_signal_set(uint64_t *sig_addr, uint64_t signal, int pe)
+.. cpp:function:: __device__ void rocshmem_ctx_signal_set(rocshmem_ctx_t ctx, uint64_t *sig_addr, uint64_t signal, int pe)
+
+  :param ctx:      Context with which to perform this operation.
+  :param sig_addr: Symmetric address of the signal data object on the remote PE.
+  :param signal:   Value to store in the signal data object.
+  :param pe:       PE containing the signal data object.
+  :returns:        None.
+
+**Description:**
+This function atomically stores ``signal`` in the signal data object at
+``sig_addr`` on PE ``pe``.
+
+ROCSHMEM_SIGNAL_WAIT_UNTIL
+--------------------------
+
+.. cpp:function:: __device__ uint64_t rocshmem_signal_wait_until(uint64_t *sig_addr, int cmp, uint64_t cmp_value)
+
+  :param sig_addr:  Local address of the signal data object.
+  :param cmp:       Comparison operator.
+  :param cmp_value: Value against which to compare the signal.
+  :returns:         The signal value that satisfied the comparison.
+
+**Description:**
+This function blocks until ``(*sig_addr cmp cmp_value)`` is true and returns
+the value that satisfied the comparison.
+
+Valid ``cmp`` values are listed in :ref:`CMP_VALUES`.
+
 Signal operators
 ----------------
 .. _SIGNAL_OPERATORS:

projects/rocshmem/include/rocshmem/rocshmem_SIG_OP.hpp

@@ -615,10 +615,37 @@ __device__ ATTR_NO_INLINE void rocshmem_ulonglong_put_signal_nbi_wave(
     uint64_t signal, int sig_op, int pe);
 
 
+__device__ ATTR_NO_INLINE void rocshmem_ctx_signal_add(
+    rocshmem_ctx_t ctx, uint64_t *sig_addr, uint64_t signal, int pe);
+__device__ ATTR_NO_INLINE void rocshmem_signal_add(
+    uint64_t *sig_addr, uint64_t signal, int pe);
+
 __device__ ATTR_NO_INLINE uint64_t rocshmem_signal_fetch(const uint64_t *sig_addr);
+
+/**
+ * @brief Atomically fetch a signal value collectively at work-group scope.
+ *
+ * @deprecated Use rocshmem_signal_fetch() instead.
+ */
+[[deprecated("Use rocshmem_signal_fetch() instead")]]
 __device__ ATTR_NO_INLINE uint64_t rocshmem_signal_fetch_wg(const uint64_t *sig_addr);
+
+/**
+ * @brief Atomically fetch a signal value collectively at wave scope.
+ *
+ * @deprecated Use rocshmem_signal_fetch() instead.
+ */
+[[deprecated("Use rocshmem_signal_fetch() instead")]]
 __device__ ATTR_NO_INLINE uint64_t rocshmem_signal_fetch_wave(const uint64_t *sig_addr);
 
+__device__ ATTR_NO_INLINE void rocshmem_ctx_signal_set(
+    rocshmem_ctx_t ctx, uint64_t *sig_addr, uint64_t signal, int pe);
+__device__ ATTR_NO_INLINE void rocshmem_signal_set(
+    uint64_t *sig_addr, uint64_t signal, int pe);
+
+__device__ ATTR_NO_INLINE uint64_t rocshmem_signal_wait_until(
+    uint64_t *sig_addr, int cmp, uint64_t cmp_value);
+
 /**
  * @brief Kernel wrapper for putmem_signal operation on stream
  *

projects/rocshmem/scripts/functional_tests/driver.sh

@@ -181,6 +181,9 @@ declare -A TEST_NUMBERS=(
   ["tile_put_wg_colmajor"]="168"
   ["tile_get_wg_rowmajor"]="169"
   ["tile_get_wg_colmajor"]="170"
+  ["signaladd"]="171"
+  ["signalset"]="172"
+  ["signalwaituntil"]="173"
 )
 
 # Detect which runtime to use
@@ -797,6 +800,10 @@ TestSigOps() {
   ExecTest  "wgsignalfetch"    2       2            32
   ExecTest  "wavesignalfetch"  2       1            32
   ExecTest  "wavesignalfetch"  2       1            64
+
+  ExecTest  "signaladd"        2       2            32
+  ExecTest  "signalset"        2       2            32
+  ExecTest  "signalwaituntil"  2       1            1
 }
 
 TestColl() {

projects/rocshmem/src/backend_bc.cpp

@@ -198,7 +198,7 @@ void Backend::dump_stats() {
     if (val) { append("  %-30s %llu\n", name, static_cast<unsigned long long>(val)); ++n_printed; }
   };
 
-  static_assert(NUM_STATS == 69,
+  static_assert(NUM_STATS == 70,
     "rocshmem_stats enum changed; update dump_stats device section");
   const auto& device_stats{globalStats};
   uint64_t device_total = 0;
@@ -237,6 +237,7 @@ void Backend::dump_stats() {
     pstat("WAVE_Sync",             device_stats.getStat(NUM_SYNC_WAVE));
     pstat("WG_Sync",               device_stats.getStat(NUM_SYNC_WG));
     pstat("Wait_Until",            device_stats.getStat(NUM_WAIT_UNTIL));
+    pstat("Signal_Wait_Until",     device_stats.getStat(NUM_SIGNAL_WAIT_UNTIL));
     pstat("Wait_Until_Any",        device_stats.getStat(NUM_WAIT_UNTIL_ANY));
     pstat("Wait_Until_All",        device_stats.getStat(NUM_WAIT_UNTIL_ALL));
     pstat("Wait_Until_Some",       device_stats.getStat(NUM_WAIT_UNTIL_SOME));

projects/rocshmem/src/context.hpp

@@ -110,9 +110,15 @@ class Context {
                                            const int *status,
                                            int cmp, T* vals);
 
+  __device__ uint64_t signal_wait_until(uint64_t *sig_addr, int cmp,
+                                        uint64_t cmp_value);
+
   template <typename T>
   __device__ int test(T *ivars, int cmp, T val);
 
+  template <typename T>
+  __device__ int test_value(T value, int cmp, T cmp_value);
+
   __device__ void threadfence_system();
 
   __device__ void putmem(void* dest, const void* source, size_t nelems, int pe);

projects/rocshmem/src/context_device.cpp

@@ -354,6 +354,16 @@ CONTEXT_SIGNAL_FETCH_DEF()
 CONTEXT_SIGNAL_FETCH_DEF(_wg)
 CONTEXT_SIGNAL_FETCH_DEF(_wave)
 
+__device__ uint64_t Context::signal_wait_until(uint64_t *sig_addr, int cmp,
+                                               uint64_t cmp_value) {
+  while (true) {
+    uint64_t value{uncached_load(sig_addr)};
+    if (test_value(value, cmp, cmp_value)) {
+      return value;
+    }
+  }
+}
+
 __device__ int Context::tile_collective_wait(rocshmem_team_t team, uint64_t flags) {
   DISPATCH_RET(tile_collective_wait(team, flags));
 }

projects/rocshmem/src/context_tmpl_device.hpp

@@ -288,6 +288,27 @@ __device__ void Context::broadcast_wg(T *dest, const T *source, int nelems,
                         pe_size, p_sync));
 }
 
+template <typename T>
+__device__ __forceinline__ int Context::test_value(T value, int cmp,
+                                                   T cmp_value) {
+  switch (cmp) {
+    case ROCSHMEM_CMP_EQ:
+      return value == cmp_value;
+    case ROCSHMEM_CMP_NE:
+      return value != cmp_value;
+    case ROCSHMEM_CMP_GT:
+      return value > cmp_value;
+    case ROCSHMEM_CMP_GE:
+      return value >= cmp_value;
+    case ROCSHMEM_CMP_LT:
+      return value < cmp_value;
+    case ROCSHMEM_CMP_LE:
+      return value <= cmp_value;
+    default:
+      return false;
+  }
+}
+
 template <typename T>
 __device__ __forceinline__ void Context::wait_until(T *ivars, int cmp,
                                                     T val) {
@@ -495,42 +516,7 @@ size_t Context::wait_until_some_vector(T *ivars, size_t nelems,
 template <typename T>
 __device__ __forceinline__ int Context::test(T *ivars, int cmp,
                                              T val) {
-  int ret = 0;
-  switch (cmp) {
-    case ROCSHMEM_CMP_EQ:
-      if (uncached_load(ivars) == val) {
-        ret = 1;
-      }
-      break;
-    case ROCSHMEM_CMP_NE:
-      if (uncached_load(ivars) != val) {
-        ret = 1;
-      }
-      break;
-    case ROCSHMEM_CMP_GT:
-      if (uncached_load(ivars) > val) {
-        ret = 1;
-      }
-      break;
-    case ROCSHMEM_CMP_GE:
-      if (uncached_load(ivars) >= val) {
-        ret = 1;
-      }
-      break;
-    case ROCSHMEM_CMP_LT:
-      if (uncached_load(ivars) < val) {
-        ret = 1;
-      }
-      break;
-    case ROCSHMEM_CMP_LE:
-      if (uncached_load(ivars) <= val) {
-        ret = 1;
-      }
-      break;
-    default:
-      break;
-  }
-  return ret;
+  return test_value(uncached_load(ivars), cmp, val);
 }
 
 template <typename T>

projects/rocshmem/src/device/rocshmem_wrapper.cc

@@ -33,7 +33,7 @@
  * - RMA: put/get/p/g + variants (wave, wg, nbi)
  * - AMO: standard, extended, bitwise
  * - Sync: wait_until variants, test
- * - Signal: put_signal variants
+ * - Signal: put_signal and standalone operations
  *
  * Intentionally excluded (internal use only):
  * - Context methods
@@ -365,6 +365,11 @@ ROCSHMEM_DEVICE_API uint64_t rocshmem_signal_fetch(
   return rocshmem::rocshmem_signal_fetch(sig_addr);
 }
 
+#if defined(__clang__)
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wdeprecated-declarations"
+#endif
+
 ROCSHMEM_DEVICE_API uint64_t rocshmem_signal_fetch_wg(
     const uint64_t *sig_addr) {
   return rocshmem::rocshmem_signal_fetch_wg(sig_addr);
@@ -375,6 +380,25 @@ ROCSHMEM_DEVICE_API uint64_t rocshmem_signal_fetch_wave(
   return rocshmem::rocshmem_signal_fetch_wave(sig_addr);
 }
 
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#endif
+
+ROCSHMEM_DEVICE_API void rocshmem_signal_add(
+    uint64_t *sig_addr, uint64_t signal, int pe) {
+  rocshmem::rocshmem_signal_add(sig_addr, signal, pe);
+}
+
+ROCSHMEM_DEVICE_API void rocshmem_signal_set(
+    uint64_t *sig_addr, uint64_t signal, int pe) {
+  rocshmem::rocshmem_signal_set(sig_addr, signal, pe);
+}
+
+ROCSHMEM_DEVICE_API uint64_t rocshmem_signal_wait_until(
+    uint64_t *sig_addr, int cmp, uint64_t cmp_value) {
+  return rocshmem::rocshmem_signal_wait_until(sig_addr, cmp, cmp_value);
+}
+
 // The explicit instantiation pattern pre-compiles all type variants into bitcode,
 // so JIT linkers don't need to instantiate templates at link time.
 #define WRAP_RMA(T, TNAME)                                                     \

projects/rocshmem/src/rocshmem_gpu.cpp

@@ -1148,6 +1148,13 @@ __device__ __forceinline__ void direct_ctx_wait_until(rocshmem_ctx_t ctx,
   get_base_internal_ctx(ctx)->wait_until(ivars, cmp, val);
 }
 
+__device__ __forceinline__ uint64_t direct_ctx_signal_wait_until(
+    rocshmem_ctx_t ctx, uint64_t *sig_addr, int cmp, uint64_t cmp_value) {
+  get_base_internal_ctx(ctx)->ctxStats.incStat(NUM_SIGNAL_WAIT_UNTIL);
+  return get_base_internal_ctx(ctx)->signal_wait_until(sig_addr, cmp,
+                                                       cmp_value);
+}
+
 template <typename T>
 __device__ __forceinline__ void direct_ctx_wait_until_all(
     rocshmem_ctx_t ctx, T *ivars, size_t nelems, const int *status, int cmp,
@@ -2297,6 +2304,26 @@ ROCSHMEM_SIGNAL_FETCH_DEF()
 ROCSHMEM_SIGNAL_FETCH_DEF(_wg)
 ROCSHMEM_SIGNAL_FETCH_DEF(_wave)
 
+#define ROCSHMEM_SIGNAL_OP_DEF(OP)                                           \
+  __device__ void rocshmem_ctx_signal_##OP(                                  \
+      rocshmem_ctx_t ctx, uint64_t *sig_addr, uint64_t signal, int pe) {     \
+    rocshmem_ctx_uint64_atomic_##OP(ctx, sig_addr, signal, pe);              \
+  }                                                                          \
+                                                                             \
+  __device__ void rocshmem_signal_##OP(                                      \
+      uint64_t *sig_addr, uint64_t signal, int pe) {                         \
+    rocshmem_uint64_atomic_##OP(sig_addr, signal, pe);                       \
+  }
+
+ROCSHMEM_SIGNAL_OP_DEF(add)
+ROCSHMEM_SIGNAL_OP_DEF(set)
+
+__device__ uint64_t rocshmem_signal_wait_until(uint64_t *sig_addr, int cmp,
+                                                uint64_t cmp_value) {
+  return direct_ctx_signal_wait_until(ROCSHMEM_CTX_DEFAULT, sig_addr, cmp,
+                                      cmp_value);
+}
+
 /******************************************************************************
  ****************************** Teams Interface *******************************
  *****************************************************************************/

projects/rocshmem/src/stats.hpp

@@ -52,6 +52,7 @@ enum rocshmem_stats {
   NUM_BARRIER_WAVE,
   NUM_BARRIER_WG,
   NUM_WAIT_UNTIL,
+  NUM_SIGNAL_WAIT_UNTIL,
   NUM_WAIT_UNTIL_ANY,
   NUM_WAIT_UNTIL_ALL,
   NUM_WAIT_UNTIL_SOME,

projects/rocshmem/tests/functional_tests/CTestFunctionalTests.cmake

@@ -182,6 +182,9 @@ set(TEST_tile_put_wg_rowmajor 167)
 set(TEST_tile_put_wg_colmajor 168)
 set(TEST_tile_get_wg_rowmajor 169)
 set(TEST_tile_get_wg_colmajor 170)
+set(TEST_signaladd 171)
+set(TEST_signalset 172)
+set(TEST_signalwaituntil 173)
 
 # MPI should already be found by the parent CMakeLists.txt
 # Use standard CMake MPI variables set by find_package(MPI)
@@ -1034,6 +1037,10 @@ function(add_sigops_tests)
         add_rocshmem_functional_test(NAME wgsignalfetch RANKS 2 WORKGROUPS 2 THREADS 32)
         add_rocshmem_functional_test(NAME wavesignalfetch RANKS 2 WORKGROUPS 1 THREADS 32)
         add_rocshmem_functional_test(NAME wavesignalfetch RANKS 2 WORKGROUPS 1 THREADS 64)
+
+        add_rocshmem_functional_test(NAME signaladd RANKS 2 WORKGROUPS 2 THREADS 32)
+        add_rocshmem_functional_test(NAME signalset RANKS 2 WORKGROUPS 2 THREADS 32)
+        add_rocshmem_functional_test(NAME signalwaituntil RANKS 2 WORKGROUPS 1 THREADS 1)
     end_test_group()
 endfunction()