add tensor-product metadata to meshmode from grudge

meshmode/array_context.py

@@ -47,7 +47,7 @@
     _PytestPytatoPyOpenCLArrayContextFactory,
     register_pytest_array_context_factory,
 )
-from grudge.transform.metadata import (
+from meshmode.transform_metadata import (
     FaceMassOperatorTag,
     MassInverseOperatorTag,
     TensorProductDOFAxisTag,
@@ -1193,7 +1193,7 @@ class EinsumTag(UniqueTag):
 
 
 def _prepare_kernel_for_parallelization(kernel):
-    from grudge.transform.metadata import TensorProductDOFAxisTag
+    from meshmode.transform_metadata import TensorProductDOFAxisTag
     discr_tag_to_prefix = {DiscretizationElementAxisTag:        "iel",
                            DiscretizationDOFAxisTag:            "idof",
                            TensorProductDOFAxisTag:             "idof_tp",

meshmode/transform_metadata.py

@@ -36,6 +36,7 @@
 """
 
 from pytools.tag import Tag, UniqueTag, tag_dataclass
+from pytato.transform.metadata import AxisIgnoredForPropagationTag
 
 
 class FirstAxisIsElementsTag(Tag):
@@ -131,3 +132,103 @@ class DiscretizationDOFPickListAxisTag(DiscretizationEntityAxisTag):
     DOF pick lists. See :mod:`meshmode.discretization.connection.direct` for
     details.
     """
+
+# {{{ tensor-product and operator metadata
+
+class OperatorTag(Tag):
+    """
+    Used to signify that an array is an operator.
+    """
+
+
+class FaceMassOperatorTag(OperatorTag):
+    """
+    Used to signify than an array is a face mass operator.
+    """
+
+
+class MassOperatorTag(OperatorTag):
+    """
+    Used to signify that an array is a mass operator.
+    """
+
+
+class MassInverseOperatorTag(OperatorTag):
+    """
+    Used to signify that an array is an inverse mass operator.
+    """
+
+
+class DifferentiationOperatorTag(OperatorTag):
+    """
+    Used to signify that an array is a *strong* differentiation operator.
+    """
+
+
+class StiffnessOperatorTag(OperatorTag):
+    """
+    Used to signify that an array is a *weak* differentiation operator.
+    """
+
+
+@tag_dataclass
+class TensorProductDOFAxisTag(DiscretizationEntityAxisTag):
+    """
+    Signify an axis as containing the DOFs of a tensor product discretization.
+    `iaxis` is later interpreted to determine the relative update speed (i.e.
+    the stride) of each axis.
+    """
+    iaxis: int
+
+
+class TensorProductOperatorAxisTag(
+        DiscretizationEntityAxisTag,
+        AxisIgnoredForPropagationTag
+    ):
+    """
+    Signify an axis is part of a 1D operator applied to a tensor product
+    discretization. No tags will be propagated to or along axes containing this
+    tag.
+    """
+    pass
+
+
+class TensorProductOperatorTag(Tag):
+    """
+    Used to tag an operator as one that acts on DOFs from a tensor-product
+    discretization. Used to make decisions about how to handle prefetching and
+    precomputing these operators.
+    """
+    pass
+
+
+class TensorProductMassOperatorTag(TensorProductOperatorTag):
+    """
+    Tag an operator as being a reference mass operator. Used to realize an
+    algebraic simplification of redundant mass-times-mass-inverse operations
+    when using a tensor product discretization.
+    """
+    pass
+
+
+class TensorProductMassInverseOperatorTag(TensorProductOperatorTag):
+    """
+    See `TensorProductMassOperatorTag`.
+    """
+    pass
+
+
+class TensorProductDifferentiationOperatorTag(OperatorTag):
+    """
+    See `DifferentiationOperatorTag`.
+    """
+
+
+class TensorProductStiffnessOperatorTag(TensorProductOperatorTag):
+    """
+    Similar to `TensorProductMassOperatorTag`. Used to implement an
+    associativity DAG transformation.
+    """
+    pass
+
+# }}}