inducer · nchristensen · Oct 25, 2020 · Nov 10, 2020 · Nov 10, 2020 · Nov 10, 2020
diff --git a/meshmode/discretization/__init__.py b/meshmode/discretization/__init__.py
@@ -33,7 +33,8 @@
 
 import loopy as lp
 from meshmode.transform_metadata import (
-        ConcurrentElementInameTag, ConcurrentDOFInameTag, FirstAxisIsElementsTag)
+        ConcurrentElementInameTag, ConcurrentDOFInameTag, FirstAxisIsElementsTag,
+        IsDOFArray, ParameterValue, IsOpArray, KernelDataTag)
 
 from warnings import warn
 
@@ -553,6 +554,10 @@ def prg():
             t_unit = make_loopy_program(
                 "{[iel,idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}",
                 "result[iel,idof] = weights[idof]",
+                kernel_data=[
+                    lp.GlobalArg("result", None, shape=lp.auto, tags=[IsDOFArray()]),
+                    ...
+                ],
                 name="quad_weights")
             return lp.tag_inames(t_unit, {
                 "iel": ConcurrentElementInameTag(),
@@ -601,10 +606,30 @@ def resample_mesh_nodes(grp, iaxis):
                     and np.linalg.norm(grp_unit_nodes - meg_unit_nodes) < tol):
                 return nodes
 
+            mat = actx.from_numpy(grp.from_mesh_interp_matrix())
+            fp_format = nodes.dtype
+            ne, nj = nodes.shape
+            ni, nj = mat.shape
+
+            kernel_data = [
+                lp.GlobalArg("arg1", fp_format, shape=(ne, nj),
+                    offset=lp.auto),  # In default data layout apparently
+                lp.GlobalArg("arg0", fp_format, shape=(ni, nj), offset=lp.auto,
+                    tags=[IsOpArray()]),
+                lp.GlobalArg("out",  fp_format, shape=(ne, ni), offset=lp.auto,
+                    tags=[IsDOFArray()], is_output=True),
+                lp.ValueArg("Ni", tags=[ParameterValue(ni)]),
+                lp.ValueArg("Nj", tags=[ParameterValue(nj)]),
+                lp.ValueArg("Ne", tags=[ParameterValue(ne)]),
+                ...
+            ]
+
+            kd_tag = KernelDataTag(kernel_data)
+
             return actx.einsum("ij,ej->ei",
-                               actx.from_numpy(grp.from_mesh_interp_matrix()),
+                               mat,
                                nodes,
-                               tagged=(FirstAxisIsElementsTag(),))
+                               tagged=(FirstAxisIsElementsTag(), kd_tag,))
 
         result = make_obj_array([
             _DOFArray(None, tuple([
@@ -660,12 +685,35 @@ def get_mat(grp, gref_axes):
 
         return actx.from_numpy(mat)
 
-    return _DOFArray(actx, tuple(
-            actx.einsum("ij,ej->ei",
-                        get_mat(grp, ref_axes),
+    data = []
+    for grp in discr.groups:
+
+        mat = get_mat(grp, ref_axes)
+        fp_format = vec[grp.index].dtype
+        ne, nj = vec[grp.index].shape
+        ni, nj = mat.shape
+
+        kernel_data = [
+            lp.GlobalArg("arg1", fp_format, shape=(ne, nj), offset=lp.auto,
+                tags=[IsDOFArray()]),
+            lp.GlobalArg("arg0", fp_format, shape=(ni, nj), offset=lp.auto,
+                tags=[IsOpArray()]),
+            lp.GlobalArg("out",  fp_format, shape=(ne, ni), offset=lp.auto,
+                tags=[IsDOFArray()], is_output=True),
+            lp.ValueArg("Ni", tags=[ParameterValue(ni)]),
+            lp.ValueArg("Nj", tags=[ParameterValue(nj)]),
+            lp.ValueArg("Ne", tags=[ParameterValue(ne)]),
+            ...
+        ]
+
+        kd_tag = KernelDataTag(kernel_data)
+
+        data.append(actx.einsum("ij,ej->ei",
+                        mat,
                         vec[grp.index],
-                        tagged=(FirstAxisIsElementsTag(),))
-            for grp in discr.groups))
+                        tagged=(FirstAxisIsElementsTag(), kd_tag,)))
+
+    return _DOFArray(actx, tuple(data))
 
 # }}}
 

diff --git a/meshmode/discretization/connection/direct.py b/meshmode/discretization/connection/direct.py
@@ -26,7 +26,8 @@
 
 import loopy as lp
 from meshmode.transform_metadata import (
-        ConcurrentElementInameTag, ConcurrentDOFInameTag)
+        ConcurrentElementInameTag, ConcurrentDOFInameTag,
+        IsOpArray, IsDOFArray, ParameterValue)
 from pytools import memoize_in, keyed_memoize_method
 from arraycontext import (
         ArrayContext, NotAnArrayContainerError,
@@ -504,7 +505,8 @@ def _apply_with_inplace_updates(self, ary):
 
         @memoize_in(actx, (DirectDiscretizationConnection,
             "resample_by_mat_knl_inplace"))
-        def mat_knl():
+        def mat_knl(nelements_vec, nelements_result, n_to_nodes, n_from_nodes,
+                result_dtype, rmat_dtype, ary_dtype, index_dtype):
             t_unit = make_loopy_program(
                 """{[iel, idof, j]:
                     0<=iel<nelements and
@@ -514,14 +516,23 @@ def mat_knl():
                     = sum(j, resample_mat[idof, j] \
                     * ary[from_element_indices[iel], j])",
                 [
-                    lp.GlobalArg("result", None,
+                    lp.GlobalArg("result", result_dtype,
                         shape="nelements_result, n_to_nodes",
-                        offset=lp.auto),
-                    lp.GlobalArg("ary", None,
+                        offset=lp.auto, tags=[IsDOFArray()]),
+                    lp.GlobalArg("resample_mat", rmat_dtype,
+                        shape="n_to_nodes, n_from_nodes",
+                        offset=lp.auto, tags=[IsOpArray()]),
+                    lp.GlobalArg("ary", ary_dtype,
                         shape="nelements_vec, n_from_nodes",
-                        offset=lp.auto),
-                    lp.ValueArg("nelements_result", np.int32),
-                    lp.ValueArg("nelements_vec", np.int32),
+                        offset=lp.auto, tags=[IsDOFArray()]),
+                    lp.ValueArg("n_to_nodes", np.int32,
+                        tags=[ParameterValue(n_to_nodes)]),
+                    lp.ValueArg("n_from_nodes", np.int32,
+                        tags=[ParameterValue(n_from_nodes)]),
+                    lp.ValueArg("nelements_result", np.int32,
+                        tags=[ParameterValue(nelements_result)]),
+                    lp.ValueArg("nelements_vec", np.int32,
+                        tags=[ParameterValue(nelements_vec)]),
                     "...",
                     ],
                 name="resample_by_mat_inplace")
@@ -532,7 +543,7 @@ def mat_knl():
 
         @memoize_in(actx,
                 (DirectDiscretizationConnection, "resample_by_picking_knl_inplace"))
-        def pick_knl():
+        def pick_knl(n_to_nodes):
             t_unit = make_loopy_program(
                 """{[iel, idof]:
                     0<=iel<nelements and
@@ -542,13 +553,15 @@ def pick_knl():
                 [
                     lp.GlobalArg("result", None,
                         shape="nelements_result, n_to_nodes",
-                        offset=lp.auto),
+                        offset=lp.auto, tags=[IsDOFArray()]),
                     lp.GlobalArg("ary", None,
                         shape="nelements_vec, n_from_nodes",
-                        offset=lp.auto),
+                        offset=lp.auto, tags=[IsDOFArray()]),
                     lp.ValueArg("nelements_result", np.int32),
                     lp.ValueArg("nelements_vec", np.int32),
                     lp.ValueArg("n_from_nodes", np.int32),
+                    lp.ValueArg("n_to_nodes", np.int32,
+                        tags=[ParameterValue(n_to_nodes)]),
                     "...",
                     ],
                 name="resample_by_picking_inplace")
@@ -571,16 +584,27 @@ def pick_knl():
                         actx, i_tgrp, i_batch)
 
                 if point_pick_indices is None:
-                    actx.call_loopy(mat_knl(),
-                            resample_mat=self._resample_matrix(
-                                actx, i_tgrp, i_batch),
+                    resample_mat = self._resample_matrix(actx, i_tgrp, i_batch)
+                    n_to_nodes, n_from_nodes = resample_mat.shape
+                    nelements_result, _ = result[i_tgrp].shape
+                    nelements_vec, _ = ary[batch.from_group_index].shape
+                    index_dtype = batch.from_element_indices.dtype
+                    result_dtype = result[i_tgrp].dtype
+                    rmat_dtype = resample_mat.dtype
+                    ary_dtype = ary[batch.from_group_index].dtype
+
+                    actx.call_loopy(mat_knl(nelements_vec, nelements_result,
+                                n_to_nodes, n_from_nodes,
+                                result_dtype, rmat_dtype, ary_dtype, index_dtype),
+                            resample_mat=resample_mat,
                             result=result[i_tgrp],
                             ary=ary[batch.from_group_index],
                             from_element_indices=batch.from_element_indices,
                             to_element_indices=batch.to_element_indices)
 
                 else:
-                    actx.call_loopy(pick_knl(),
+                    nelements, n_to_nodes = result[i_tgrp].shape
+                    actx.call_loopy(pick_knl(n_to_nodes),
                             pick_list=point_pick_indices,
                             result=result[i_tgrp],
                             ary=ary[batch.from_group_index],
@@ -639,7 +663,7 @@ def knl():
             [
                 lp.GlobalArg("result", None,
                     shape="nnodes_tgt, nnodes_src",
-                    offset=lp.auto),
+                    offset=lp.auto, tags=[IsDOFArray()]),
                 lp.ValueArg("itgt_base, isrc_base", np.int32),
                 lp.ValueArg("nnodes_tgt, nnodes_src", np.int32),
                 "...",

diff --git a/meshmode/discretization/connection/modal.py b/meshmode/discretization/connection/modal.py
@@ -27,9 +27,10 @@
 import numpy.linalg as la
 import modepy as mp
 
-from arraycontext import (
-        NotAnArrayContainerError, serialize_container, deserialize_container)
-from meshmode.transform_metadata import FirstAxisIsElementsTag
+from arraycontext import (NotAnArrayContainerError,
+    serialize_container, deserialize_container,)
+from meshmode.transform_metadata import (FirstAxisIsElementsTag,
+    KernelDataTag, IsDOFArray, IsOpArray, ParameterValue)
 from meshmode.discretization import InterpolatoryElementGroupBase
 from meshmode.discretization.poly_element import QuadratureSimplexElementGroup
 from meshmode.discretization.connection.direct import DiscretizationConnection
@@ -162,10 +163,30 @@ def vandermonde_inverse(grp):
             vdm_inv = la.inv(vdm)
             return actx.from_numpy(vdm_inv)
 
+        fp_format = ary[grp.index].dtype
+        vi_mat = vandermonde_inverse(grp)
+        ne, nj = ary[grp.index].shape
+        ni, nj = vi_mat.shape
+
+        import loopy as lp
+        kernel_data = [
+            lp.GlobalArg("arg1", fp_format, shape=(ne, nj), offset=lp.auto,
+                tags=[IsDOFArray()]),
+            lp.GlobalArg("arg0", fp_format, shape=(ni, nj), offset=lp.auto,
+                tags=[IsOpArray()]),
+            lp.GlobalArg("out",  fp_format, shape=(ne, ni), offset=lp.auto,
+                tags=[IsDOFArray()], is_output=True),
+            lp.ValueArg("Ni", tags=[ParameterValue(ni)]),
+            lp.ValueArg("Nj", tags=[ParameterValue(nj)]),
+            lp.ValueArg("Ne", tags=[ParameterValue(ne)]),
+            ...
+        ]
+        kd_tag = KernelDataTag(kernel_data)
+
         return actx.einsum("ij,ej->ei",
-                           vandermonde_inverse(grp),
+                           vi_mat,
                            ary[grp.index],
-                           tagged=(FirstAxisIsElementsTag(),))
+                           tagged=(FirstAxisIsElementsTag(), kd_tag,))
 
     def __call__(self, ary):
         """Computes modal coefficients data from a functions

diff --git a/meshmode/discretization/connection/projection.py b/meshmode/discretization/connection/projection.py
@@ -29,7 +29,7 @@
 from arraycontext import (
         NotAnArrayContainerError,
         make_loopy_program, serialize_container, deserialize_container)
-from meshmode.transform_metadata import FirstAxisIsElementsTag
+from meshmode.transform_metadata import FirstAxisIsElementsTag, IsDOFArray
 from meshmode.discretization.connection.direct import (
         DiscretizationConnection,
         DirectDiscretizationConnection)
@@ -172,7 +172,7 @@ def kproj():
                                  shape=("n_from_elements", "n_from_nodes")),
                     lp.GlobalArg("result", None,
                                  shape=("n_to_elements", "n_to_nodes"),
-                                 is_input=False),
+                                 is_input=False, tags=[IsDOFArray()]),
                     lp.GlobalArg("basis_tabulation", None,
                                  shape=("n_to_nodes", "n_to_nodes")),
                     lp.GlobalArg("weights", None,

diff --git a/meshmode/dof_array.py b/meshmode/dof_array.py
@@ -35,7 +35,8 @@
 from pytools import single_valued, memoize_in
 
 from meshmode.transform_metadata import (
-            ConcurrentElementInameTag, ConcurrentDOFInameTag)
+            ConcurrentElementInameTag, ConcurrentDOFInameTag,
+            IsDOFArray)
 from arraycontext import (
         ArrayContext, NotAnArrayContainerError,
         make_loopy_program, with_container_arithmetic,
@@ -61,9 +62,9 @@
 .. autofunction:: check_dofarray_against_discr
 """
 
-
 # {{{ DOFArray
 
+
 @with_container_arithmetic(
         bcast_obj_array=True,
         bcast_numpy_array=True,
@@ -431,9 +432,9 @@ def prg():
             """,
             [
                 lp.GlobalArg("result", None,
-                             shape="nelements * ndofs_per_element"),
+                    shape="nelements * ndofs_per_element"),
                 lp.GlobalArg("grp_ary", None,
-                             shape=("nelements", "ndofs_per_element")),
+                    shape=("nelements", "ndofs_per_element"), tags=[IsDOFArray()]),
                 lp.ValueArg("nelements", np.int32),
                 lp.ValueArg("ndofs_per_element", np.int32),
                 "..."
@@ -508,6 +509,10 @@ def prg():
         t_unit = make_loopy_program(
             "{[iel,idof]: 0<=iel<nelements and 0<=idof<ndofs_per_element}",
             "result[iel, idof] = ary[grp_start + iel*ndofs_per_element + idof]",
+            kernel_data=[
+                lp.GlobalArg("result", None, shape=lp.auto, tags=[IsDOFArray()]),
+                ...
+            ],
             name="unflatten")
         return lp.tag_inames(t_unit, {
             "iel": ConcurrentElementInameTag(),

diff --git a/meshmode/mesh/generation.py b/meshmode/mesh/generation.py
@@ -1226,6 +1226,7 @@ def generate_regular_rect_mesh(a=(0, 0), b=(1, 1), *, nelements_per_axis=None,
                                npoints_per_axis=None,
                                periodic=None,
                                order=1,
+                               coord_dtype=np.float32,
                                boundary_tag_to_face=None,
                                group_cls=None,
                                mesh_type=None,

diff --git a/meshmode/transform_metadata.py b/meshmode/transform_metadata.py
@@ -28,7 +28,7 @@
 THE SOFTWARE.
 """
 
-from pytools.tag import Tag
+from pytools.tag import Tag, UniqueTag
 
 
 class FirstAxisIsElementsTag(Tag):
@@ -57,3 +57,39 @@ class ConcurrentDOFInameTag(Tag):
     computations for all DOFs within each element may be performed
     concurrently.
     """
+
+
+class ParameterValue(UniqueTag):
+    """A tag that applies to :class:`loopy.ValueArg`. Instances of this tag
+    are initialized with the value of the parameter and this value may be
+    later retrieved
+    for use with `loopy.fix_parameter` in `transform_loopy_program`. This allows
+    the fixing of parameters whose values cannot be set during the creation of the
+    program/kernel. This is useful for fixing the loop parameters of einsum kernels
+    for instance.
+    """
+
+    def __init__(self, value):
+        self.value = value
+
+
+class IsDOFArray(Tag):
+    """A tag that is applicable to :class:`loopy.ArrayArg` indicating the content of the
+    array comprises element DOFs.
+    """
+    pass
+
+
+class IsOpArray(Tag):
+    """A tag that is applicable to arrays indicating the array is an
+    operator (as opposed, for instance, to a DOF array)."""
+    pass
+
+
+class KernelDataTag(Tag):
+    """A tag that applies to :class:`loopy.LoopKernel`. Kernel data provided
+    with this tag can be later applied to the kernel. This is used, for
+    instance, to specify kernel data in einsum kernels."""
+
+    def __init__(self, kernel_data):
+        self.kernel_data = kernel_data
diff --git a/setup.py b/setup.py
@@ -44,12 +44,13 @@ def main():
               "numpy",
               "modepy>=2020.2",
               "gmsh_interop",
-              "pytools>=2020.4.1",
+              "pytools>=2020.4.5",
               "pytest>=2.3",
 
               # 2019.1 is required for the Firedrake CIs, which use an very specific
               # version of Loopy.
-              "loopy>=2019.1",
+              #"loopy>=2019.1",
+              "loopy>=2020.2.2",
 
               "arraycontext",