src: Add Dave's feedback

philipmarshall21 · Jun 10, 2024 · 8fb3a96 · 8fb3a96
1 parent 99609bf
commit 8fb3a96
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 36 deletions.
diff --git a/src/collectives.c b/src/collectives.c
@@ -613,19 +613,8 @@ shmem_internal_op_to_all_linear(void *target, const void *source, size_t count,
         SHMEM_WAIT_UNTIL(pSync, SHMEM_CMP_EQ, 0);
 
         /* send data, ack, and wait for completion */
-#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;
-        for (size_t i =0; i < count; i++) {
-            shmem_internal_fetch_atomic(SHMEM_CTX_DEFAULT, ((uint8_t *) target) + (i * type_size),
-                                        ((uint8_t *) source) + (i * type_size), &tmp_fetch, type_size,
-                                        PE_start, op, datatype);
-        }
-#else
         shmem_internal_atomicv(SHMEM_CTX_DEFAULT, target, source, count * type_size,
                                PE_start, op, datatype, &completion);
-#endif
         shmem_internal_put_wait(SHMEM_CTX_DEFAULT, &completion);
         shmem_internal_fence(SHMEM_CTX_DEFAULT);
 
@@ -828,21 +817,10 @@ shmem_internal_op_to_all_tree(void *target, const void *source, size_t count, si
         SHMEM_WAIT_UNTIL(pSync + 1, SHMEM_CMP_EQ, 0);
 
         /* send data, ack, and wait for completion */
-#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;
-        for (size_t i = 0; i < count; i++) {
-            shmem_internal_fetch_atomic(SHMEM_CTX_DEFAULT, ((uint8_t *) target) + (i * type_size),
-                                        (num_children == 0) ? ((uint8_t *) source) + (i * type_size) : ((uint8_t *) target) + (i * type_size),
-                                        &tmp_fetch, type_size, parent, op, datatype);
-        }
-#else
         shmem_internal_atomicv(SHMEM_CTX_DEFAULT, target,
                                (num_children == 0) ? source : target,
                                count * type_size, parent,
                                op, datatype, &completion);
-#endif
         shmem_internal_put_wait(SHMEM_CTX_DEFAULT, &completion);
         shmem_internal_fence(SHMEM_CTX_DEFAULT);
 

diff --git a/src/shmem_comm.h b/src/shmem_comm.h
@@ -302,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
 }