get mpi_communicator from actx instead of dcoll

Author: majosm

grudge/reductions.py

@@ -80,6 +80,7 @@
 from pytools import memoize_in
 
 import grudge.dof_desc as dof_desc
+from grudge.array_context import MPIBasedArrayContext
 from grudge.discretization import DiscretizationCollection
 
 
@@ -128,16 +129,17 @@ def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
         :class:`~arraycontext.ArrayContainer`.
     :returns: a device scalar denoting the nodal sum.
     """
-    comm = dcoll.mpi_communicator
-    if comm is None:
+    from arraycontext import get_container_context_recursively
+    actx = get_container_context_recursively(vec)
+
+    if not isinstance(actx, MPIBasedArrayContext):
         return nodal_sum_loc(dcoll, dd, vec)
 
+    comm = actx.mpi_communicator
+
     # NOTE: Do not move, we do not want to import mpi4py in single-rank computations
     from mpi4py import MPI
 
-    from arraycontext import get_container_context_recursively
-    actx = get_container_context_recursively(vec)
-
     return actx.from_numpy(
         comm.allreduce(actx.to_numpy(nodal_sum_loc(dcoll, dd, vec)), op=MPI.SUM))
 
@@ -174,13 +176,16 @@ def nodal_min(dcoll: DiscretizationCollection, dd, vec, *, initial=None) -> Scal
     :arg initial: an optional initial value. Defaults to `numpy.inf`.
     :returns: a device scalar denoting the nodal minimum.
     """
-    comm = dcoll.mpi_communicator
-    if comm is None:
+    from arraycontext import get_container_context_recursively
+    actx = get_container_context_recursively(vec)
+
+    if not isinstance(actx, MPIBasedArrayContext):
         return nodal_min_loc(dcoll, dd, vec, initial=initial)
 
+    comm = actx.mpi_communicator
+
     # NOTE: Do not move, we do not want to import mpi4py in single-rank computations
     from mpi4py import MPI
-    actx = vec.array_context
 
     return actx.from_numpy(
         comm.allreduce(
@@ -231,13 +236,16 @@ def nodal_max(dcoll: DiscretizationCollection, dd, vec, *, initial=None) -> Scal
     :arg initial: an optional initial value. Defaults to `-numpy.inf`.
     :returns: a device scalar denoting the nodal maximum.
     """
-    comm = dcoll.mpi_communicator
-    if comm is None:
+    from arraycontext import get_container_context_recursively
+    actx = get_container_context_recursively(vec)
+
+    if not isinstance(actx, MPIBasedArrayContext):
         return nodal_max_loc(dcoll, dd, vec, initial=initial)
 
+    comm = actx.mpi_communicator
+
     # NOTE: Do not move, we do not want to import mpi4py in single-rank computations
     from mpi4py import MPI
-    actx = vec.array_context
 
     return actx.from_numpy(
         comm.allreduce(

grudge/trace_pair.py

@@ -410,7 +410,7 @@ def __init__(self,
         bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank))
 
         local_bdry_data = project(dcoll, volume_dd, bdry_dd, array_container)
-        comm = dcoll.mpi_communicator
+        comm = actx.mpi_communicator
         assert comm is not None
 
         self.dcoll = dcoll