Merge pull request #13030 from hppritcha/add_support_for_user_defined_bigops

ops: add support for user-defined big count ops

Author: hppritcha

ompi/mpi/c/op_create.c

@@ -12,6 +12,8 @@
  * Copyright (c) 2008-2009 Cisco Systems, Inc.  All rights reserved.
  * Copyright (c) 2015      Research Organization for Information Science
  *                         and Technology (RIST). All rights reserved.
+ * Copyright (c) 2025      Triad National Security, LLC. All rights
+ *                         reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -57,6 +59,7 @@ int MPI_Op_create(MPI_User_function * function, int commute, MPI_Op * op)
     /* Create and cache the op.  Sets a refcount of 1. */
 
     *op = ompi_op_create_user(OPAL_INT_TO_BOOL(commute),
+                              false,
                               (ompi_op_fortran_handler_fn_t *) function);
     if (NULL == *op) {
         err = MPI_ERR_INTERN;

ompi/op/op.c

@@ -17,7 +17,7 @@
  * Copyright (c) 2015      Research Organization for Information Science
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2018      FUJITSU LIMITED.  All rights reserved.
- * Copyright (c) 2018      Triad National Security, LLC. All rights
+ * Copyright (c) 2018-2025 Triad National Security, LLC. All rights
  *                         reserved.
  * $COPYRIGHT$
  *
@@ -353,6 +353,7 @@ static int ompi_op_finalize (void)
  * Create a new MPI_Op
  */
 ompi_op_t *ompi_op_create_user(bool commute,
+                               bool bigcount,
                                ompi_op_fortran_handler_fn_t func)
 {
     ompi_op_t *new_op;
@@ -382,6 +383,9 @@ ompi_op_t *ompi_op_create_user(bool commute,
     if (commute) {
         new_op->o_flags |= OMPI_OP_FLAGS_COMMUTE;
     }
+    if(bigcount) {
+        new_op->o_flags |= OMPI_OP_FLAGS_BIGCOUNT;
+    }
 
     opal_string_copy(new_op->o_name, "USER OP", sizeof(new_op->o_name));
     new_op->o_name[sizeof(new_op->o_name) - 1] = '\0';

ompi/op/op.h

@@ -18,7 +18,7 @@
  * Copyright (c) 2018      Amazon.com, Inc. or its affiliates.  All Rights reserved.
  * Copyright (c) 2019      Research Organization for Information Science
  *                         and Technology (RIST).  All rights reserved.
- * Copyright (c) 2018      Triad National Security, LLC. All rights
+ * Copyright (c) 2018-2025 Triad National Security, LLC. All rights
  *                         reserved.
  * Copyright (c) 2021      IBM Corporation.  All rights reserved.
  * $COPYRIGHT$
@@ -61,12 +61,16 @@ BEGIN_C_DECLS
  */
 typedef void (ompi_op_c_handler_fn_t)(const void *, void *, int *,
                                       struct ompi_datatype_t **);
+typedef void (ompi_op_c_handler_bc_fn_t)(const void *, void *, size_t *,
+                                         struct ompi_datatype_t **);
 
 /**
  * Typedef for fortran user-defined MPI_Ops.
  */
 typedef void (ompi_op_fortran_handler_fn_t)(const void *, void *,
                                             MPI_Fint *, MPI_Fint *);
+typedef void (ompi_op_fortran_handler_bc_fn_t)(const void *, void *,
+                                               size_t *, MPI_Fint *);
 
 /**
  * Typedef for Java op functions intercept (used for user-defined
@@ -98,8 +102,8 @@ typedef void (ompi_op_java_handler_fn_t)(const void *, void *, int *,
 #define OMPI_OP_FLAGS_FLOAT_ASSOC  0x0020
 /** Set if the callback function is communative */
 #define OMPI_OP_FLAGS_COMMUTE      0x0040
-
-
+/** Set if the callback function is using bigcount */
+#define OMPI_OP_FLAGS_BIGCOUNT     0x0080
 
 
 /*
@@ -152,8 +156,12 @@ struct ompi_op_t {
         ompi_op_base_op_fns_t intrinsic;
         /** C handler function pointer */
         ompi_op_c_handler_fn_t *c_fn;
+        /** C handler function pointer  - bigcount*/
+        ompi_op_c_handler_bc_fn_t *c_fn_bc;
         /** Fortran handler function pointer */
         ompi_op_fortran_handler_fn_t *fort_fn;
+        /** Fortran handler function pointer  - bigcount*/
+        ompi_op_fortran_handler_bc_fn_t *fort_fn_bc;
         /** Java intercept function data */
         struct {
             /* The OMPI C++ callback/intercept function */
@@ -333,6 +341,8 @@ int ompi_op_init(void);
  *
  * @param commute Boolean indicating whether the operation is
  *        communative or not
+ * @param bigcount Boolean indicating whether or not the op is
+ *        using the bigcount (MPI_Count) interface
  * @param func Function pointer of the error handler
  *
  * @returns op Pointer to the ompi_op_t that will be
@@ -355,6 +365,7 @@ int ompi_op_init(void);
  * manually.
  */
 ompi_op_t *ompi_op_create_user(bool commute,
+                               bool bigcount,
                                ompi_op_fortran_handler_fn_t func);
 
 /**
@@ -512,11 +523,9 @@ static inline void ompi_op_reduce(ompi_op_t * op, const void *source,
      * 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) ) {
+    if(OPAL_UNLIKELY((full_count > INT_MAX) &&
+        (0 == (op->o_flags & OMPI_OP_FLAGS_BIGCOUNT)))) {
         size_t done_count = 0, shift;
         int iter_count;
         ptrdiff_t ext, lb;
@@ -578,8 +587,12 @@ static inline void ompi_op_reduce(ompi_op_t * op, const void *source,
     /* 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);
-        op->o_func.fort_fn(source, target, &f_count, &f_dtype);
+        if (0 == (op->o_flags & OMPI_OP_FLAGS_BIGCOUNT)) {
+            f_count = OMPI_INT_2_FINT(count);
+            op->o_func.fort_fn(source, target, &f_count, &f_dtype);
+        } else {
+            op->o_func.fort_fn_bc(source, target, &full_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,
@@ -588,15 +601,25 @@ static inline void ompi_op_reduce(ompi_op_t * op, const void *source,
                                           op->o_func.java_data.object);
         return;
     }
-    op->o_func.c_fn(source, target, &count, &dtype);
+    if (0 == (op->o_flags & OMPI_OP_FLAGS_BIGCOUNT)) {
+        op->o_func.c_fn(source, target, &count, &dtype);
+    } else {
+        op->o_func.c_fn_bc(source, target, &full_count, &dtype);
+    }
     return;
 }
 
 static inline void ompi_3buff_op_user (ompi_op_t *op, void * restrict source1, void * restrict source2,
-                                       void * restrict result, int count, struct ompi_datatype_t *dtype)
+                                       void * restrict result, size_t full_count, struct ompi_datatype_t *dtype)
 {
-    ompi_datatype_copy_content_same_ddt (dtype, count, (char*)result, (char*)source1);
-    op->o_func.c_fn (source2, result, &count, &dtype);
+    ompi_datatype_copy_content_same_ddt (dtype, full_count, (char*)result, (char*)source1);
+    if (0 == (op->o_flags & OMPI_OP_FLAGS_BIGCOUNT)) {
+        assert(full_count <= INT_MAX);
+        int count = (int)full_count;  /* protected by loop in only caller of this function */
+        op->o_func.c_fn (source2, result, &count, &dtype);
+    } else {
+        op->o_func.c_fn_bc (source2, result, &full_count, &dtype);
+    }
 }
 
 /**
@@ -618,13 +641,11 @@ static inline void ompi_3buff_op_user (ompi_op_t *op, void * restrict source1, v
  * with the values in the source buffer and the result is stored in
  * the target buffer).
  *
- * This function will *only* be invoked on intrinsic MPI_Ops.
- *
  * Otherwise, this function is the same as ompi_op_reduce.
  */
 static inline void ompi_3buff_op_reduce(ompi_op_t * op, void *source1,
                                         void *source2, void *target,
-                                        int count, ompi_datatype_t * dtype)
+                                        size_t full_count, ompi_datatype_t * dtype)
 {
     void *restrict src1;
     void *restrict src2;
@@ -633,13 +654,36 @@ static inline void ompi_3buff_op_reduce(ompi_op_t * op, void *source1,
     src2 = source2;
     tgt = target;
 
+    if(OPAL_UNLIKELY((full_count > INT_MAX) &&
+        (0 == (op->o_flags & OMPI_OP_FLAGS_BIGCOUNT)))) {
+        size_t done_count = 0, shift, 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_3buff_op_reduce(op, (char*)source1 + shift, (char *)source2 + shift,
+                                (char*)target + shift, iter_count, dtype);
+            done_count += iter_count;
+        }
+        return;
+    }
+
     if (OPAL_LIKELY(ompi_op_is_intrinsic (op))) {
+        int count = (int)full_count;
         op->o_3buff_intrinsic.fns[ompi_op_ddt_map[dtype->id]](src1, src2,
                                                               tgt, &count,
                                                               &dtype,
                                                               op->o_3buff_intrinsic.modules[ompi_op_ddt_map[dtype->id]]);
     } else {
-        ompi_3buff_op_user (op, src1, src2, tgt, count, dtype);
+        ompi_3buff_op_user (op, src1, src2, tgt, full_count, dtype);
     }
 }