Add a CommTag class with a stable hash

Co-authored-by: Matt Smith <[email protected]>

examples/wave/wave-op-mpi.py

@@ -42,7 +42,7 @@
 from meshmode.mesh import BTAG_ALL, BTAG_NONE # noqa
 
 from grudge.dof_desc import as_dofdesc, DOFDesc, DISCR_TAG_BASE, DISCR_TAG_QUAD
-from grudge.trace_pair import TracePair
+from grudge.trace_pair import TracePair, CommTag
 from grudge.discretization import DiscretizationCollection
 from grudge.shortcuts import make_visualizer, compiled_lsrk45_step
 
@@ -95,7 +95,7 @@ def wave_flux(actx, dcoll, c, w_tpair):
  return op.project(dcoll, dd, dd.with_dtag("all_faces"), c*flux_weak)
 
 
-class _WaveStateTag:
+class _WaveStateTag(CommTag):
  pass
 
 
@@ -144,7 +144,7 @@ def interp_to_surf_quad(utpair):
  ) + sum(
  wave_flux(actx, dcoll, c=c, w_tpair=interp_to_surf_quad(tpair))
  for tpair in op.interior_trace_pairs(dcoll, w,
- comm_tag=_WaveStateTag)
+ comm_tag=_WaveStateTag())
  )
  )
  )

grudge/trace_pair.py

@@ -22,6 +22,7 @@
 ---------------------------------------
 
 .. autofunction:: interior_trace_pairs
+.. autoclass:: CommTag
 .. autofunction:: local_interior_trace_pair
 .. autofunction:: cross_rank_trace_pairs
 """
@@ -70,7 +71,7 @@
 
 from numbers import Number
 
-from pytools import memoize_on_first_arg
+from pytools import memoize_on_first_arg, memoize_method
 
 from grudge.discretization import DiscretizationCollection
 from grudge.projection import project
@@ -318,8 +319,19 @@ def interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair:
  return local_interior_trace_pair(dcoll, vec)
 
 
+class CommTag:
+ """A communication tag with a hash value that is stable across
+ runs, even without setting ``PYTHONHASHSEED``."""
+ @memoize_method
+ def __hash__(self):
+ return hash(tuple(str(type(self)).encode("ascii")))
+
+ def __eq__(self, other):
+ return isinstance(other, type(self))
+
+
 def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *,
- comm_tag: Hashable = None, tag: Hashable = None,
+ comm_tag: CommTag = None, tag: Hashable = None,
  volume_dd: Optional[DOFDesc] = None) -> List[TracePair]:
  r"""Return a :class:`list` of :class:`TracePair` objects
  defined on the interior faces of *dcoll* and any faces connected to a
@@ -331,10 +343,10 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *,
 
  :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an
  :class:`~arraycontext.ArrayContainer` of them.
- :arg comm_tag: a hashable object used to match sent and received data
+ :arg comm_tag: a :class:`CommTag` used to match sent and received data
  across ranks. Communication will only match if both endpoints specify
  objects that compare equal. A generalization of MPI communication
- tags to arbitary, potentially composite objects.
+ tags to arbitrary, potentially composite objects.
  :returns: a :class:`list` of :class:`TracePair` objects.
  """
 
@@ -379,7 +391,7 @@ def connected_ranks(
  dcoll._volume_discrs[volume_dd.domain_tag.tag].mesh)
 
 
-def _sym_tag_to_num_tag(comm_tag: Optional[Hashable]) -> Optional[int]:
+def _sym_tag_to_num_tag(comm_tag: Optional[CommTag]) -> Optional[int]:
  if comm_tag is None:
  return comm_tag
 
@@ -498,10 +510,14 @@ class _RankBoundaryCommunicationLazy:
  def __init__(self,
  dcoll: DiscretizationCollection,
  array_container: ArrayOrContainer,
- remote_rank: int, comm_tag: Hashable,
+ remote_rank: int, comm_tag: Optional[CommTag],
  volume_dd=DD_VOLUME_ALL):
  if comm_tag is None:
- raise ValueError("lazy communication requires 'tag' to be supplied")
+ raise ValueError("lazy communication requires 'comm_tag' to be supplied")
+
+ if not isinstance(comm_tag, CommTag):
+ from warnings import warn
+ warn(f"comm_tag {comm_tag} should be an instance of CommTag")
 
  bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank))
 
@@ -544,7 +560,7 @@ def finish(self):
 def cross_rank_trace_pairs(
  dcoll: DiscretizationCollection, ary: ArrayOrContainer,
  tag: Hashable = None,
- *, comm_tag: Hashable = None,
+ *, comm_tag: CommTag = None,
  volume_dd: Optional[DOFDesc] = None) -> List[TracePair]:
  r"""Get a :class:`list` of *ary* trace pairs for each partition boundary.
 
@@ -570,7 +586,7 @@ def cross_rank_trace_pairs(
  :arg comm_tag: a hashable object used to match sent and received data
  across ranks. Communication will only match if both endpoints specify
  objects that compare equal. A generalization of MPI communication
- tags to arbitary, potentially composite objects.
+ tags to arbitrary, potentially composite objects.
 
  :returns: a :class:`list` of :class:`TracePair` objects.
  """