TEST/MPI: fix global memh and mem-map test lifecycle

Author: wfaderhold21

test/mpi/test_case.cc

@@ -230,30 +230,40 @@ ucc_status_t TestCase::register_memh_global(void *buf, size_t size,
 
 ucc_status_t TestCase::register_memhs(void *sbuf, size_t ssize, void *dbuf, size_t dsize)
 {
+    int all_have_dbuf, has_dbuf;
+
     if (memh_mode == UCC_TEST_MEMH_NONE) {
         return UCC_OK;
     }
 
     if (memh_mode == UCC_TEST_MEMH_GLOBAL) {
-        if (sbuf && ssize > 0) {
-            UCC_CHECK(register_memh_global(sbuf, ssize, &src_memh,
-                                           &src_global_memh));
-            args.src_memh.global_memh = src_global_memh;
-            args.mask  |= UCC_COLL_ARGS_FIELD_MEM_MAP_SRC_MEMH;
-            args.flags |= UCC_COLL_ARGS_FLAG_SRC_MEMH_GLOBAL;
-        }
-        if (dbuf && dsize > 0) {
-            UCC_CHECK(register_memh_global(dbuf, dsize, &dst_memh,
-                                           &dst_global_memh));
-            args.dst_memh.global_memh = dst_global_memh;
-            args.mask  |= UCC_COLL_ARGS_FIELD_MEM_MAP_DST_MEMH;
-            args.flags |= UCC_COLL_ARGS_FLAG_DST_MEMH_GLOBAL;
+        /* Global destination handles are consumed by alltoall/alltoallv
+         * algorithms, including TL/CUDA push and TL/UCP one-sided paths.
+         * Other collective types have consumers that treat a present memh
+         * field as a local handle and must not be given an array.
+         */
+        if (args.coll_type != UCC_COLL_TYPE_ALLTOALL &&
+            args.coll_type != UCC_COLL_TYPE_ALLTOALLV) {
+            return UCC_OK;
         }
-        /* Flags are only honored when FIELD_FLAGS is present in the mask. */
-        if (args.flags & (UCC_COLL_ARGS_FLAG_SRC_MEMH_GLOBAL |
-                          UCC_COLL_ARGS_FLAG_DST_MEMH_GLOBAL)) {
-            args.mask |= UCC_COLL_ARGS_FIELD_FLAGS;
+
+        /* register_memh_global contains MPI collectives. Enter it only when
+         * every rank has a destination buffer so all ranks execute the same
+         * exchange sequence (including zero-count alltoallv cases).
+         */
+        has_dbuf = dbuf && dsize > 0;
+        MPI_Allreduce(&has_dbuf, &all_have_dbuf, 1, MPI_INT, MPI_MIN,
+                      team.comm);
+        if (!all_have_dbuf) {
+            return UCC_OK;
         }
+
+        UCC_CHECK(register_memh_global(dbuf, dsize, &dst_memh,
+                                       &dst_global_memh));
+        args.dst_memh.global_memh = dst_global_memh;
+        args.mask  |= UCC_COLL_ARGS_FIELD_FLAGS |
+                      UCC_COLL_ARGS_FIELD_MEM_MAP_DST_MEMH;
+        args.flags |= UCC_COLL_ARGS_FLAG_DST_MEMH_GLOBAL;
         return UCC_OK;
     }
 
@@ -305,6 +315,7 @@ TestCase::TestCase(ucc_test_team_t &_team, ucc_coll_type_t ct,
     src_global_memh  = NULL;
     dst_global_memh  = NULL;
     memh_global_size = 0;
+    req              = NULL;
     test_skip      = TEST_SKIP_NONE;
     args.flags     = 0;
     args.mask      = 0;
@@ -331,7 +342,7 @@ TestCase::~TestCase()
     MPI_Cancel(&progress_request);
     MPI_Wait(&progress_request, &status);
 
-    if (TEST_SKIP_NONE == test_skip) {
+    if (TEST_SKIP_NONE == test_skip && req) {
         UCC_CHECK(ucc_collective_finalize(req));
     }
 

test/mpi/test_mem_map.cc

@@ -6,6 +6,7 @@
 
 #include "test_mpi.h"
 #include "mpi_util.h"
+#include <climits>
 #include <cstring>
 #include <vector>
 #include "components/mc/ucc_mc.h"
@@ -85,24 +86,73 @@ class TestMemMap : public TestCase
     void run(bool triggered) override
     {
         (void)triggered; /* mem_map tests don't support triggered mode */
-        ucc_mem_map_params_t map_params;
-        ucc_mem_map_t        segment;
-        int                  rank;
+        ucc_mem_map_mem_h    export_memh = nullptr;
+        ucc_mem_map_params_t map_params = {};
+        ucc_mem_map_t        segment    = {};
+        size_t               export_size = 0;
+        uint64_t             recv_size   = 0;
+        int                  comm_size, peer, rank;
+        ucc_status_t         status;
 
         MPI_Comm_rank(team.comm, &rank);
+        MPI_Comm_size(team.comm, &comm_size);
 
         /* Set up memory map parameters */
         segment.address       = test_buffer;
         segment.len           = buffer_size;
         map_params.segments   = &segment;
         map_params.n_segments = 1;
 
+        if (!is_export_test) {
+            /* Import operates on an exchanged serialized export handle. Send
+             * each rank's blob to its next peer so import never receives a
+             * NULL or otherwise uninitialized handle.
+             */
+            status = ucc_mem_map(team.ctx, UCC_MEM_MAP_MODE_EXPORT,
+                                 &map_params, &export_size, &export_memh);
+            if (status != UCC_OK) {
+                if (status == UCC_ERR_NOT_SUPPORTED ||
+                    status == UCC_ERR_NOT_IMPLEMENTED) {
+                    test_skip = TEST_SKIP_NOT_SUPPORTED;
+                    return;
+                }
+                UCC_CHECK(status);
+            }
+            if (!export_memh || export_size == 0) {
+                UCC_CHECK(UCC_ERR_INVALID_PARAM);
+            }
+
+            peer = (rank + 1) % comm_size;
+            uint64_t send_size = export_size;
+            MPI_Sendrecv(&send_size, 1, MPI_UINT64_T, peer, 1001,
+                         &recv_size, 1, MPI_UINT64_T,
+                         (rank + comm_size - 1) % comm_size, 1001,
+                         team.comm, MPI_STATUS_IGNORE);
+
+            if (send_size > INT_MAX || recv_size > INT_MAX) {
+                UCC_CHECK(UCC_ERR_INVALID_PARAM);
+            }
+            memh = ucc_malloc(recv_size, "import memh blob");
+            UCC_MALLOC_CHECK(memh);
+            MPI_Sendrecv(export_memh, static_cast<int>(send_size), MPI_BYTE,
+                         peer, 1002, memh, static_cast<int>(recv_size), MPI_BYTE,
+                         (rank + comm_size - 1) % comm_size, 1002,
+                         team.comm, MPI_STATUS_IGNORE);
+            memh_size = recv_size;
+        }
+
         /* Test memory map */
-        ucc_status_t status = ucc_mem_map(team.ctx, mode, &map_params,
-                                          &memh_size, &memh);
+        status = ucc_mem_map(team.ctx, mode, &map_params, &memh_size, &memh);
         if (status != UCC_OK) {
+            if (export_memh) {
+                UCC_CHECK(ucc_mem_unmap(&export_memh));
+            }
             if (status == UCC_ERR_NOT_SUPPORTED ||
                 status == UCC_ERR_NOT_IMPLEMENTED) {
+                if (memh) {
+                    ucc_free(memh);
+                    memh = nullptr;
+                }
                 test_skip = TEST_SKIP_NOT_SUPPORTED;
                 return;
             }
@@ -112,19 +162,22 @@ class TestMemMap : public TestCase
         if (!memh) {
             std::cerr << "Rank " << rank << ": Memory handle is NULL"
                       << std::endl;
-            return;
+            UCC_CHECK(UCC_ERR_INVALID_PARAM);
         }
-        if (memh_size == 0) {
+        if (is_export_test && memh_size == 0) {
             std::cerr << "Rank " << rank << ": Memory handle size is 0"
                       << std::endl;
-            return;
+            UCC_CHECK(UCC_ERR_INVALID_PARAM);
         }
 
         /* Verify data integrity after mapping */
         UCC_CHECK(check());
 
         /* Test unmap */
         UCC_CHECK(ucc_mem_unmap(&memh));
+        if (export_memh) {
+            UCC_CHECK(ucc_mem_unmap(&export_memh));
+        }
         if (memh != nullptr) {
             std::cerr << "Rank " << rank
                       << ": Memory handle not NULL after unmap" << std::endl;
@@ -205,9 +258,6 @@ class TestMemMapStress : public TestCase
                 ucc_mem_unmap(&memh);
             }
         }
-        if (test_buffer) {
-            ucc_free(test_buffer);
-        }
     }
 
     ucc_status_t set_input(int iter_persistent = 0) override
@@ -260,8 +310,8 @@ class TestMemMapStress : public TestCase
             ucc_mem_map_mem_h memh;
             size_t             memh_size;
 
-            /* Fill buffer with iteration-specific pattern */
-            memset(test_buffer, 0xCC + rank + i, buffer_size);
+            /* Keep the contents stable while repeatedly mapping the buffer. */
+            memset(test_buffer, 0xBB + rank, buffer_size);
 
             ucc_status_t status = ucc_mem_map(team.ctx, UCC_MEM_MAP_MODE_EXPORT,
                                               &map_params, &memh_size, &memh);
@@ -528,26 +578,21 @@ void run_mem_map_tests(UccTestMpi *test,
             if (!tc) continue;
 
             if (TEST_SKIP_NONE == tc->test_skip) {
+                UCC_CHECK(tc->set_input());
                 tc->run(false);
             } else {
                 results.push_back(
                     std::make_tuple(tc->args.coll_type, UCC_ERR_LAST));
                 continue;
             }
 
-            ucc_status_t status;
-            do {
-                tc->mpi_progress();
-                status = tc->test();
-            } while (status == UCC_INPROGRESS);
-
-            if (UCC_OK != status) {
-                std::cerr << "error during mem_map test: "
-                          << ucc_status_string(status) << std::endl;
-                MPI_Abort(MPI_COMM_WORLD, -1);
+            if (TEST_SKIP_NONE != tc->test_skip) {
+                results.push_back(
+                    std::make_tuple(tc->args.coll_type, UCC_ERR_LAST));
+                continue;
             }
 
-            status = tc->check();
+            ucc_status_t status = tc->check();
             results.push_back(
                 std::make_tuple(tc->args.coll_type, status));
         }