Merge remote-tracking branch 'origin/main' into multi-volume

grudge/array_context.py

@@ -32,12 +32,14 @@
 # {{{ imports
 
 from typing import (
-        TYPE_CHECKING, Mapping, Tuple, Any, Callable, Optional, Type)
+        TYPE_CHECKING, Mapping, Tuple, Any, Callable, Optional, Type,
+        FrozenSet)
 from dataclasses import dataclass
-
+from pytools.tag import Tag
 from meshmode.array_context import (
         PyOpenCLArrayContext as _PyOpenCLArrayContextBase,
         PytatoPyOpenCLArrayContext as _PytatoPyOpenCLArrayContextBase)
+from pyrsistent import pmap
 
 import logging
 logger = logging.getLogger(__name__)
@@ -160,21 +162,38 @@ def _dag_to_compiled_func(self, dict_of_named_arrays,
         # }}}
 
         part_id_to_prg = {}
+        name_in_program_to_tags = pmap()
+        name_in_program_to_axes = pmap()
 
         from pytato import DictOfNamedArrays
         for part in distributed_partition.parts.values():
             d = DictOfNamedArrays(
                         {var_name: distributed_partition.var_name_to_result[var_name]
                             for var_name in part.output_names
                          })
-            part_id_to_prg[part.pid], _, _ = self._dag_to_transformed_loopy_prg(d)
+            (
+                part_id_to_prg[part.pid],
+                part_prg_name_to_tags,
+                part_prg_name_to_axes
+            ) = self._dag_to_transformed_loopy_prg(d)
+
+            assert not (set(name_in_program_to_tags.keys())
+                        & set(part_prg_name_to_tags.keys()))
+            assert not (set(name_in_program_to_axes.keys())
+                        & set(part_prg_name_to_axes.keys()))
+            name_in_program_to_tags = name_in_program_to_tags.update(
+                part_prg_name_to_tags)
+            name_in_program_to_axes = name_in_program_to_axes.update(
+                part_prg_name_to_axes)
 
         return _DistributedCompiledFunction(
                 actx=self.actx,
                 distributed_partition=distributed_partition,
                 part_id_to_prg=part_id_to_prg,
                 input_id_to_name_in_program=input_id_to_name_in_program,
                 output_id_to_name_in_program=output_id_to_name_in_program,
+                name_in_program_to_tags=name_in_program_to_tags,
+                name_in_program_to_axes=name_in_program_to_axes,
                 output_template=output_template)
 
 
@@ -213,6 +232,8 @@ class _DistributedCompiledFunction:
     part_id_to_prg: "Mapping[PartId, pt.target.BoundProgram]"
     input_id_to_name_in_program: Mapping[Tuple[Any, ...], str]
     output_id_to_name_in_program: Mapping[Tuple[Any, ...], str]
+    name_in_program_to_tags: Mapping[str, FrozenSet[Tag]]
+    name_in_program_to_axes: Mapping[str, Tuple["pt.Axis", ...]]
     output_template: ArrayContainer
 
     def __call__(self, arg_id_to_arg) -> ArrayContainer:
@@ -223,6 +244,8 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
         """
 
         from arraycontext.impl.pytato.compile import _args_to_cl_buffers
+        from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array
+        from arraycontext.impl.pytato.utils import get_cl_axes_from_pt_axes
         input_args_for_prg = _args_to_cl_buffers(
                 self.actx, self.input_id_to_name_in_program, arg_id_to_arg)
 
@@ -234,7 +257,12 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
                 input_args=input_args_for_prg)
 
         def to_output_template(keys, _):
-            return self.actx.thaw(out_dict[self.output_id_to_name_in_program[keys]])
+            ary_name_in_prg = self.output_id_to_name_in_program[keys]
+            return self.actx.thaw(to_tagged_cl_array(
+                out_dict[ary_name_in_prg],
+                axes=get_cl_axes_from_pt_axes(
+                    self.name_in_program_to_axes[ary_name_in_prg]),
+                tags=self.name_in_program_to_tags[ary_name_in_prg]))
 
         from arraycontext.container.traversal import rec_keyed_map_array_container
         return rec_keyed_map_array_container(to_output_template,

grudge/dt_utils.py

@@ -46,7 +46,7 @@
 from typing import Optional, Sequence
 import numpy as np
 
-from arraycontext import ArrayContext, thaw, freeze, DeviceScalar
+from arraycontext import ArrayContext, thaw, freeze, Scalar
 from meshmode.transform_metadata import FirstAxisIsElementsTag
 
 from grudge.dof_desc import (
@@ -163,7 +163,7 @@ def dt_non_geometric_factors(
 @memoize_on_first_arg
 def h_max_from_volume(
         dcoll: DiscretizationCollection, dim=None,
-        dd: Optional[DOFDesc] = None) -> "DeviceScalar":
+        dd: Optional[DOFDesc] = None) -> Scalar:
     """Returns a (maximum) characteristic length based on the volume of the
     elements. This length may not be representative if the elements have very
     high aspect ratios.
@@ -195,7 +195,7 @@ def h_max_from_volume(
 @memoize_on_first_arg
 def h_min_from_volume(
         dcoll: DiscretizationCollection, dim=None,
-        dd: Optional[DOFDesc] = None) -> "DeviceScalar":
+        dd: Optional[DOFDesc] = None) -> Scalar:
     """Returns a (minimum) characteristic length based on the volume of the
     elements. This length may not be representative if the elements have very
     high aspect ratios.

grudge/reductions.py

@@ -63,7 +63,7 @@
     make_loopy_program,
     map_array_container,
     serialize_container,
-    DeviceScalar
+    Scalar
 )
 from arraycontext.container import ArrayOrContainerT
 
@@ -79,7 +79,7 @@
 
 # {{{ Nodal reductions
 
-def norm(dcoll: DiscretizationCollection, vec, p, dd=None) -> "DeviceScalar":
+def norm(dcoll: DiscretizationCollection, vec, p, dd=None) -> Scalar:
     r"""Return the vector p-norm of a function represented
     by its vector of degrees of freedom *vec*.
 
@@ -113,7 +113,7 @@ def norm(dcoll: DiscretizationCollection, vec, p, dd=None) -> "DeviceScalar":
         raise ValueError("unsupported norm order")
 
 
-def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     r"""Return the nodal sum of a vector of degrees of freedom *vec*.
 
     :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value
@@ -136,7 +136,7 @@ def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
         comm.allreduce(actx.to_numpy(nodal_sum_loc(dcoll, dd, vec)), op=MPI.SUM))
 
 
-def nodal_sum_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def nodal_sum_loc(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     r"""Return the rank-local nodal sum of a vector of degrees of freedom *vec*.
 
     :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value
@@ -156,7 +156,7 @@ def nodal_sum_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
     return sum([actx.np.sum(grp_ary) for grp_ary in vec])
 
 
-def nodal_min(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def nodal_min(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     r"""Return the nodal minimum of a vector of degrees of freedom *vec*.
 
     :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value
@@ -177,7 +177,7 @@ def nodal_min(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
         comm.allreduce(actx.to_numpy(nodal_min_loc(dcoll, dd, vec)), op=MPI.MIN))
 
 
-def nodal_min_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def nodal_min_loc(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     r"""Return the rank-local nodal minimum of a vector of degrees
     of freedom *vec*.
 
@@ -200,7 +200,7 @@ def nodal_min_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
             vec, actx.from_numpy(np.array(np.inf)))
 
 
-def nodal_max(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def nodal_max(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     r"""Return the nodal maximum of a vector of degrees of freedom *vec*.
 
     :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value
@@ -221,7 +221,7 @@ def nodal_max(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
         comm.allreduce(actx.to_numpy(nodal_max_loc(dcoll, dd, vec)), op=MPI.MAX))
 
 
-def nodal_max_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def nodal_max_loc(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     r"""Return the rank-local nodal maximum of a vector of degrees
     of freedom *vec*.
 
@@ -244,7 +244,7 @@ def nodal_max_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
             vec, actx.from_numpy(np.array(-np.inf)))
 
 
-def integral(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
+def integral(dcoll: DiscretizationCollection, dd, vec) -> Scalar:
     """Numerically integrates a function represented by a
     :class:`~meshmode.dof_array.DOFArray` of degrees of freedom.
 

test/test_grudge_sym_old.py

@@ -144,7 +144,7 @@ def _get_variables_on(dd):
 
     num_integral_1 = np.dot(ones_volm, actx.to_numpy(flatten(mass_op(f=f_quad))))
     err_1 = abs(num_integral_1 - true_integral)
-    assert err_1 < 2e-9, err_1
+    assert err_1 < 9e-9, err_1
 
     num_integral_2 = np.dot(f_volm, actx.to_numpy(flatten(mass_op(f=ones_quad))))
     err_2 = abs(num_integral_2 - true_integral)