Updates to TP operators

examples/tensor-product-examples/acoustic_pulse.py

@@ -216,14 +216,12 @@ def main(ctx_factory, order=3, final_time=1, resolution=16,
     queue = cl.CommandQueue(cl_ctx)
 
     if lazy:
-        from grudge.array_context import PytatoTensorProductArrayContext
-        actx = PytatoTensorProductArrayContext(
+        actx = PytatoPyOpenCLArrayContext(
             queue,
             allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)),
         )
     else:
-        from grudge.array_context import TensorProductArrayContext
-        actx = TensorProductArrayContext(
+        actx = PyOpenCLArrayContext(
             queue,
             allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)),
             force_device_scalars=True,

grudge/array_context.py

@@ -109,8 +109,7 @@
 
 class PyOpenCLArrayContext(_PyOpenCLArrayContextBase):
     """Inherits from :class:`meshmode.array_context.PyOpenCLArrayContext`. Extends it
-    to understand :mod:`grudge`-specific transform metadata. (Of which there isn't
-    any, for now.)
+    to understand :mod:`grudge`-specific transform metadata.
     """
     def __init__(self, queue: "pyopencl.CommandQueue",
             allocator: Optional["pyopencl.tools.AllocatorBase"] = None,
@@ -125,15 +124,38 @@ def __init__(self, queue: "pyopencl.CommandQueue",
         super().__init__(queue, allocator,
                          wait_event_queue_length, force_device_scalars)
 
+    def transform_loopy_program(self, t_unit):
+        knl = t_unit.default_entrypoint
+
+        # {{{ process tensor product specific metadata
+
+        if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered):
+            new_args = []
+            for arg in knl.args:
+                if arg.is_output:
+                    arg = arg.copy(dim_tags=(
+                        f"N{len(arg.shape)-1},"
+                        + ",".join(f"N{i}"
+                                   for i in range(len(arg.shape)-1))
+                        ))
+
+                new_args.append(arg)
+
+            knl = knl.copy(args=new_args)
+            t_unit = t_unit.with_kernel(knl)
+
+        # }}}
+
+        return super().transform_loopy_program(t_unit)
+
 # }}}
 
 
 # {{{ pytato
 
 class PytatoPyOpenCLArrayContext(_PytatoPyOpenCLArrayContextBase):
     """Inherits from :class:`meshmode.array_context.PytatoPyOpenCLArrayContext`.
-    Extends it to understand :mod:`grudge`-specific transform metadata. (Of
-    which there isn't any, for now.)
+    Extends it to understand :mod:`grudge`-specific transform metadata.
     """
     def __init__(self, queue, allocator=None,
             *,
@@ -154,6 +176,29 @@ def __init__(self, queue, allocator=None,
         super().__init__(queue, allocator,
                 compile_trace_callback=compile_trace_callback)
 
+    def transform_loopy_program(self, t_unit):
+        knl = t_unit.default_entrypoint
+
+        # {{{ process tensor product specific metadata
+
+        if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered):
+            new_args = []
+            for arg in knl.args:
+                if arg.is_output:
+                    arg = arg.copy(dim_tags=(
+                        f"N{len(arg.shape)-1},"
+                        + ",".join(f"N{i}"
+                                   for i in range(len(arg.shape)-1))
+                        ))
+
+                new_args.append(arg)
+
+            knl = knl.copy(args=new_args)
+
+        # }}}
+
+        return super().transform_loopy_program(t_unit)
+
 # }}}
 
 
@@ -605,8 +650,8 @@ def get_reasonable_array_context_class(
 
 # }}}
 
-
 # {{{ distributed + numpy
+
 try:
     from arraycontext import NumpyArrayContext
 
@@ -626,123 +671,32 @@ def clone(self):
 except ImportError:
     print("Failed to import numpy array context.")
     pass
-# }}}
-
-
-# {{{ Tensor product array contexts
-
-# {{{ Relevant tags
-
-class OutputIsTensorProductDOFArrayOrdered(Tag):
-    """Signify that the strides will not be of order "C" or "F". See
-    :class:`grudge.array_context.TensorProductArrayContext` for more details.
-    """
-    pass
 
 # }}}
 
 
-# {{{ Eager TP array contexts
+# {{{ tensor product-specific machinery
 
-class TensorProductArrayContext(_PyOpenCLArrayContextBase):
-    """Specialized array context for use with tensor product elements.
+class OutputIsTensorProductDOFArrayOrdered(Tag):
+    """Signify that the strides will not be of order "C" or "F".
 
     The strides for the arrays containing tensor product element data are of the
     form (slow, fastest, faster, fast). These strides are not "C" or "F" order.
     Hence, this specialized array context takes care of specifying the
     particular strides required.
     """
-
-    def transform_loopy_program(self, t_unit):
-        knl = t_unit.default_entrypoint
-        if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered):
-            new_args = []
-            for arg in knl.args:
-                if arg.is_output:
-                    arg = arg.copy(dim_tags=(
-                        f"N{len(arg.shape)-1},"
-                        + ",".join(f"N{i}"
-                                   for i in range(len(arg.shape)-1))
-                        ))
-
-                new_args.append(arg)
-
-            knl = knl.copy(args=new_args)
-            t_unit = t_unit.with_kernel(knl)
-
-        return super().transform_loopy_program(t_unit)
-
-
-class TensorProductMPIPyOpenCLArrayContext(MPIPyOpenCLArrayContext,
-                                         TensorProductArrayContext):
     pass
 
-# }}}
-
-
-# {{{ Lazy tensor product array contexts
-
-class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext):
-    def transform_dag(self, dag):
-        return super().transform_dag(dag)
-
-    def transform_loopy_program(self, t_unit):
-        knl = t_unit.default_entrypoint
-
-        # {{{ adjust strides according to tensor product structure
-        if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered):
-            new_args = []
-            for arg in knl.args:
-                if arg.is_output:
-                    arg = arg.copy(dim_tags=(
-                        f"N{len(arg.shape)-1},"
-                        + ",".join(f"N{i}"
-                                   for i in range(len(arg.shape)-1))
-                        ))
-
-                new_args.append(arg)
-
-            knl = knl.copy(args=new_args)
-        # }}}
-
-        return super().transform_loopy_program(t_unit)
-
-# }}}
-
-
-# {{{ TP fusion actx
-
-from meshmode.array_context import FusionContractorArrayContext
-
-
-class TensorProductFusionContractorArrayContext(FusionContractorArrayContext):
 
-    def transform_loopy_program(self, t_unit):
-        knl = t_unit.default_entrypoint
-        if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered):
-            new_args = []
-            for arg in knl.args:
-                if arg.is_output:
-                    arg = arg.copy(dim_tags=(
-                        f"N{len(arg.shape)-1},"
-                        + ",".join(f"N{i}"
-                                   for i in range(len(arg.shape)-1))
-                        ))
-
-                new_args.append(arg)
-
-            knl = knl.copy(args=new_args)
-            t_unit = t_unit.with_kernel(knl)
-
-        return super().transform_loopy_program(t_unit)
-
-
-class TensorProductMPIFusionContractorArrayContext(
-        MPIPytatoArrayContext, TensorProductFusionContractorArrayContext):
+class MassMatrix1d(Tag):
+    """Used in DAG transformation to realize algebraic simplification of 1D
+    inverse mass operator times mass operator.
+    """
     pass
 
-# }}}
-
+class InverseMassMatrix1d(Tag):
+    """See MassMatrix1d.
+    """
 
 # }}}