TEST: do local check in mpi tests (#789)

* TEST: do local check in mpi tests * REVIEW: fix review comments
openucx · Nov 1, 2023 · 7aefb67 · 7aefb67
1 parent 483b91b
commit 7aefb67
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 34 deletions.
diff --git a/test/mpi/buffer.cc b/test/mpi/buffer.cc
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -25,7 +25,7 @@ void init_buffer_host(void *buf, size_t count, int _value)
 }
 
 void init_buffer(void *_buf, size_t count, ucc_datatype_t dt,
-                 ucc_memory_type_t mt, int value)
+                 ucc_memory_type_t mt, int value, int offset)
 {
     void *buf = NULL;
     if (mt == UCC_MEMORY_TYPE_CUDA || mt == UCC_MEMORY_TYPE_ROCM) {
@@ -37,6 +37,8 @@ void init_buffer(void *_buf, size_t count, ucc_datatype_t dt,
         std::cerr << "Unsupported mt\n";
         MPI_Abort(MPI_COMM_WORLD, -1);
     }
+
+    value += offset;
     switch(dt) {
     case UCC_DT_INT8:
         init_buffer_host<int8_t>(buf, count, value);

diff --git a/test/mpi/test_allgather.cc b/test/mpi/test_allgather.cc
@@ -53,6 +53,7 @@ ucc_status_t TestAllgather::set_input(int iter_persistent)
     int    rank;
     void  *buf, *check;
 
+    this->iter_persistent = iter_persistent;
     MPI_Comm_rank(team.comm, &rank);
     if (inplace) {
         buf   = PTR_OFFSET(rbuf, rank * single_rank_size);
@@ -70,18 +71,16 @@ ucc_status_t TestAllgather::set_input(int iter_persistent)
 
 ucc_status_t TestAllgather::check()
 {
-    int size, completed;
-    MPI_Comm_size(team.comm, &size);
-    size_t       single_rank_count = args.dst.info.count / size;
-    MPI_Datatype mpi_dt            = ucc_dt_to_mpi(dt);
-    MPI_Request  req;
+    size_t dt_size, single_rank_count;
+    int    size, i;
 
-    MPI_Iallgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, check_buf,
-                   single_rank_count, mpi_dt, team.comm, &req);
-    do {
-        MPI_Test(&req, &completed, MPI_STATUS_IGNORE);
-        ucc_context_progress(team.ctx);
-    } while(!completed);
+    MPI_Comm_size(team.comm, &size);
+    single_rank_count = args.dst.info.count / size;
+    dt_size = ucc_dt_size(dt);
+    for (i = 0; i < size; i++) {
+        init_buffer(PTR_OFFSET(check_buf, i * single_rank_count * dt_size),
+                    single_rank_count, dt, mem_type, i * (iter_persistent + 1));
+    }
 
     return compare_buffers(rbuf, check_buf, single_rank_count * size, dt,
                            mem_type);

diff --git a/test/mpi/test_alltoall.cc b/test/mpi/test_alltoall.cc
@@ -74,6 +74,7 @@ ucc_status_t TestAlltoall::set_input(int iter_persistent)
     void *      buf;
     int         rank, nprocs, completed;
 
+    this->iter_persistent = iter_persistent;
     MPI_Comm_rank(team.comm, &rank);
     MPI_Comm_size(team.comm, &nprocs);
     if (inplace) {
@@ -99,19 +100,18 @@ ucc_status_t TestAlltoall::set_input(int iter_persistent)
 
 ucc_status_t TestAlltoall::check()
 {
-    int         size, completed;
-    size_t      single_rank_count;
-    MPI_Request req;
+    int    size, rank, i;
+    size_t single_rank_count;
 
+    MPI_Comm_rank(team.comm, &rank);
     MPI_Comm_size(team.comm, &size);
     single_rank_count = args.src.info.count / size;
 
-    MPI_Ialltoall(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, check_buf,
-                  single_rank_count, ucc_dt_to_mpi(dt), team.comm, &req);
-    do {
-        MPI_Test(&req, &completed, MPI_STATUS_IGNORE);
-        ucc_context_progress(team.ctx);
-    } while(!completed);
+    for (   i = 0; i < size; i++) {
+        init_buffer(PTR_OFFSET(check_buf, i * single_rank_count * ucc_dt_size(dt)),
+                    single_rank_count, dt, mem_type, i * (iter_persistent + 1),
+                    single_rank_count * rank);
+    }
 
     return compare_buffers(rbuf, check_buf, single_rank_count * size, dt,
                            mem_type);

diff --git a/test/mpi/test_bcast.cc b/test/mpi/test_bcast.cc
@@ -45,6 +45,7 @@ ucc_status_t TestBcast::set_input(int iter_persistent)
     size_t count   = msgsize / dt_size;
     int    rank;
 
+    this->iter_persistent = iter_persistent;
     MPI_Comm_rank(team.comm, &rank);
     if (rank == root) {
         init_buffer(sbuf, count, dt, mem_type, rank * (iter_persistent + 1));
@@ -56,18 +57,12 @@ ucc_status_t TestBcast::set_input(int iter_persistent)
 
 ucc_status_t TestBcast::check()
 {
-    size_t       count = args.src.info.count;
-    MPI_Datatype mpi_dt = ucc_dt_to_mpi(dt);
-    int          rank, completed;
-    MPI_Request  req;
+    size_t count = args.src.info.count;
+    int rank;
 
     MPI_Comm_rank(team.comm, &rank);
-    MPI_Ibcast(check_buf, count, mpi_dt, root, team.comm, &req);
-    do {
-        MPI_Test(&req, &completed, MPI_STATUS_IGNORE);
-        ucc_context_progress(team.ctx);
-    } while(!completed);
-
+    init_buffer(check_buf, count, dt, UCC_MEMORY_TYPE_HOST,
+                root * (iter_persistent + 1));
     return (rank == root)
                ? UCC_OK
                : compare_buffers(sbuf, check_buf, count, dt, mem_type);

diff --git a/test/mpi/test_mpi.h b/test/mpi/test_mpi.h
@@ -278,7 +278,7 @@ class TestCase {
     uint8_t     progress_buf[1];
     size_t test_max_size;
     ucc_datatype_t dt;
-
+    int iter_persistent;
 public:
     void mpi_progress(void);
     test_skip_cause_t test_skip;
@@ -523,7 +523,7 @@ class TestScatterv : public TestCase {
 };
 
 void init_buffer(void *buf, size_t count, ucc_datatype_t dt,
-                 ucc_memory_type_t mt, int value);
+                 ucc_memory_type_t mt, int value, int offset = 0);
 
 ucc_status_t compare_buffers(void *rst, void *expected, size_t count,
                              ucc_datatype_t dt, ucc_memory_type_t mt);