Merge pull request #1134 from philipmarshall21/collectives_use_fetchi…

…ng_amo Add DISABLE_NONFETCH_AMO support to op_to_all implementations
Sandia-OpenSHMEM · Jun 11, 2024 · 42cd8ad · 42cd8ad
2 parents fab23f4 + 8fb3a96
commit 42cd8ad
Showing 1 changed file with 30 additions and 14 deletions.
diff --git a/src/shmem_comm.h b/src/shmem_comm.h
@@ -245,6 +245,8 @@ shmem_internal_atomic(shmem_ctx_t ctx, void *target, const void *source, size_t
         shmem_shr_transport_atomic(ctx, target, source, len, pe, op, datatype);
     } else {
 #ifdef DISABLE_NONFETCH_AMO
+        /* FIXME: This is a temporary workaround to resolve a known issue with non-fetching AMOs when using
+           the CXI provider */
         unsigned long long tmp_fetch = 0;
         shmem_transport_fetch_atomic((shmem_transport_ctx_t *)ctx, target,
                                      source, &tmp_fetch, len, pe, op, datatype);
@@ -284,6 +286,8 @@ shmem_internal_atomic_set(shmem_ctx_t ctx, void *target, const void *source, siz
         shmem_shr_transport_atomic_set(ctx, target, source, len, pe, datatype);
     } else {
 #ifdef DISABLE_NONFETCH_AMO
+        /* FIXME: This is a temporary workaround to resolve a known issue with non-fetching AMOs when using
+           the CXI provider */
         unsigned long long tmp_fetch = 0;
         shmem_transport_fetch_atomic((shmem_transport_ctx_t *)ctx, target,
                                      source, &tmp_fetch, len, pe, FI_ATOMIC_WRITE, datatype);
@@ -298,37 +302,49 @@ shmem_internal_atomic_set(shmem_ctx_t ctx, void *target, const void *source, siz
 
 static inline
 void
-shmem_internal_atomicv(shmem_ctx_t ctx, void *target, const void *source,
-                       size_t len, int pe, shm_internal_op_t op,
-                       shm_internal_datatype_t datatype, long *completion)
+shmem_internal_fetch_atomic(shmem_ctx_t ctx, void *target, void *source, void *dest, size_t len,
+                            int pe, shm_internal_op_t op,
+                            shm_internal_datatype_t datatype)
 {
     shmem_internal_assert(len > 0);
 
     if (shmem_shr_transport_use_atomic(ctx, target, len, pe, datatype)) {
-        shmem_shr_transport_atomicv(ctx, target, source, len, pe, op, datatype);
+        shmem_shr_transport_fetch_atomic(ctx, target, source, dest, len, pe,
+                                         op, datatype);
     } else {
-        shmem_transport_atomicv((shmem_transport_ctx_t *)ctx, target, source, len,
-                                pe, op, datatype, completion);
+        shmem_transport_fetch_atomic((shmem_transport_ctx_t *)ctx, target,
+                                     source, dest, len, pe, op, datatype);
     }
 }
 
 
-
 static inline
 void
-shmem_internal_fetch_atomic(shmem_ctx_t ctx, void *target, void *source, void *dest, size_t len,
-                            int pe, shm_internal_op_t op,
-                            shm_internal_datatype_t datatype)
+shmem_internal_atomicv(shmem_ctx_t ctx, void *target, const void *source,
+                       size_t len, int pe, shm_internal_op_t op,
+                       shm_internal_datatype_t datatype, long *completion)
 {
     shmem_internal_assert(len > 0);
 
+#ifdef DISABLE_NONFETCH_AMO
+    /* FIXME: This is a temporary workaround to resolve a known issue with non-fetching AMOs when using
+        the CXI provider */
+    unsigned long long tmp_fetch = 0;
+    size_t type_size = SHMEM_Dtsize[SHMEM_TRANSPORT_DTYPE(datatype)];
+    size_t count = len / type_size;
+    for (size_t i = 0; i < count; i++) {
+        shmem_internal_fetch_atomic(ctx, ((uint8_t *) target) + (i * type_size),
+                                    ((uint8_t *) source) + (i * type_size), &tmp_fetch, type_size,
+                                    pe, op, datatype);
+    }
+#else
     if (shmem_shr_transport_use_atomic(ctx, target, len, pe, datatype)) {
-        shmem_shr_transport_fetch_atomic(ctx, target, source, dest, len, pe,
-                                         op, datatype);
+        shmem_shr_transport_atomicv(ctx, target, source, len, pe, op, datatype);
     } else {
-        shmem_transport_fetch_atomic((shmem_transport_ctx_t *)ctx, target,
-                                     source, dest, len, pe, op, datatype);
+        shmem_transport_atomicv((shmem_transport_ctx_t *)ctx, target, source, len,
+                                pe, op, datatype, completion);
     }
+#endif
 }