Add optional user-specified comm_tag to operators

mirgecom/artificial_viscosity.py

@@ -149,10 +149,6 @@
 import grudge.op as op
 
 
-class _AVCVTag:
-    pass
-
-
 class _AVRTag:
     pass
 
@@ -161,7 +157,7 @@ def av_laplacian_operator(dcoll, boundaries, fluid_state, alpha, gas_model=None,
                           kappa=1., s0=-6., time=0, operator_states_quad=None,
                           grad_cv=None, quadrature_tag=None, boundary_kwargs=None,
                           indicator=None, divergence_numerical_flux=num_flux_central,
-                          **kwargs):
+                          comm_tag=None, **kwargs):
     r"""Compute the artificial viscosity right-hand-side.
 
     Computes the the right-hand-side term for artificial viscosity.
@@ -205,6 +201,9 @@ def av_laplacian_operator(dcoll, boundaries, fluid_state, alpha, gas_model=None,
     boundary_kwargs: :class:`dict`
         dictionary of extra arguments to pass through to the boundary conditions
 
+    comm_tag: Hashable
+        Tag for distributed communication
+
     Returns
     -------
     :class:`mirgecom.fluid.ConservedVars`
@@ -233,7 +232,7 @@ def interp_to_vol_quad(u):
     if operator_states_quad is None:
         from mirgecom.gas_model import make_operator_fluid_states
         operator_states_quad = make_operator_fluid_states(
-            dcoll, fluid_state, gas_model, boundaries, quadrature_tag)
+            dcoll, fluid_state, gas_model, boundaries, quadrature_tag, comm_tag)
 
     vol_state_quad, inter_elem_bnd_states_quad, domain_bnd_states_quad = \
         operator_states_quad
@@ -247,7 +246,8 @@ def interp_to_vol_quad(u):
         from mirgecom.navierstokes import grad_cv_operator
         grad_cv = grad_cv_operator(dcoll, gas_model, boundaries, fluid_state,
                                    time=time, quadrature_tag=quadrature_tag,
-                                   operator_states_quad=operator_states_quad)
+                                   operator_states_quad=operator_states_quad,
+                                   comm_tag=comm_tag)
 
     # Compute R = alpha*grad(Q)
     r = -alpha * indicator * grad_cv
@@ -264,7 +264,8 @@ def central_flux_div(utpair):
     r_bnd = (
         # Rank-local and cross-rank (across parallel partitions) contributions
         + sum(central_flux_div(interp_to_surf_quad(tpair=tpair))
-              for tpair in interior_trace_pairs(dcoll, r, tag=_AVRTag))
+              for tpair in interior_trace_pairs(dcoll, r,
+                    comm_tag=(_AVRTag, comm_tag)))
 
         # Contributions from boundary fluxes
         + sum(boundaries[btag].av_flux(

mirgecom/diffusion.py

@@ -226,7 +226,8 @@ class _DiffusionGradTag:
     pass
 
 
-def grad_operator(dcoll, boundaries, u, quadrature_tag=DISCR_TAG_BASE):
+def grad_operator(
+        dcoll, boundaries, u, quadrature_tag=DISCR_TAG_BASE, comm_tag=None):
     r"""
     Compute the gradient of *u*.
 
@@ -245,6 +246,8 @@ def grad_operator(dcoll, boundaries, u, quadrature_tag=DISCR_TAG_BASE):
     quadrature_tag:
         quadrature tag indicating which discretization in *dcoll* to use for
         overintegration
+    comm_tag: Hashable
+        Tag for distributed communication
 
     Returns
     -------
@@ -276,7 +279,7 @@ def grad_operator(dcoll, boundaries, u, quadrature_tag=DISCR_TAG_BASE):
             sum(
                 grad_flux(dcoll, u_tpair, quadrature_tag=quadrature_tag)
                 for u_tpair in interior_trace_pairs(
-                    dcoll, u, comm_tag=_DiffusionStateTag))
+                    dcoll, u, comm_tag=(_DiffusionStateTag, comm_tag)))
             + sum(
                 bdry.get_grad_flux(
                     dcoll, as_dofdesc(btag), u, quadrature_tag=quadrature_tag)
@@ -325,7 +328,7 @@ def _normalize_arguments(*args, **kwargs):
     return kappa, boundaries, u, quadrature_tag
 
 
-def diffusion_operator(dcoll, *args, return_grad_u=False, **kwargs):
+def diffusion_operator(dcoll, *args, return_grad_u=False, comm_tag=None, **kwargs):
     r"""
     Compute the diffusion operator.
 
@@ -352,6 +355,8 @@ def diffusion_operator(dcoll, *args, return_grad_u=False, **kwargs):
     quadrature_tag:
         quadrature tag indicating which discretization in *dcoll* to use for
         overintegration
+    comm_tag: Hashable
+        Tag for distributed communication
 
     Returns
     -------
@@ -395,8 +400,10 @@ def diffusion_operator(dcoll, *args, return_grad_u=False, **kwargs):
                 diffusion_flux(
                     dcoll, kappa_tpair, grad_u_tpair, quadrature_tag=quadrature_tag)
                 for kappa_tpair, grad_u_tpair in zip(
-                    interior_trace_pairs(dcoll, kappa, comm_tag=_DiffusionKappaTag),
-                    interior_trace_pairs(dcoll, grad_u, comm_tag=_DiffusionGradTag)))
+                    interior_trace_pairs(
+                        dcoll, kappa, comm_tag=(_DiffusionKappaTag, comm_tag)),
+                    interior_trace_pairs(
+                        dcoll, grad_u, comm_tag=(_DiffusionGradTag, comm_tag))))
             + sum(
                 bdry.get_diffusion_flux(
                     dcoll, as_dofdesc(btag), kappa, grad_u,

mirgecom/euler.py

@@ -68,7 +68,7 @@
 
 def euler_operator(dcoll, state, gas_model, boundaries, time=0.0,
                    inviscid_numerical_flux_func=inviscid_facial_flux_rusanov,
-                   quadrature_tag=None, operator_states_quad=None):
+                   quadrature_tag=None, comm_tag=None, operator_states_quad=None):
     r"""Compute RHS of the Euler flow equations.
 
     Returns
@@ -107,13 +107,16 @@ def euler_operator(dcoll, state, gas_model, boundaries, time=0.0,
         An optional identifier denoting a particular quadrature
         discretization to use during operator evaluations.
         The default value is *None*.
+
+    comm_tag: Hashable
+        Tag for distributed communication
     """
     dd_quad_vol = DOFDesc("vol", quadrature_tag)
     dd_quad_faces = DOFDesc("all_faces", quadrature_tag)
 
     if operator_states_quad is None:
-        operator_states_quad = make_operator_fluid_states(dcoll, state, gas_model,
-                                                          boundaries, quadrature_tag)
+        operator_states_quad = make_operator_fluid_states(
+            dcoll, state, gas_model, boundaries, quadrature_tag, comm_tag)
 
     volume_state_quad, interior_state_pairs_quad, domain_boundary_states_quad = \
         operator_states_quad

mirgecom/gas_model.py

@@ -385,7 +385,7 @@ class _FluidTemperatureTag:
 
 
 def make_operator_fluid_states(dcoll, volume_state, gas_model, boundaries,
-                               quadrature_tag=None):
+                               quadrature_tag=None, comm_tag=None):
     """Prepare gas model-consistent fluid states for use in fluid operators.
 
     This routine prepares a model-consistent fluid state for each of the volume and
@@ -421,6 +421,9 @@ def make_operator_fluid_states(dcoll, volume_state, gas_model, boundaries,
         discretization to use during operator evaluations.
         The default value is *None*.
 
+    comm_tag: Hashable
+        Tag for distributed communication
+
     Returns
     -------
     (:class:`~mirgecom.gas_model.FluidState`, :class:`~grudge.trace_pair.TracePair`,
@@ -448,7 +451,8 @@ def make_operator_fluid_states(dcoll, volume_state, gas_model, boundaries,
         # Get the interior trace pairs onto the surface quadrature
         # discretization (if any)
         interp_to_surf_quad(tpair=tpair)
-        for tpair in interior_trace_pairs(dcoll, volume_state.cv, tag=_FluidCVTag)
+        for tpair in interior_trace_pairs(dcoll, volume_state.cv,
+            comm_tag=(_FluidCVTag, comm_tag))
     ]
 
     tseed_interior_pairs = None
@@ -461,8 +465,9 @@ def make_operator_fluid_states(dcoll, volume_state, gas_model, boundaries,
             # Get the interior trace pairs onto the surface quadrature
             # discretization (if any)
             interp_to_surf_quad(tpair=tpair)
-            for tpair in interior_trace_pairs(dcoll, volume_state.temperature,
-                                              tag=_FluidTemperatureTag)]
+            for tpair in interior_trace_pairs(
+                dcoll, volume_state.temperature,
+                comm_tag=(_FluidTemperatureTag, comm_tag))]
 
     interior_boundary_states_quad = \
         make_fluid_state_trace_pairs(cv_interior_pairs, gas_model,

mirgecom/navierstokes.py

@@ -107,7 +107,7 @@ def _gradient_flux_interior(dcoll, numerical_flux_func, tpair):
 def grad_cv_operator(
         dcoll, gas_model, boundaries, state, *, time=0.0,
         numerical_flux_func=num_flux_central,
-        quadrature_tag=DISCR_TAG_BASE,
+        quadrature_tag=DISCR_TAG_BASE, comm_tag=None,
         # Added to avoid repeated computation
         # FIXME: See if there's a better way to do this
         operator_states_quad=None):
@@ -140,6 +140,9 @@ def grad_cv_operator(
         An identifier denoting a particular quadrature discretization to use during
         operator evaluations.
 
+    comm_tag: Hashable
+        Tag for distributed communication
+
     Returns
     -------
     :class:`~mirgecom.fluid.ConservedVars`
@@ -152,7 +155,8 @@ def grad_cv_operator(
 
     if operator_states_quad is None:
         operator_states_quad = make_operator_fluid_states(
-            dcoll, state, gas_model, boundaries, quadrature_tag)
+            dcoll, state, gas_model, boundaries, quadrature_tag,
+            comm_tag)
 
     vol_state_quad, inter_elem_bnd_states_quad, domain_bnd_states_quad = \
         operator_states_quad
@@ -194,7 +198,7 @@ def grad_cv_operator(
 def grad_t_operator(
         dcoll, gas_model, boundaries, state, *, time=0.0,
         numerical_flux_func=num_flux_central,
-        quadrature_tag=DISCR_TAG_BASE,
+        quadrature_tag=DISCR_TAG_BASE, comm_tag=None,
         # Added to avoid repeated computation
         # FIXME: See if there's a better way to do this
         operator_states_quad=None):
@@ -227,6 +231,9 @@ def grad_t_operator(
         An identifier denoting a particular quadrature discretization to use during
         operator evaluations.
 
+    comm_tag: Hashable
+        Tag for distributed communication
+
     Returns
     -------
     :class:`numpy.ndarray`
@@ -239,7 +246,7 @@ def grad_t_operator(
 
     if operator_states_quad is None:
         operator_states_quad = make_operator_fluid_states(
-            dcoll, state, gas_model, boundaries, quadrature_tag)
+            dcoll, state, gas_model, boundaries, quadrature_tag, comm_tag)
 
     vol_state_quad, inter_elem_bnd_states_quad, domain_bnd_states_quad = \
         operator_states_quad
@@ -287,6 +294,7 @@ def ns_operator(dcoll, gas_model, state, boundaries, *, time=0.0,
                 gradient_numerical_flux_func=num_flux_central,
                 viscous_numerical_flux_func=viscous_facial_flux_central,
                 quadrature_tag=DISCR_TAG_BASE, return_gradients=False,
+                comm_tag=None,
                 # Added to avoid repeated computation
                 # FIXME: See if there's a better way to do this
                 operator_states_quad=None,
@@ -352,6 +360,9 @@ def ns_operator(dcoll, gas_model, state, boundaries, *, time=0.0,
         $\nabla(\text{CV})$ and $\nabla(T)$ along with the RHS for the Navier-Stokes
         equations.  Useful for debugging and visualization.
 
+    comm_tag: Hashable
+        Tag for distributed communication
+
     Returns
     -------
     :class:`mirgecom.fluid.ConservedVars`
@@ -378,7 +389,7 @@ def ns_operator(dcoll, gas_model, state, boundaries, *, time=0.0,
     # otherwise they stay on the interpolatory/base domain.
     if operator_states_quad is None:
         operator_states_quad = make_operator_fluid_states(
-            dcoll, state, gas_model, boundaries, quadrature_tag)
+            dcoll, state, gas_model, boundaries, quadrature_tag, comm_tag)
 
     vol_state_quad, inter_elem_bnd_states_quad, domain_bnd_states_quad = \
         operator_states_quad
@@ -404,7 +415,8 @@ def ns_operator(dcoll, gas_model, state, boundaries, *, time=0.0,
         # Get the interior trace pairs onto the surface quadrature
         # discretization (if any)
         interp_to_surf_quad(tpair=tpair)
-        for tpair in interior_trace_pairs(dcoll, grad_cv, tag=_NSGradCVTag)
+        for tpair in interior_trace_pairs(
+            dcoll, grad_cv, comm_tag=(_NSGradCVTag, comm_tag))
     ]
 
     # }}} Compute grad(CV)
@@ -423,7 +435,8 @@ def ns_operator(dcoll, gas_model, state, boundaries, *, time=0.0,
         # Get the interior trace pairs onto the surface quadrature
         # discretization (if any)
         interp_to_surf_quad(tpair=tpair)
-        for tpair in interior_trace_pairs(dcoll, grad_t, tag=_NSGradTemperatureTag)
+        for tpair in interior_trace_pairs(
+            dcoll, grad_t, comm_tag=(_NSGradTemperatureTag, comm_tag))
     ]
 
     # }}} compute grad(temperature)

mirgecom/wave.py

@@ -60,7 +60,7 @@ class _WaveTag:
     pass
 
 
-def wave_operator(dcoll, c, w):
+def wave_operator(dcoll, c, w, *, comm_tag=None):
     """Compute the RHS of the wave equation.
 
     Parameters
@@ -71,6 +71,8 @@ def wave_operator(dcoll, c, w):
         the (constant) wave speed
     w: numpy.ndarray
         an object array of DOF arrays, representing the state vector
+    comm_tag: Hashable
+        Tag for distributed communication
 
     Returns
     -------
@@ -98,7 +100,8 @@ def wave_operator(dcoll, c, w):
                                         exterior=dir_bc))
                 + sum(
                     _flux(dcoll, c=c, w_tpair=tpair)
-                    for tpair in interior_trace_pairs(dcoll, w, comm_tag=_WaveTag))
+                    for tpair in interior_trace_pairs(
+                        dcoll, w, comm_tag=(_WaveTag, comm_tag)))
                 )
             )
         )