CORE: Fix gtest

nsarka · Aug 22, 2024 · a1d55ea · a1d55ea
1 parent a4b09de
commit a1d55ea
Showing 1 changed file with 17 additions and 14 deletions.
diff --git a/test/gtest/asym_mem/test_asymmetric_memory.cc b/test/gtest/asym_mem/test_asymmetric_memory.cc
@@ -117,17 +117,17 @@ class test_asymmetric_memory : public ucc::test,
 
             root = ctxs[i]->args->root;
 
-            if (i == root) {
-                dst_mem_type = ctxs[root]->args->dst.info.mem_type;
+            if (i == root || ctxs[i]->args->coll_type == UCC_COLL_TYPE_SCATTER) {
+                dst_mem_type = ctxs[i]->args->dst.info.mem_type;
 
-                rst = (uint8_t*) ucc_malloc(ctxs[root]->rbuf_size, "validation buf");
+                rst = (uint8_t*) ucc_malloc(ctxs[i]->rbuf_size, "validation buf");
                 EXPECT_NE(rst, nullptr);
 
-                UCC_CHECK(ucc_mc_memcpy(rst, ctxs[root]->args->dst.info.buffer,
-                                        ctxs[root]->rbuf_size,
+                UCC_CHECK(ucc_mc_memcpy(rst, ctxs[i]->args->dst.info.buffer,
+                                        ctxs[i]->rbuf_size,
                                         UCC_MEMORY_TYPE_HOST, dst_mem_type));
 
-                for (int j = 0; j < ctxs[root]->rbuf_size; j++) {
+                for (int j = 0; j < ctxs[i]->rbuf_size; j++) {
                     if (result != rst[j]) {
                         ret = false;
                         break;
@@ -146,22 +146,25 @@ class test_asymmetric_memory : public ucc::test,
         size_t msglen = 2048;
         size_t src_modifier = 1;
         ucc_coll_type_t coll_type = ctxs[0]->args->coll_type;
+        int root = ctxs[0]->args->root;
         ucc_memory_type_t src_mem_type = ctxs[0]->args->src.info.mem_type;
 
         if (coll_type == UCC_COLL_TYPE_SCATTER) {
             src_modifier = tsize;
         }
 
         for (int i = 0; i < tsize; i++) {
-            ucc_coll_args_t *coll = ctxs[i]->args;
-            uint8_t *sbuf = (uint8_t*)ctxs[i]->init_buf;
-            for (int j = 0; j < msglen * src_modifier; j++) {
-                sbuf[j] = (uint8_t) data;
+            if (i == root || coll_type != UCC_COLL_TYPE_SCATTER) {
+                ucc_coll_args_t *coll = ctxs[i]->args;
+                uint8_t *sbuf = (uint8_t*)ctxs[i]->init_buf;
+                for (int j = 0; j < msglen * src_modifier; j++) {
+                    sbuf[j] = (uint8_t) data;
+                }
+                UCC_CHECK(ucc_mc_memcpy(coll->src.info.buffer,
+                                        ctxs[i]->init_buf,
+                                        msglen * src_modifier, src_mem_type,
+                                        UCC_MEMORY_TYPE_HOST));
             }
-            UCC_CHECK(ucc_mc_memcpy(coll->src.info.buffer,
-                                    ctxs[i]->init_buf,
-                                    msglen * src_modifier, src_mem_type,
-                                    UCC_MEMORY_TYPE_HOST));
         }
     }
 };