Merge pull request #7902 from vspetrov/v4.1.x_hcoll_reduce_scatter

jsquyres · web-flow · commit 36bcc48dc120 · 2020-07-27T15:50:09.000-04:00
V4.1.x hcoll reduce scatter
diff --git a/ompi/mca/coll/hcoll/coll_hcoll.h b/ompi/mca/coll/hcoll/coll_hcoll.h
@@ -142,6 +142,8 @@ struct mca_coll_hcoll_module_t {
     mca_coll_base_module_t *previous_scatterv_module;
     mca_coll_base_module_reduce_scatter_fn_t previous_reduce_scatter;
     mca_coll_base_module_t *previous_reduce_scatter_module;
+    mca_coll_base_module_reduce_scatter_block_fn_t previous_reduce_scatter_block;
+    mca_coll_base_module_t *previous_reduce_scatter_block_module;
     mca_coll_base_module_ibcast_fn_t previous_ibcast;
     mca_coll_base_module_t *previous_ibcast_module;
     mca_coll_base_module_ibarrier_fn_t previous_ibarrier;
@@ -212,6 +214,18 @@ int mca_coll_hcoll_allreduce(const void *sbuf, void *rbuf, int count,
                             struct ompi_communicator_t *comm,
                             mca_coll_base_module_t *module);
 
+#if HCOLL_API > HCOLL_VERSION(4,5)
+int mca_coll_hcoll_reduce_scatter_block(const void *sbuf, void *rbuf, int rcount,
+                                        struct ompi_datatype_t *dtype,
+                                        struct ompi_op_t *op,
+                                        struct ompi_communicator_t *comm,
+                                        mca_coll_base_module_t *module);
+int mca_coll_hcoll_reduce_scatter(const void *sbuf, void *rbuf, const int* rcounts,
+                                  struct ompi_datatype_t *dtype,
+                                  struct ompi_op_t *op,
+                                  struct ompi_communicator_t *comm,
+                                  mca_coll_base_module_t *module);
+#endif
 int mca_coll_hcoll_reduce(const void *sbuf, void *rbuf, int count,
                             struct ompi_datatype_t *dtype,
                             struct ompi_op_t *op,
@@ -303,11 +317,11 @@ int mca_coll_hcoll_ialltoall(const void *sbuf, int scount,
                             mca_coll_base_module_t *module);
 
 #if HCOLL_API >= HCOLL_VERSION(3,7)
-int mca_coll_hcoll_ialltoallv(const void *sbuf, int *scounts,
-                            int *sdisps,
+int mca_coll_hcoll_ialltoallv(const void *sbuf, const int *scounts,
+                            const int *sdisps,
                             struct ompi_datatype_t *sdtype,
-                            void *rbuf, int *rcounts,
-                            int *rdisps,
+                            void *rbuf, const int *rcounts,
+                            const int *rdisps,
                             struct ompi_datatype_t *rdtype,
                             struct ompi_communicator_t *comm,
                             ompi_request_t **req,
diff --git a/ompi/mca/coll/hcoll/coll_hcoll_module.c b/ompi/mca/coll/hcoll/coll_hcoll_module.c
@@ -51,6 +51,7 @@ static void mca_coll_hcoll_module_clear(mca_coll_hcoll_module_t *hcoll_module)
     hcoll_module->previous_alltoallw  = NULL;
     hcoll_module->previous_reduce     = NULL;
     hcoll_module->previous_reduce_scatter  = NULL;
+    hcoll_module->previous_reduce_scatter_block  = NULL;
     hcoll_module->previous_ibarrier    = NULL;
     hcoll_module->previous_ibcast      = NULL;
     hcoll_module->previous_iallreduce  = NULL;
@@ -119,6 +120,8 @@ static void mca_coll_hcoll_module_destruct(mca_coll_hcoll_module_t *hcoll_module
         OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_barrier_module);
         OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_bcast_module);
         OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_allreduce_module);
+        OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_reduce_scatter_block_module);
+        OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_reduce_scatter_module);
         OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_allgather_module);
         OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_allgatherv_module);
         OBJ_RELEASE_IF_NOT_NULL(hcoll_module->previous_gatherv_module);
@@ -173,6 +176,8 @@ static int mca_coll_hcoll_save_coll_handlers(mca_coll_hcoll_module_t *hcoll_modu
     HCOL_SAVE_PREV_COLL_API(barrier);
     HCOL_SAVE_PREV_COLL_API(bcast);
     HCOL_SAVE_PREV_COLL_API(allreduce);
+    HCOL_SAVE_PREV_COLL_API(reduce_scatter_block);
+    HCOL_SAVE_PREV_COLL_API(reduce_scatter);
     HCOL_SAVE_PREV_COLL_API(reduce);
     HCOL_SAVE_PREV_COLL_API(allgather);
     HCOL_SAVE_PREV_COLL_API(allgatherv);
@@ -419,6 +424,12 @@ mca_coll_hcoll_comm_query(struct ompi_communicator_t *comm, int *priority)
     hcoll_module->super.coll_ialltoallv = hcoll_collectives.coll_ialltoallv ? mca_coll_hcoll_ialltoallv : NULL;
 #else
     hcoll_module->super.coll_ialltoallv = NULL;
+#endif
+#if HCOLL_API > HCOLL_VERSION(4,5)
+    hcoll_module->super.coll_reduce_scatter_block = hcoll_collectives.coll_reduce_scatter_block ?
+        mca_coll_hcoll_reduce_scatter_block : NULL;
+    hcoll_module->super.coll_reduce_scatter = hcoll_collectives.coll_reduce_scatter ?
+        mca_coll_hcoll_reduce_scatter : NULL;
 #endif
     *priority = cm->hcoll_priority;
     module = &hcoll_module->super;
diff --git a/ompi/mca/coll/hcoll/coll_hcoll_ops.c b/ompi/mca/coll/hcoll/coll_hcoll_ops.c
@@ -136,7 +136,8 @@ int mca_coll_hcoll_allgatherv(const void *sbuf, int scount,
                                              hcoll_module->previous_allgatherv_module);
         return rc;
     }
-    rc = hcoll_collectives.coll_allgatherv((void *)sbuf,scount,stype,rbuf,rcount,displs,rtype,hcoll_module->hcoll_context);
+    rc = hcoll_collectives.coll_allgatherv((void *)sbuf,scount,stype,rbuf,(int*)rcount,
+                                           (int*)displs,rtype,hcoll_module->hcoll_context);
     if (HCOLL_SUCCESS != rc){
         HCOL_VERBOSE(20,"RUNNING FALLBACK ALLGATHERV");
         rc = hcoll_module->previous_allgatherv(sbuf,scount,sdtype,
@@ -558,7 +559,7 @@ int mca_coll_hcoll_iallgatherv(const void *sbuf, int scount,
                                              hcoll_module->previous_iallgatherv_module);
         return rc;
     }
-    rc = hcoll_collectives.coll_iallgatherv((void *)sbuf,scount,stype,rbuf,rcount,displs,rtype,
+    rc = hcoll_collectives.coll_iallgatherv((void *)sbuf,scount,stype,rbuf,(int*)rcount,(int*)displs,rtype,
             hcoll_module->hcoll_context, rt_handle);
     if (HCOLL_SUCCESS != rc){
        HCOL_VERBOSE(20,"RUNNING FALLBACK NON-BLOCKING ALLGATHER");
@@ -724,9 +725,9 @@ int mca_coll_hcoll_igatherv(const void* sbuf, int scount,
 
 
 #if HCOLL_API >= HCOLL_VERSION(3,7)
-int mca_coll_hcoll_ialltoallv(const void *sbuf, int *scounts, int *sdisps,
+int mca_coll_hcoll_ialltoallv(const void *sbuf, const int *scounts, const int *sdisps,
                               struct ompi_datatype_t *sdtype,
-                              void *rbuf, int *rcounts, int *rdisps,
+                              void *rbuf, const int *rcounts, const int *rdisps,
                               struct ompi_datatype_t *rdtype,
                               struct ompi_communicator_t *comm,
                               ompi_request_t ** request,
@@ -760,3 +761,88 @@ int mca_coll_hcoll_ialltoallv(const void *sbuf, int *scounts, int *sdisps,
     return rc;
 }
 #endif
+
+#if HCOLL_API > HCOLL_VERSION(4,5)
+int mca_coll_hcoll_reduce_scatter_block(const void *sbuf, void *rbuf, int rcount,
+                                        struct ompi_datatype_t *dtype,
+                                        struct ompi_op_t *op,
+                                        struct ompi_communicator_t *comm,
+                                        mca_coll_base_module_t *module) {
+    dte_data_representation_t Dtype;
+    hcoll_dte_op_t *Op;
+    int rc;
+    HCOL_VERBOSE(20,"RUNNING HCOL REDUCE SCATTER BLOCK");
+    mca_coll_hcoll_module_t *hcoll_module = (mca_coll_hcoll_module_t*)module;
+    Dtype = ompi_dtype_2_hcoll_dtype(dtype, NO_DERIVED);
+    if (OPAL_UNLIKELY(HCOL_DTE_IS_ZERO(Dtype))){
+        /*If we are here then datatype is not simple predefined datatype */
+        /*In future we need to add more complex mapping to the dte_data_representation_t */
+        /* Now use fallback */
+        HCOL_VERBOSE(20,"Ompi_datatype is not supported: dtype = %s; calling fallback allreduce;",
+                     dtype->super.name);
+        goto fallback;
+    }
+
+    Op = ompi_op_2_hcolrte_op(op);
+    if (OPAL_UNLIKELY(HCOL_DTE_OP_NULL == Op->id)){
+        /*If we are here then datatype is not simple predefined datatype */
+        /*In future we need to add more complex mapping to the dte_data_representation_t */
+        /* Now use fallback */
+        HCOL_VERBOSE(20,"ompi_op_t is not supported: op = %s; calling fallback allreduce;",
+                     op->o_name);
+        goto fallback;
+    }
+
+    rc = hcoll_collectives.coll_reduce_scatter_block((void *)sbuf,rbuf,rcount,Dtype,Op,hcoll_module->hcoll_context);
+    if (HCOLL_SUCCESS != rc){
+    fallback:        
+        HCOL_VERBOSE(20,"RUNNING FALLBACK ALLREDUCE");
+        rc = hcoll_module->previous_reduce_scatter_block(sbuf,rbuf,
+                                             rcount,dtype,op,
+                                             comm, hcoll_module->previous_allreduce_module);
+    }
+    return rc;
+}
+
+int mca_coll_hcoll_reduce_scatter(const void *sbuf, void *rbuf, const int* rcounts,
+                                  struct ompi_datatype_t *dtype,
+                                  struct ompi_op_t *op,
+                                  struct ompi_communicator_t *comm,
+                                  mca_coll_base_module_t *module) {
+    dte_data_representation_t Dtype;
+    hcoll_dte_op_t *Op;
+    int rc;
+    HCOL_VERBOSE(20,"RUNNING HCOL REDUCE SCATTER");
+    mca_coll_hcoll_module_t *hcoll_module = (mca_coll_hcoll_module_t*)module;
+    Dtype = ompi_dtype_2_hcoll_dtype(dtype, NO_DERIVED);
+    if (OPAL_UNLIKELY(HCOL_DTE_IS_ZERO(Dtype))){
+        /*If we are here then datatype is not simple predefined datatype */
+        /*In future we need to add more complex mapping to the dte_data_representation_t */
+        /* Now use fallback */
+        HCOL_VERBOSE(20,"Ompi_datatype is not supported: dtype = %s; calling fallback allreduce;",
+                     dtype->super.name);
+        goto fallback;
+    }
+
+    Op = ompi_op_2_hcolrte_op(op);
+    if (OPAL_UNLIKELY(HCOL_DTE_OP_NULL == Op->id)){
+        /*If we are here then datatype is not simple predefined datatype */
+        /*In future we need to add more complex mapping to the dte_data_representation_t */
+        /* Now use fallback */
+        HCOL_VERBOSE(20,"ompi_op_t is not supported: op = %s; calling fallback allreduce;",
+                     op->o_name);
+        goto fallback;
+    }
+
+    rc = hcoll_collectives.coll_reduce_scatter((void*)sbuf, rbuf, (int*)rcounts,
+                                               Dtype, Op, hcoll_module->hcoll_context);
+    if (HCOLL_SUCCESS != rc){
+    fallback:        
+        HCOL_VERBOSE(20,"RUNNING FALLBACK ALLREDUCE");
+        rc = hcoll_module->previous_reduce_scatter(sbuf,rbuf,
+                                             rcounts,dtype,op,
+                                             comm, hcoll_module->previous_allreduce_module);
+    }
+    return rc;
+}
+#endif
