Merge pull request #9488 from wfaderhold21/topic/scoll_ucc_a2a

scoll/ucc: adjust count/dtype for alltoall
open-mpi · Oct 19, 2021 · 0a8942b · 0a8942b
2 parents ff6b034 + 3d96daf
commit 0a8942b
Showing 1 changed file with 21 additions and 10 deletions.
diff --git a/oshmem/mca/scoll/ucc/scoll_ucc_alltoall.c b/oshmem/mca/scoll/ucc/scoll_ucc_alltoall.c
@@ -12,24 +12,34 @@
 
 #include <ucc/api/ucc.h>
 
-static inline ucc_status_t mca_scoll_ucc_alltoall_init(const void *sbuf, void *rbuf,
-                                                       int count,
-                                                       mca_scoll_ucc_module_t * ucc_module,
-                                                       ucc_coll_req_h * req)
+static inline ucc_status_t mca_scoll_ucc_alltoall_init(const void *sbuf, void *rbuf, 
+                                                       int count, size_t elem_size,
+                                                       mca_scoll_ucc_module_t *ucc_module,
+                                                       ucc_coll_req_h *req)
 {
+    ucc_datatype_t dt;
+
+    if (elem_size == 8) {
+        dt = UCC_DT_INT64;
+    } else if (elem_size == 4) {
+        dt = UCC_DT_INT32;
+    } else {
+        dt = UCC_DT_INT8;
+    }
+
     ucc_coll_args_t coll = {
         .mask = 0,
         .coll_type = UCC_COLL_TYPE_ALLTOALL,
         .src.info = {
             .buffer = (void *)sbuf,
-            .count = count,
-            .datatype = UCC_DT_UINT8,
+            .count = count * ucc_module->group->proc_count,
+            .datatype = dt,
             .mem_type = UCC_MEMORY_TYPE_UNKNOWN
         },
         .dst.info = {
             .buffer = rbuf,
-            .count = count,
-            .datatype = UCC_DT_UINT8,
+            .count = count * ucc_module->group->proc_count,
+            .datatype = dt,
             .mem_type = UCC_MEMORY_TYPE_UNKNOWN
         },
     };
@@ -56,14 +66,15 @@ int mca_scoll_ucc_alltoall(struct oshmem_group_t *group,
 
     UCC_VERBOSE(3, "running ucc alltoall");
     ucc_module = (mca_scoll_ucc_module_t *) group->g_scoll.scoll_alltoall_module;
-    count = nelems * element_size;
+    count = nelems;
 
     /* Do nothing on zero-length request */
     if (OPAL_UNLIKELY(!nelems)) {
         return OSHMEM_SUCCESS;
     }
 
-    SCOLL_UCC_CHECK(mca_scoll_ucc_alltoall_init(source, target, count, ucc_module, &req));
+    SCOLL_UCC_CHECK(mca_scoll_ucc_alltoall_init(source, target, count, 
+                                                element_size, ucc_module, &req));
     SCOLL_UCC_CHECK(ucc_collective_post(req));
     SCOLL_UCC_CHECK(scoll_ucc_req_wait(req));
     return OSHMEM_SUCCESS;