Redesign how we interface with communicators and PETSc (#4766)

* Redesign how we interface with communicators and PETSc.

Motivation:
* These changes will make it much harder to accidentally stash PETSc
  objects on the wrong communicators.
* It avoids a situation where there were 2 PETSc communicators holding
  objects to destroy (one dup'ed from MPI_COMM_WORLD and the other
  dup-ed from PYOP2_COMM_WORLD).

Core changes:
* firedrake.COMM_WORLD is now *exactly* MPI.COMM_WORLD, not a duplicate.
* Objects do not have user and internal comm attributes. If
  Firedrake/PyOP2 needs to send a message we should use the
  `temp_internal_comm` method.
* Added a test.

Author: connorjward

docs/notebooks/03-elasticity.ipynb

@@ -328,7 +328,7 @@
       "\u001b[36mFile \u001b[39m\u001b[32mpetsc4py/PETSc/Log.pyx:188\u001b[39m, in \u001b[36mpetsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func\u001b[39m\u001b[34m()\u001b[39m\n",
       "\u001b[36mFile \u001b[39m\u001b[32mpetsc4py/PETSc/Log.pyx:189\u001b[39m, in \u001b[36mpetsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func\u001b[39m\u001b[34m()\u001b[39m\n",
       "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/adjoint_utils/variational_solver.py:108\u001b[39m, in \u001b[36mNonlinearVariationalSolverMixin._ad_annotate_solve.<locals>.wrapper\u001b[39m\u001b[34m(self, **kwargs)\u001b[39m\n\u001b[32m    105\u001b[39m     tape.add_block(block)\n\u001b[32m    107\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m stop_annotating():\n\u001b[32m--> \u001b[39m\u001b[32m108\u001b[39m     out = \u001b[43msolve\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    110\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m annotate:\n\u001b[32m    111\u001b[39m     block.add_output(\u001b[38;5;28mself\u001b[39m._ad_problem._ad_u.create_block_variable())\n",
-      "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/variational_solver.py:361\u001b[39m, in \u001b[36mNonlinearVariationalSolver.solve\u001b[39m\u001b[34m(self, bounds)\u001b[39m\n\u001b[32m    359\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m problem.restrict:\n\u001b[32m    360\u001b[39m     problem.u.interpolate(problem.u_restrict)\n\u001b[32m--> \u001b[39m\u001b[32m361\u001b[39m \u001b[43msolving_utils\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcheck_snes_convergence\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43msnes\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    363\u001b[39m \u001b[38;5;66;03m# Grab the comm associated with the `_problem` and call PETSc's garbage cleanup routine\u001b[39;00m\n\u001b[32m    364\u001b[39m comm = \u001b[38;5;28mself\u001b[39m._problem.u_restrict.function_space().mesh()._comm\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/variational_solver.py:361\u001b[39m, in \u001b[36mNonlinearVariationalSolver.solve\u001b[39m\u001b[34m(self, bounds)\u001b[39m\n\u001b[32m    359\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m problem.restrict:\n\u001b[32m    360\u001b[39m     problem.u.interpolate(problem.u_restrict)\n\u001b[32m--> \u001b[39m\u001b[32m361\u001b[39m \u001b[43msolving_utils\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcheck_snes_convergence\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43msnes\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    363\u001b[39m \u001b[38;5;66;03m# Grab the comm associated with the `_problem` and call PETSc's garbage cleanup routine\u001b[39;00m\n\u001b[32m    364\u001b[39m comm = \u001b[38;5;28mself\u001b[39m._problem.u_restrict.function_space().mesh().comm\n",
       "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/solving_utils.py:128\u001b[39m, in \u001b[36mcheck_snes_convergence\u001b[39m\u001b[34m(snes)\u001b[39m\n\u001b[32m    126\u001b[39m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    127\u001b[39m             msg = reason\n\u001b[32m--> \u001b[39m\u001b[32m128\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m ConvergenceError(\u001b[33mr\u001b[39m\u001b[33m\"\"\"\u001b[39m\u001b[33mNonlinear solve failed to converge after \u001b[39m\u001b[38;5;132;01m%d\u001b[39;00m\u001b[33m nonlinear iterations.\u001b[39m\n\u001b[32m    129\u001b[39m \u001b[33mReason:\u001b[39m\n\u001b[32m    130\u001b[39m \u001b[33m   \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[33m\"\"\"\u001b[39m % (snes.getIterationNumber(), msg))\n",
       "\u001b[31mConvergenceError\u001b[39m: Nonlinear solve failed to converge after 0 nonlinear iterations.\nReason:\n   DIVERGED_LINEAR_SOLVE"
      ]

docs/source/parallelism.rst

@@ -96,9 +96,5 @@ different simulations on the two halves we would write.
 To access the communicator a mesh was created on, we can use the
 ``mesh.comm`` property, or the function ``mesh.mpi_comm``.
 
-.. warning::
-  Do not use the internal ``mesh._comm`` attribute for communication.
-  This communicator is for internal Firedrake MPI communication only.
-
 .. _MPI: http://mpi-forum.org/
 .. _STREAMS: http://www.cs.virginia.edu/stream/

firedrake/adjoint_utils/function.py

@@ -1,4 +1,5 @@
 from functools import wraps
+from pyop2.mpi import temp_internal_comm
 import ufl
 from ufl.domain import extract_unique_domain
 from pyadjoint.overloaded_type import create_overloaded_object, FloatingType
@@ -280,8 +281,9 @@ def _ad_assign_numpy(dst, src, offset):
         m_a_local = src[offset + range_begin:offset + range_end]
         if dst.function_space().ufl_element().family() == "Real":
             # Real space keeps a redundant copy of the data on every rank
-            comm = dst.function_space().mesh()._comm
-            dst.dat.data_wo[...] = comm.bcast(m_a_local, root=0)
+            comm = dst.function_space().mesh().comm
+            with temp_internal_comm(comm) as icomm:
+                dst.dat.data_wo[...] = icomm.bcast(m_a_local, root=0)
         else:
             dst.dat.data_wo[...] = m_a_local.reshape(dst.dat.data_wo.shape)
         offset += dst.dat.dataset.layout_vec.size

firedrake/assemble.py

@@ -1168,13 +1168,13 @@ def __init__(self, form, form_compiler_parameters=None):
 
     def allocate(self):
         # Getting the comm attribute of a form isn't straightforward
-        # form.ufl_domains()[0]._comm seems the most robust method
+        # form.ufl_domains()[0].comm seems the most robust method
         # revisit in a refactor
         return op2.Global(
             1,
             [0.0],
             dtype=utils.ScalarType,
-            comm=self._form.ufl_domains()[0]._comm
+            comm=self._form.ufl_domains()[0].comm
         )
 
     def _apply_bc(self, tensor, bc, u=None):

firedrake/checkpointing.py

@@ -3,7 +3,7 @@
 from petsc4py.PETSc import ViewerHDF5
 import finat.ufl
 from pyop2 import op2
-from pyop2.mpi import COMM_WORLD, internal_comm, MPI
+from pyop2.mpi import COMM_WORLD, MPI
 from petsctools import OptionsManager
 from firedrake.cython import hdf5interface as h5i
 from firedrake.cython import dmcommon
@@ -104,7 +104,6 @@ def __init__(self, basename, single_file=True,
             warnings.warn("DumbCheckpoint class will soon be deprecated; use CheckpointFile class instead.",
                           DeprecationWarning)
         self.comm = comm or COMM_WORLD
-        self._comm = internal_comm(self.comm, self)
         self.mode = mode
 
         self._single = single_file
@@ -195,7 +194,7 @@ def new_file(self, name=None):
         if mode == FILE_UPDATE and not exists:
             mode = FILE_CREATE
         self._vwr = PETSc.ViewerHDF5().create(name, mode=mode,
-                                              comm=self._comm)
+                                              comm=self.comm)
         if self.mode == FILE_READ:
             nprocs = self.read_attribute("/", "nprocs")
             if nprocs != self.comm.size:
@@ -379,7 +378,6 @@ def __init__(self, filename, file_mode, comm=None):
             warnings.warn("HDF5File class will soon be deprecated; use CheckpointFile class instead.",
                           DeprecationWarning)
         self.comm = comm or COMM_WORLD
-        self._comm = internal_comm(self.comm, self)
 
         self._filename = filename
         self._mode = file_mode
@@ -397,7 +395,7 @@ def __init__(self, filename, file_mode, comm=None):
 
         # Try to use MPI
         try:
-            self._h5file = h5py.File(filename, file_mode, driver="mpio", comm=self._comm)
+            self._h5file = h5py.File(filename, file_mode, driver="mpio", comm=self.comm)
         except NameError:  # the error you get if h5py isn't compiled against parallel HDF5
             raise RuntimeError("h5py *must* be installed with MPI support")
 
@@ -527,10 +525,9 @@ def __init__(self, filename, mode, comm=COMM_WORLD):
         self.viewer = ViewerHDF5()
         self.filename = filename
         self.comm = comm
-        self._comm = internal_comm(comm, self)
         r"""The neme of the checkpoint file."""
-        self.viewer.create(filename, mode=mode, comm=self._comm)
-        self.commkey = self._comm.py2f()
+        self.viewer.create(filename, mode=mode, comm=self.comm)
+        self.commkey = self.comm.py2f()
         assert self.commkey != MPI.COMM_NULL.py2f()
         self._function_spaces = {}
         self._function_load_utils = {}
@@ -597,7 +594,7 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None):
                     layers_tV = impl.FunctionSpace(base_tmesh, element)
                     self._save_function_space_topology(layers_tV)
                     # Note that _cell_numbering coincides with DG0 section, so we can use tmesh.layers directly.
-                    layers_iset = PETSc.IS().createGeneral(tmesh.layers[:tmesh.cell_set.size, :], comm=tmesh._comm)
+                    layers_iset = PETSc.IS().createGeneral(tmesh.layers[:tmesh.cell_set.size, :], comm=tmesh.comm)
                     layers_iset.setName("_".join([PREFIX_EXTRUDED, "layers_iset"]))
                     self.viewer.pushGroup(path)
                     layers_iset.view(self.viewer)
@@ -661,7 +658,7 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None):
                     reflected = o_r_map[tmesh.entity_orientations[:tmesh.cell_set.size, -1]]
                     reflected_indices = (reflected == 1)
                     canonical_cell_orientations[reflected_indices] = 1 - canonical_cell_orientations[reflected_indices]
-                    cell_orientations_iset = PETSc.IS().createGeneral(canonical_cell_orientations, comm=tmesh._comm)
+                    cell_orientations_iset = PETSc.IS().createGeneral(canonical_cell_orientations, comm=tmesh.comm)
                     cell_orientations_iset.setName("_".join([PREFIX_IMMERSED, "cell_orientations_iset"]))
                     self.viewer.pushGroup(path)
                     cell_orientations_iset.view(self.viewer)
@@ -1065,7 +1062,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
                 _, _, lsf = self._function_load_utils[base_tmesh_key + sd_key]
                 nroots, _, _ = lsf.getGraph()
                 layers_a = np.empty(nroots, dtype=utils.IntType)
-                layers_a_iset = PETSc.IS().createGeneral(layers_a, comm=self._comm)
+                layers_a_iset = PETSc.IS().createGeneral(layers_a, comm=self.comm)
                 layers_a_iset.setName("_".join([PREFIX_EXTRUDED, "layers_iset"]))
                 self.viewer.pushGroup(path)
                 layers_a_iset.load(self.viewer)
@@ -1128,7 +1125,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
                 _, _, lsf = self._function_load_utils[tmesh_key + sd_key]
                 nroots, _, _ = lsf.getGraph()
                 cell_orientations_a = np.empty(nroots, dtype=utils.IntType)
-                cell_orientations_a_iset = PETSc.IS().createGeneral(cell_orientations_a, comm=self._comm)
+                cell_orientations_a_iset = PETSc.IS().createGeneral(cell_orientations_a, comm=self.comm)
                 cell_orientations_a_iset.setName("_".join([PREFIX_IMMERSED, "cell_orientations_iset"]))
                 self.viewer.pushGroup(path)
                 cell_orientations_a_iset.load(self.viewer)
@@ -1165,7 +1162,7 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
             _distribution_name, = self.h5pyfile[path].keys()
             path = self._path_to_distribution(tmesh_name, _distribution_name)
             _comm_size = self.get_attr(path, "comm_size")
-            if _comm_size == self._comm.size and \
+            if _comm_size == self.comm.size and \
                distribution_parameters is None and reorder is None:
                 load_distribution_permutation = True
         if load_distribution_permutation:
@@ -1185,7 +1182,7 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
             permutation_name = None
             perm_is = None
         plex = PETSc.DMPlex()
-        plex.create(comm=self._comm)
+        plex.create(comm=self.comm)
         plex.setName(tmesh_name)
         # Check format
         path = os.path.join(self._path_to_topology(tmesh_name), "topology")
@@ -1203,15 +1200,15 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
         plex.removeLabel("pyop2_ghost")
         if load_distribution_permutation:
             chart_size = np.empty(1, dtype=utils.IntType)
-            chart_sizes_iset = PETSc.IS().createGeneral(chart_size, comm=self._comm)
+            chart_sizes_iset = PETSc.IS().createGeneral(chart_size, comm=self.comm)
             chart_sizes_iset.setName("chart_sizes")
             path = self._path_to_distribution(tmesh_name, distribution_name)
             self.viewer.pushGroup(path)
             chart_sizes_iset.load(self.viewer)
             self.viewer.popGroup()
             chart_size = chart_sizes_iset.getIndices().item()
             perm = np.empty(chart_size, dtype=utils.IntType)
-            perm_is = PETSc.IS().createGeneral(perm, comm=self._comm)
+            perm_is = PETSc.IS().createGeneral(perm, comm=self.comm)
             path = self._path_to_permutation(tmesh_name, distribution_name, permutation_name)
             self.viewer.pushGroup(path)
             perm_is.setName("permutation")
@@ -1276,10 +1273,10 @@ def _load_function_space_topology(self, tmesh, element):
         sd_key = self._get_shared_data_key_for_checkpointing(tmesh, element)
         if tmesh_key + sd_key not in self._function_load_utils:
             topology_dm = tmesh.topology_dm
-            dm = PETSc.DMShell().create(comm=tmesh._comm)
+            dm = PETSc.DMShell().create(comm=tmesh.comm)
             dm.setName(self._get_dm_name_for_checkpointing(tmesh, element))
             dm.setPointSF(topology_dm.getPointSF())
-            section = PETSc.Section().create(comm=tmesh._comm)
+            section = PETSc.Section().create(comm=tmesh.comm)
             section.setPermutation(tmesh._dm_renumbering)
             dm.setSection(section)
             base_tmesh = tmesh._base_mesh if isinstance(tmesh, ExtrudedMeshTopology) else tmesh
@@ -1446,7 +1443,7 @@ def _get_dm_for_checkpointing(self, tV):
             nodes_per_entity, real_tensorproduct, block_size = sd_key
             global_numbering, _ = tV.mesh().create_section(nodes_per_entity, real_tensorproduct, block_size=block_size)
             topology_dm = tV.mesh().topology_dm
-            dm = PETSc.DMShell().create(tV.mesh()._comm)
+            dm = PETSc.DMShell().create(tV.mesh().comm)
             dm.setPointSF(topology_dm.getPointSF())
             dm.setSection(global_numbering)
         else:

firedrake/cofunction.py

@@ -3,7 +3,7 @@
 import ufl
 
 from ufl.form import BaseForm
-from pyop2 import op2, mpi
+from pyop2 import op2
 from pyadjoint.tape import stop_annotating, annotate_tape, get_working_tape
 from finat.ufl import MixedElement
 import firedrake.assemble
@@ -65,18 +65,16 @@ def __init__(self, function_space, val=None, name=None, dtype=ScalarType,
 
         # User comm
         self.comm = V.comm
-        # Internal comm
-        self._comm = mpi.internal_comm(V.comm, self)
         self._function_space = V
-        self.uid = utils._new_uid(self._comm)
+        self.uid = utils._new_uid(self.comm)
         self._name = name or 'cofunction_%d' % self.uid
         self._label = "a cofunction"
 
         if isinstance(val, Cofunction):
             val = val.dat
 
         if isinstance(val, (op2.Dat, op2.DatView, op2.MixedDat, op2.Global)):
-            assert val.comm == self._comm
+            assert val.comm == self.comm
             self.dat = val
         else:
             self.dat = function_space.make_dat(val, dtype, self.name())

firedrake/cython/dmcommon.pyx

@@ -1263,7 +1263,7 @@ def create_section(mesh, nodes_per_entity, on_base=False, block_size=1, boundary
                 nodes = sum(nodes_per_entity[:, i]*(mesh.layers - i) for i in range(2)).reshape(dimension + 1, -1)
     else:
         nodes = nodes_per_entity.reshape(dimension + 1, -1)
-    section = PETSc.Section().create(comm=mesh._comm)
+    section = PETSc.Section().create(comm=mesh.comm)
     get_chart(dm.dm, &pStart, &pEnd)
     section.setChart(pStart, pEnd)
 

firedrake/dmhooks.py

@@ -393,7 +393,7 @@ def create_subdm(dm, fields, *args, **kwargs):
         # Index set mapping from W into subspace.
         iset = PETSc.IS().createGeneral(numpy.concatenate([W._ises[f].indices
                                                            for f in fields]),
-                                        comm=W._comm)
+                                        comm=W.comm)
         if ctx is not None:
             ctx, = ctx.split([fields])
             add_hook(parent, setup=partial(push_appctx, subspace.dm, ctx),

firedrake/ensemble/ensemble.py

@@ -21,14 +21,12 @@ def _mpi_dispatch(self, *args, **kwargs):
                for arg in [*args, *kwargs.values()]):
             return func(self, *args, **kwargs)
         else:
-            mpicall = getattr(
-                self._ensemble_comm,
-                func.__name__)
+            mpicall = getattr(self._ensemble_comm, func.__name__)
             return mpicall(*args, **kwargs)
     return _mpi_dispatch
 
 
-class Ensemble(object):
+class Ensemble:
     def __init__(self, comm, M, **kwargs):
         """
         Create a set of space and ensemble subcommunicators.
@@ -47,22 +45,23 @@ def __init__(self, comm, M, **kwargs):
 
         # Global comm
         self.global_comm = comm
-        # Internal global comm
-        self._comm = internal_comm(comm, self)
 
         ensemble_name = kwargs.get("ensemble_name", "Ensemble")
         # User and internal communicator for spatial parallelism, contains a
         # contiguous chunk of M processes from `global_comm`.
         self.comm = self.global_comm.Split(color=(rank // M), key=rank)
         self.comm.name = f"{ensemble_name} spatial comm"
         weakref.finalize(self, self.comm.Free)
-        self._spatial_comm = internal_comm(self.comm, self)
 
         # User and internal communicator for ensemble parallelism, contains all
         # processes in `global_comm` which have the same rank in `comm`.
         self.ensemble_comm = self.global_comm.Split(color=(rank % M), key=rank)
         self.ensemble_comm.name = f"{ensemble_name} ensemble comm"
         weakref.finalize(self, self.ensemble_comm.Free)
+        # Keep a reference to the internal communicator because some methods return
+        # non-blocking requests and we need to avoid cleaning up the communicator before
+        # they complete. Note that this communicator should *never* be passed to PETSc, as
+        # objects created with the communicator will never get cleaned up.
         self._ensemble_comm = internal_comm(self.ensemble_comm, self)
 
         assert self.comm.size == M
@@ -90,11 +89,11 @@ def _check_function(self, f, g=None):
         :raises ValueError: if function communicators mismatch each other or the ensemble
             spatial communicator, or is the functions are in different spaces
         """
-        if MPI.Comm.Compare(f._comm, self._spatial_comm) not in {MPI.CONGRUENT, MPI.IDENT}:
+        if MPI.Comm.Compare(f.comm, self.comm) not in {MPI.CONGRUENT, MPI.IDENT}:
             raise ValueError("Function communicator does not match space communicator")
 
         if g is not None:
-            if MPI.Comm.Compare(f._comm, g._comm) not in {MPI.CONGRUENT, MPI.IDENT}:
+            if MPI.Comm.Compare(f.comm, g.comm) not in {MPI.CONGRUENT, MPI.IDENT}:
                 raise ValueError("Mismatching communicators for functions")
             if f.function_space() != g.function_space():
                 raise ValueError("Mismatching function spaces for functions")
@@ -206,7 +205,6 @@ def ibcast(self, f, root=0):
         :raises ValueError: if function communicator mismatches the ensemble spatial communicator.
         """
         self._check_function(f)
-
         return [self._ensemble_comm.Ibcast(dat.data, root=root)
                 for dat in f.dat]