ompi/coll/cuda: implement reduce_local

Signed-off-by: Akshay Venkatesh <akvenkatesh@nvidia.com>

bot:notacherrypick

Author: Akshay-Venkatesh

ompi/mca/coll/cuda/coll_cuda.h

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2024      NVIDIA Corporation.  All rights reserved.
  * Copyright (c) 2014      The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
@@ -45,6 +46,11 @@ mca_coll_cuda_allreduce(const void *sbuf, void *rbuf, int count,
                         struct ompi_communicator_t *comm,
                         mca_coll_base_module_t *module);
 
+int mca_coll_cuda_reduce_local(const void *sbuf, void *rbuf, int count,
+                               struct ompi_datatype_t *dtype,
+                               struct ompi_op_t *op,
+                               mca_coll_base_module_t *module);
+
 int mca_coll_cuda_reduce(const void *sbuf, void *rbuf, int count,
                          struct ompi_datatype_t *dtype,
                          struct ompi_op_t *op,

ompi/mca/coll/cuda/coll_cuda_module.c

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2024      NVIDIA Corporation.  All rights reserved.
  * Copyright (c) 2014-2017 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
@@ -104,6 +105,7 @@ mca_coll_cuda_comm_query(struct ompi_communicator_t *comm,
     cuda_module->super.coll_gather     = NULL;
     cuda_module->super.coll_gatherv    = NULL;
     cuda_module->super.coll_reduce     = mca_coll_cuda_reduce;
+    cuda_module->super.coll_reduce_local   = mca_coll_cuda_reduce_local;
     cuda_module->super.coll_reduce_scatter = NULL;
     cuda_module->super.coll_reduce_scatter_block = mca_coll_cuda_reduce_scatter_block;
     cuda_module->super.coll_scan       = mca_coll_cuda_scan;

ompi/mca/coll/cuda/coll_cuda_reduce.c

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2024      NVIDIA Corporation.  All rights reserved.
  * Copyright (c) 2004-2015 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
@@ -34,7 +35,7 @@ mca_coll_cuda_reduce(const void *sbuf, void *rbuf, int count,
                      mca_coll_base_module_t *module)
 {
     mca_coll_cuda_module_t *s = (mca_coll_cuda_module_t*) module;
-    int rank = ompi_comm_rank(comm);
+    int rank = (comm == NULL) ? -1 : ompi_comm_rank(comm);
     ptrdiff_t gap;
     char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
     const char *sbuf2;
@@ -64,9 +65,15 @@ mca_coll_cuda_reduce(const void *sbuf, void *rbuf, int count,
         rbuf2 = rbuf; /* save away original buffer */
         rbuf = rbuf1 - gap;
     }
-    rc = s->c_coll.coll_reduce((void *) sbuf, rbuf, count,
-                               dtype, op, root, comm,
-                               s->c_coll.coll_reduce_module);
+
+    if ((comm != NULL) && (root == -1)) {
+        ompi_op_reduce(op, (void *)sbuf, rbuf, count, dtype);
+        rc = OMPI_SUCCESS;
+    } else {
+        rc = s->c_coll.coll_reduce((void *) sbuf, rbuf, count,
+                                   dtype, op, root, comm,
+                                   s->c_coll.coll_reduce_module);
+    }
 
     if (NULL != sbuf1) {
         free(sbuf1);
@@ -78,3 +85,12 @@ mca_coll_cuda_reduce(const void *sbuf, void *rbuf, int count,
     }
     return rc;
 }
+
+int
+mca_coll_cuda_reduce_local(const void *sbuf, void *rbuf, int count,
+                           struct ompi_datatype_t *dtype,
+                           struct ompi_op_t *op,
+                           mca_coll_base_module_t *module)
+{
+    return mca_coll_cuda_reduce(sbuf, rbuf, count, dtype, op, -1, NULL ,module);
+}