Update ompi_op_reduce to take size_t count

 * If the `count` is greater than `INT_MAX` then we call the
   operation in chunks that fit into an `int`.
 * This moves the functionality out of libnbc and into the common
   reduction operation so that all collectives may pass larger
   counts than `INT_MAX` into the internal reduction operation.

Signed-off-by: Joshua Hursey <jhursey@us.ibm.com>
(cherry picked from commit 6075048)

Author: jjhursey

ompi/mca/coll/libnbc/nbc.c

@@ -536,30 +536,7 @@ static inline int NBC_Start_round(NBC_Handle *handle) {
           buf2=opargs.buf2;
         }
 
-        /* If the count is > INT_MAX then we need to call ompi_op_reduce()
-         * in iterations of counts <= INT_MAX since it has an `int count`
-         * parameter.
-         */
-        if( OPAL_UNLIKELY(opargs.count > INT_MAX) ) {
-          size_t done_count = 0, shift;
-          int iter_count;
-          ptrdiff_t ext, lb;
-
-          ompi_datatype_get_extent (opargs.datatype, &lb, &ext);
-
-          while(done_count < opargs.count) {
-            if( done_count + INT_MAX > opargs.count ) {
-              iter_count = opargs.count - done_count;
-            } else {
-              iter_count = INT_MAX;
-            }
-            shift = done_count * ext;
-            ompi_op_reduce(opargs.op, buf1 + shift, buf2 + shift, iter_count, opargs.datatype);
-            done_count += iter_count;
-          }
-        } else {
-          ompi_op_reduce(opargs.op, buf1, buf2, opargs.count, opargs.datatype);
-        }
+        ompi_op_reduce(opargs.op, buf1, buf2, opargs.count, opargs.datatype);
         break;
       case COPY:
         NBC_DEBUG(5, "  COPY   (offset %li) ", offset);

ompi/op/op.h

@@ -20,6 +20,7 @@
  *                         and Technology (RIST).  All rights reserved.
  * Copyright (c) 2018      Triad National Security, LLC. All rights
  *                         reserved.
+ * Copyright (c) 2021      IBM Corporation.  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -510,10 +511,41 @@ static inline bool ompi_op_is_valid(ompi_op_t * op, ompi_datatype_t * ddt,
  * is not defined to have that operation, it is likely to seg fault.
  */
 static inline void ompi_op_reduce(ompi_op_t * op, void *source,
-                                  void *target, int count,
+                                  void *target, size_t full_count,
                                   ompi_datatype_t * dtype)
 {
     MPI_Fint f_dtype, f_count;
+    int count = full_count;
+
+    /*
+     * If the full_count is > INT_MAX then we need to call the reduction op
+     * in iterations of counts <= INT_MAX since it has an `int *len`
+     * parameter.
+     *
+     * Note: When we add BigCount support then we can distinguish between
+     * a reduction operation with `int *len` and `MPI_Count *len`. At which
+     * point we can avoid this loop.
+     */
+    if( OPAL_UNLIKELY(full_count > INT_MAX) ) {
+        size_t done_count = 0, shift;
+        int iter_count;
+        ptrdiff_t ext, lb;
+
+        ompi_datatype_get_extent(dtype, &lb, &ext);
+
+        while(done_count < full_count) {
+            if(done_count + INT_MAX > full_count) {
+                iter_count = full_count - done_count;
+            } else {
+                iter_count = INT_MAX;
+            }
+            shift = done_count * ext;
+            // Recurse one level in iterations of 'int'
+            ompi_op_reduce(op, (char*)source + shift, (char*)target + shift, iter_count, dtype);
+            done_count += iter_count;
+        }
+        return;
+    }
 
     /*
      * Call the reduction function.  Two dimensions: a) if both the op