Revoke coll/han subcomms

Signed-off-by: Matthew Whitlock <mwhitlo@sandia.gov>

Author: Matthew-Whitlock

ompi/communicator/ft/comm_ft_revoke.c

@@ -91,6 +91,8 @@ static int ompi_comm_revoke_local(ompi_communicator_t* comm, ompi_comm_rbcast_me
     comm->any_source_enabled = false;
     /* purge the communicator unexpected fragments and matching logic */
     MCA_PML_CALL(revoke_comm(comm, false));
+    /* revoke any subcomms created by coll */
+    comm->c_coll->coll_revoke_local(comm);
     /* Signal the point-to-point stack to recheck requests */
     wait_sync_global_wakeup(MPI_ERR_REVOKED);
     return true;

ompi/mca/coll/base/Makefile.am

@@ -31,6 +31,7 @@ libmca_coll_la_SOURCES += \
         base/coll_base_comm_unselect.c \
         base/coll_base_find_available.c \
         base/coll_base_frame.c \
+        base/coll_base_revoke_local.c \
         base/coll_base_bcast.c \
         base/coll_base_scatter.c \
         base/coll_base_topo.c \

ompi/mca/coll/base/coll_base_comm_select.c

@@ -49,6 +49,7 @@
 #include "ompi/mca/coll/coll.h"
 #include "ompi/mca/coll/base/base.h"
 #include "ompi/mca/coll/base/coll_base_util.h"
+#include "ompi/mca/coll/base/coll_base_functions.h"
 
 /*
  * Stuff for the OBJ interface
@@ -227,6 +228,7 @@ int mca_coll_base_comm_select(ompi_communicator_t * comm)
     /* Initialize all the relevant pointers, since they're used as
      * sentinel values */
     comm->c_coll = (mca_coll_base_comm_coll_t*)calloc(1, sizeof(mca_coll_base_comm_coll_t));
+    comm->c_coll->coll_revoke_local = mca_coll_base_revoke_local;
 
     opal_output_verbose(10, ompi_coll_base_framework.framework_output,
                         "coll:base:comm_select: Checking all available modules");

ompi/mca/coll/base/coll_base_functions.h

@@ -306,6 +306,8 @@ int mca_coll_base_reduce_local(const void *inbuf, void *inoutbuf, size_t count,
                                struct ompi_datatype_t * dtype, struct ompi_op_t * op,
                                mca_coll_base_module_t *module);
 
+int mca_coll_base_revoke_local(struct ompi_communicator_t *comm);
+
 #if OPAL_ENABLE_FT_MPI
 /* Agreement */
 int ompi_coll_base_agree_noft(void *contrib,

ompi/mca/coll/base/coll_base_revoke_local.c

@@ -0,0 +1,47 @@
+/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil -*- */
+/*
+ * Copyright (c) 2004-2005 The Trustees of Indiana University and Indiana
+ *                         University Research and Technology
+ *                         Corporation.  All rights reserved.
+ * Copyright (c) 2004-2017 The University of Tennessee and The University
+ *                         of Tennessee Research Foundation.  All rights
+ *                         reserved.
+ * Copyright (c) 2004-2005 High Performance Computing Center Stuttgart,
+ *                         University of Stuttgart.  All rights reserved.
+ * Copyright (c) 2004-2005 The Regents of the University of California.
+ *                         All rights reserved.
+ * Copyright (c) 2013      Los Alamos National Security, LLC. All Rights
+ *                         reserved.
+ * Copyright (c) 2015-2016 Research Organization for Information Science
+ *                         and Technology (RIST). All rights reserved.
+ * Copyright (c) 2016-2017 IBM Corporation.  All rights reserved.
+ * Copyright (c) 2018      Siberian State University of Telecommunications
+ *                         and Information Science. All rights reserved.
+ * Copyright (c) 2022      Cisco Systems, Inc.  All rights reserved.
+ * Copyright (c) 2024      Stony Brook University.  All rights reserved.
+ * $COPYRIGHT$
+ *
+ * Additional copyrights may follow
+ *
+ * $HEADER$
+ */
+
+#include "ompi_config.h"
+
+#include "mpi.h"
+#include "opal/class/opal_list.h"
+#include "ompi/communicator/communicator.h"
+#include "ompi/mca/coll/base/coll_base_functions.h"
+#include "ompi/mca/coll/base/coll_base_util.h"
+
+int mca_coll_base_revoke_local(ompi_communicator_t* comm){
+    // Called on each initialized component, to give each the opportunity to
+    // revoke any subcomms
+    mca_coll_base_avail_coll_t* avail;
+    OPAL_LIST_FOREACH(avail, comm->c_coll->module_list, mca_coll_base_avail_coll_t){
+        if(NULL == avail->ac_module) continue;
+        if(NULL == avail->ac_module->coll_revoke_local) continue;
+        avail->ac_module->coll_revoke_local(comm, avail->ac_module);
+    }
+    return OMPI_SUCCESS;
+}

ompi/mca/coll/coll.h

@@ -484,6 +484,19 @@ typedef int (*mca_coll_base_module_reduce_local_fn_t)
     struct ompi_datatype_t * dtype, struct ompi_op_t * op,
     struct mca_coll_base_module_3_0_0_t *module);
 
+/*
+ * revoke_local
+ * Even though this is not a collective operation, it is related to the
+ * collectives. Adding to the framework allows a collective component the
+ * option of intercepting it to, e.g., also revoke sub-communicators
+ */
+typedef int (*mca_coll_base_module_revoke_local_fn_t)
+  (struct ompi_communicator_t* comm,
+   struct mca_coll_base_module_3_0_0_t *module);
+/* revoke applies to all coll modules, so the comm's function differs */
+typedef int (*mca_coll_base_comm_revoke_local_fn_t)
+  (struct ompi_communicator_t* comm);
+
 
 /* ******************************************************************** */
 
@@ -627,6 +640,8 @@ struct mca_coll_base_module_3_0_0_t {
 
     mca_coll_base_module_reduce_local_fn_t coll_reduce_local;
 
+    mca_coll_base_module_revoke_local_fn_t coll_revoke_local;
+
     /** Data storage for all the algorithms defined in the base. Should
         not be used by other modules */
     struct mca_coll_base_comm_t* base_data;
@@ -802,6 +817,8 @@ struct mca_coll_base_comm_coll_t {
     mca_coll_base_module_iagree_fn_t coll_iagree;
     mca_coll_base_module_3_0_0_t *coll_iagree_module;
 
+    mca_coll_base_comm_revoke_local_fn_t coll_revoke_local;
+
     /* List of modules initialized, queried and enabled */
     opal_list_t *module_list;
 };

ompi/mca/coll/han/coll_han.h

@@ -551,6 +551,9 @@ mca_coll_han_scatter_intra_dynamic(SCATTER_BASE_ARGS,
 int
 mca_coll_han_scatterv_intra_dynamic(SCATTERV_BASE_ARGS,
                                     mca_coll_base_module_t *module);
+int
+mca_coll_han_revoke_local(struct ompi_communicator_t *comm,
+                          mca_coll_base_module_t *module);
 
 int mca_coll_han_barrier_intra_simple(struct ompi_communicator_t *comm,
                                       mca_coll_base_module_t *module);

ompi/mca/coll/han/coll_han_module.c

@@ -92,6 +92,8 @@ static void mca_coll_han_module_construct(mca_coll_han_module_t * module)
     module->dynamic_errors = 0;
 
     han_module_clear(module);
+
+    module->super.coll_revoke_local = mca_coll_han_revoke_local;
 }
 
 /*
@@ -257,6 +259,7 @@ mca_coll_han_comm_query(struct ompi_communicator_t * comm, int *priority)
         /* We are on a topologic sub-communicator, return only the selector */
         han_module->super.coll_allgatherv = mca_coll_han_allgatherv_intra_dynamic;
     }
+    han_module->super.coll_revoke_local = mca_coll_han_revoke_local;
 
     opal_output_verbose(10, ompi_coll_base_framework.framework_output,
                         "coll:han:comm_query (%s/%s): pick me! pick me!",

ompi/mca/coll/han/coll_han_subcomms.c

@@ -414,3 +414,26 @@ int mca_coll_han_comm_create(struct ompi_communicator_t *comm,
     OBJ_DESTRUCT(&comm_info);
     return OMPI_SUCCESS;
 }
+
+int mca_coll_han_revoke_local(ompi_communicator_t *comm,
+                              mca_coll_base_module_t *module)
+{
+    mca_coll_han_module_t *han_module = (mca_coll_han_module_t*) module;
+    for(int i = 0; i < NB_TOPO_LVL; i++){
+        if(NULL == han_module->sub_comm[i]) continue;
+        ompi_comm_revoke_internal(han_module->sub_comm[i]);
+    }
+    if(han_module->cached_low_comms != NULL){
+        for(int i = 0; i < COLL_HAN_LOW_MODULES; i++){
+            if(NULL == han_module->cached_low_comms[i]) continue;
+            ompi_comm_revoke_internal(han_module->cached_low_comms[i]);
+        }
+    }
+    if(han_module->cached_up_comms != NULL){
+        for(int i = 0; i < COLL_HAN_LOW_MODULES; i++){
+            if(NULL == han_module->cached_low_comms[i]) continue;
+            ompi_comm_revoke_internal(han_module->cached_low_comms[i]);
+        }
+    }
+    return MPI_SUCCESS;
+}