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 <[email protected]>

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 count,
                                   ompi_datatype_t * dtype)
 {
     MPI_Fint f_dtype, f_count;
+    int int_count = count;
+
+    /*
+     * If the 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(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 < count) {
+            if(done_count + INT_MAX > count) {
+                iter_count = 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
@@ -548,25 +580,25 @@ static inline void ompi_op_reduce(ompi_op_t * op, void *source,
             dtype_id = ompi_op_ddt_map[dtype->id];
         }
         op->o_func.intrinsic.fns[dtype_id](source, target,
-                                           &count, &dtype,
+                                           &int_count, &dtype,
                                            op->o_func.intrinsic.modules[dtype_id]);
         return;
     }
 
     /* User-defined function */
     if (0 != (op->o_flags & OMPI_OP_FLAGS_FORTRAN_FUNC)) {
         f_dtype = OMPI_INT_2_FINT(dtype->d_f_to_c_index);
-        f_count = OMPI_INT_2_FINT(count);
+        f_count = OMPI_INT_2_FINT(int_count);
         op->o_func.fort_fn(source, target, &f_count, &f_dtype);
         return;
     } else if (0 != (op->o_flags & OMPI_OP_FLAGS_JAVA_FUNC)) {
-        op->o_func.java_data.intercept_fn(source, target, &count, &dtype,
+        op->o_func.java_data.intercept_fn(source, target, &int_count, &dtype,
                                           op->o_func.java_data.baseType,
                                           op->o_func.java_data.jnienv,
                                           op->o_func.java_data.object);
         return;
     }
-    op->o_func.c_fn(source, target, &count, &dtype);
+    op->o_func.c_fn(source, target, &int_count, &dtype);
     return;
 }