TEST: add 2step reduce to gtest

openucx · Mar 4, 2024 · 6fecb92 · 6fecb92
1 parent 3af837c
commit 6fecb92
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 41 deletions.
diff --git a/src/components/cl/hier/reduce/reduce_2step.c b/src/components/cl/hier/reduce/reduce_2step.c
@@ -117,7 +117,9 @@ ucc_cl_hier_reduce_2step_init_schedule(ucc_base_coll_args_t *coll_args,
             ? find_root_node_rank(root, cl_team) : 0;
 
         if ((root != rank) && SBGP_ENABLED(cl_team, NODE_LEADERS)) {
-            status = ucc_mc_alloc(&cl_schedule->scratch, args.max_frag_count,
+            status = ucc_mc_alloc(&cl_schedule->scratch,
+                                  args.max_frag_count *
+                                  ucc_dt_size(args.args.src.info.datatype),
                                   args.args.src.info.mem_type);
             if (ucc_unlikely(UCC_OK != status)) {
                 goto out;
@@ -273,7 +275,8 @@ UCC_CL_HIER_PROFILE_FUNC(ucc_status_t, ucc_cl_hier_reduce_2step_init,
     int                       n_frags, pipeline_depth;
     ucc_status_t              status;
 
-    if (coll_args->args.op == UCC_OP_AVG) {
+    if (UCC_IS_PERSISTENT(coll_args->args) ||
+        (coll_args->args.op == UCC_OP_AVG)) {
         return UCC_ERR_NOT_SUPPORTED;
     }
 

diff --git a/test/gtest/coll/test_reduce.cc b/test/gtest/coll/test_reduce.cc
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -283,55 +283,66 @@ template <typename T> class test_reduce_avg_order : public test_reduce<T> {
 template <typename T> class test_reduce_dbt : public test_reduce<T> {
 };
 
-#define TEST_DECLARE_WITH_ENV(_env, _n_procs)                                    \
-    {                                                                            \
-        UccJob        job(_n_procs, UccJob::UCC_JOB_CTX_GLOBAL, _env);           \
-        UccTeam_h     team   = job.create_team(_n_procs);                        \
-        int           repeat = 3;                                                \
-        UccCollCtxVec ctxs;                                                      \
-        std::vector<ucc_memory_type_t> mt = {UCC_MEMORY_TYPE_HOST};              \
-        if (UCC_OK == ucc_mc_available(UCC_MEMORY_TYPE_CUDA)) {                  \
-            mt.push_back(UCC_MEMORY_TYPE_CUDA);                                  \
-        }                                                                        \
-        if (UCC_OK == ucc_mc_available(UCC_MEMORY_TYPE_CUDA_MANAGED)) {          \
-            mt.push_back(UCC_MEMORY_TYPE_CUDA_MANAGED);                          \
-        }                                                                        \
-        for (auto count : {5, 256, 65536}) {                                     \
-            for (auto inplace : {TEST_NO_INPLACE, TEST_INPLACE}) {               \
-                for (auto m : mt) {                                              \
-                    CHECK_TYPE_OP_SKIP(TypeParam::dt, TypeParam::redop, m);      \
-                    SET_MEM_TYPE(m);                                             \
-                    this->set_inplace(inplace);                                  \
-                    this->data_init(_n_procs, TypeParam::dt, count, ctxs, true); \
-                    UccReq req(team, ctxs);                                      \
-                    CHECK_REQ_NOT_SUPPORTED_SKIP(req, this->data_fini(ctxs));    \
-                    for (auto i = 0; i < repeat; i++) {                          \
-                        req.start();                                             \
-                        req.wait();                                              \
-                        EXPECT_EQ(true, this->data_validate(ctxs));              \
-                        this->reset(ctxs);                                       \
-                    }                                                            \
-                    this->data_fini(ctxs);                                       \
-                }                                                                \
-            }                                                                    \
-        }                                                                        \
+template <typename T> class test_reduce_2step : public test_reduce<T> {
+};
+
+#define TEST_DECLARE_WITH_ENV(_env, _n_procs, _persistent)                     \
+    {                                                                          \
+        UccJob        job(_n_procs, UccJob::UCC_JOB_CTX_GLOBAL, _env);         \
+        UccTeam_h     team   = job.create_team(_n_procs);                      \
+        int           repeat = _persistent ? 3 : 1;                            \
+        UccCollCtxVec ctxs;                                                    \
+        std::vector<ucc_memory_type_t> mt = {UCC_MEMORY_TYPE_HOST};            \
+        if (UCC_OK == ucc_mc_available(UCC_MEMORY_TYPE_CUDA)) {                \
+            mt.push_back(UCC_MEMORY_TYPE_CUDA);                                \
+        }                                                                      \
+        if (UCC_OK == ucc_mc_available(UCC_MEMORY_TYPE_CUDA_MANAGED)) {        \
+            mt.push_back(UCC_MEMORY_TYPE_CUDA_MANAGED);                        \
+        }                                                                      \
+        for (auto count : {5, 256, 65536}) {                                   \
+            for (auto inplace : {TEST_NO_INPLACE, TEST_INPLACE}) {             \
+                for (auto m : mt) {                                            \
+                    CHECK_TYPE_OP_SKIP(TypeParam::dt, TypeParam::redop, m);    \
+                    SET_MEM_TYPE(m);                                           \
+                    this->set_inplace(inplace);                                \
+                    this->data_init(_n_procs, TypeParam::dt, count, ctxs,      \
+                                    _persistent);                              \
+                    UccReq req(team, ctxs);                                    \
+                    CHECK_REQ_NOT_SUPPORTED_SKIP(req, this->data_fini(ctxs));  \
+                    for (auto i = 0; i < repeat; i++) {                        \
+                        req.start();                                           \
+                        req.wait();                                            \
+                        EXPECT_EQ(true, this->data_validate(ctxs));            \
+                        this->reset(ctxs);                                     \
+                    }                                                          \
+                    this->data_fini(ctxs);                                     \
+                }                                                              \
+            }                                                                  \
+        }                                                                      \
     }
 
 TYPED_TEST_CASE(test_reduce_avg_order, CollReduceTypeOpsAvg);
 TYPED_TEST_CASE(test_reduce_dbt, CollReduceTypeOpsHost);
+TYPED_TEST_CASE(test_reduce_2step, CollReduceTypeOpsHost);
 
-ucc_job_env_t post_op_env    = {{"UCC_TL_UCP_REDUCE_AVG_PRE_OP", "0"}};
-ucc_job_env_t reduce_dbt_env = {{"UCC_TL_UCP_TUNE", "reduce:@dbt:0-inf:inf"},
-                                {"UCC_CLS", "basic"}};
+ucc_job_env_t post_op_env      = {{"UCC_TL_UCP_REDUCE_AVG_PRE_OP", "0"}};
+ucc_job_env_t reduce_dbt_env   = {{"UCC_TL_UCP_TUNE", "reduce:@dbt:0-inf:inf"},
+                                  {"UCC_CLS", "basic"}};
+ucc_job_env_t reduce_2step_env = {{"UCC_CL_HIER_TUNE", "reduce:@2step:0-inf:inf"},
+                                  {"UCC_CLS", "all"}};
 
 TYPED_TEST(test_reduce_avg_order, avg_post_op) {
-    TEST_DECLARE_WITH_ENV(post_op_env, 15);
+    TEST_DECLARE_WITH_ENV(post_op_env, 15, true);
 }
 
 TYPED_TEST(test_reduce_dbt, reduce_dbt_shift) {
-    TEST_DECLARE_WITH_ENV(reduce_dbt_env, 15);
+    TEST_DECLARE_WITH_ENV(reduce_dbt_env, 15, true);
 }
 
 TYPED_TEST(test_reduce_dbt, reduce_dbt_mirror) {
-    TEST_DECLARE_WITH_ENV(reduce_dbt_env, 16);
+    TEST_DECLARE_WITH_ENV(reduce_dbt_env, 16, true);
+}
+
+TYPED_TEST(test_reduce_2step, 2step) {
+    TEST_DECLARE_WITH_ENV(reduce_2step_env, 16, false);
 }