Fix for UCX tests

Sandia-OpenSHMEM · Jan 16, 2025 · dd5fca9 · dd5fca9
1 parent caffaca
commit dd5fca9
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 8 deletions.
diff --git a/src/shmem_comm.h b/src/shmem_comm.h
@@ -346,7 +346,7 @@ shmem_internal_atomicv(shmem_ctx_t ctx, void *target, const void *source,
                                        pe, op, datatype);
         }
     } else {
-        shmem_transport_atomicv((shmem_transport_ctx_t *)ctx, target, source, len,
+        shmem_transport_atomicv((shmem_transport_ctx_t *)ctx, target, source, count, type_size,
                                 pe, op, datatype, completion);
     }
 #endif

diff --git a/src/transport_none.h b/src/transport_none.h
@@ -216,7 +216,7 @@ shmem_transport_atomic(shmem_transport_ctx_t* ctx, void *target, const void *sou
 
 static inline
 void
-shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source, size_t len,
+shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source, size_t count, size_t type_size,
                         int pe, shm_internal_op_t op, shm_internal_datatype_t datatype, long *completion)
 {
     RAISE_ERROR_STR("No path to peer");

diff --git a/src/transport_ofi.h b/src/transport_ofi.h
@@ -1166,9 +1166,10 @@ void shmem_transport_atomic(shmem_transport_ctx_t* ctx, void *target, const void
 
 static inline
 void shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source,
-                             size_t full_len, int pe, int op, int datatype,
+                             size_t count, size_t type_size, int pe, int op, int datatype,
                              long *completion)
 {
+    size_t full_len = count*type_size;
     int ret = 0;
     uint64_t dst = (uint64_t) pe;
     int dt = SHMEM_TRANSPORT_DTYPE(datatype);

diff --git a/src/transport_portals4.h b/src/transport_portals4.h
@@ -892,9 +892,10 @@ shmem_transport_atomic(shmem_transport_ctx_t* ctx, void *target, const void *sou
 
 static inline
 void
-shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source, size_t len, int pe,
+shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source, size_t count, size_t type_size, int pe,
                         ptl_op_t op, ptl_datatype_t datatype, long *completion)
 {
+    size_t len = count*type_size;
     int ret;
     ptl_pt_index_t pt;
     long offset;

diff --git a/src/transport_ucx.h b/src/transport_ucx.h
@@ -519,12 +519,14 @@ shmem_transport_atomic(shmem_transport_ctx_t* ctx, void *target, const void *sou
 
 static inline
 void
-shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source, size_t len,
+shmem_transport_atomicv(shmem_transport_ctx_t* ctx, void *target, const void *source, size_t count, size_t type_size,
                         int pe, shm_internal_op_t op, shm_internal_datatype_t datatype, long *completion)
 {
-    /* Used only by reductions, currently redirected to softwre reductions via
-     * the shmem_transport_atomic_supported query below. */
-    RAISE_ERROR_STR("Unsupported operation");
+    for (size_t i = 0; i < count; i++) {
+            shmem_transport_atomic(ctx, ((uint8_t *) target) + (i * type_size),
+                                   ((uint8_t *) source) + (i * type_size), type_size,
+                                   pe, op, datatype);
+    }
 }
 
 static inline