diff --git a/grudge/discretization.py b/grudge/discretization.py index 64c00e94f7b3668b8a8f1ae96ac59a1b8f8daf66..a8bee7e11a046aaeaf810f3f47c0e01ae6df995d 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -25,9 +25,9 @@ THE SOFTWARE. import six from pytools import memoize_method -import pyopencl as cl from grudge import sym -import numpy as np +import numpy as np # noqa: F401 +from meshmode.array_context import ArrayContext # FIXME Naming not ideal @@ -40,7 +40,6 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): .. automethod :: discr_from_dd .. automethod :: connection_from_dds - .. autoattribute :: cl_context .. autoattribute :: dim .. autoattribute :: ambient_dim .. autoattribute :: mesh @@ -49,7 +48,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): .. automethod :: zeros """ - def __init__(self, cl_ctx, mesh, order, quad_tag_to_group_factory=None, + def __init__(self, array_context, mesh, order, quad_tag_to_group_factory=None, mpi_communicator=None): """ :param quad_tag_to_group_factory: A mapping from quadrature tags (typically @@ -60,6 +59,8 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): be carried out with the standard volume discretization. """ + self._setup_actx = array_context + if quad_tag_to_group_factory is None: quad_tag_to_group_factory = {} @@ -68,7 +69,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): from meshmode.discretization import Discretization - self._volume_discr = Discretization(cl_ctx, mesh, + self._volume_discr = Discretization(array_context, mesh, self.group_factory_for_quadrature_tag(sym.QTAG_NONE)) # {{{ management of discretization-scoped common subexpressions @@ -81,9 +82,9 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): # }}} - with cl.CommandQueue(cl_ctx) as queue: - self._dist_boundary_connections = \ - self._set_up_distributed_communication(mpi_communicator, queue) + self._dist_boundary_connections = \ + self._set_up_distributed_communication( + mpi_communicator, array_context) self.mpi_communicator = mpi_communicator @@ -97,7 +98,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): return self.mpi_communicator.Get_rank() \ == self._get_management_rank_index() - def _set_up_distributed_communication(self, mpi_communicator, queue): + def _set_up_distributed_communication(self, mpi_communicator, array_context): from_dd = sym.DOFDesc("vol", sym.QTAG_NONE) from meshmode.distributed import get_connected_partitions @@ -118,7 +119,8 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): from_dd, sym.DOFDesc(sym.BTAG_PARTITION(i_remote_part), sym.QTAG_NONE)) setup_helper = setup_helpers[i_remote_part] = MPIBoundaryCommSetupHelper( - mpi_communicator, queue, conn, i_remote_part, grp_factory) + mpi_communicator, array_context, conn, + i_remote_part, grp_factory) setup_helper.post_sends() for i_remote_part, setup_helper in six.iteritems(setup_helpers): @@ -152,7 +154,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): from meshmode.discretization import Discretization return Discretization( - self._volume_discr.cl_context, + self._setup_actx, no_quad_discr.mesh, self.group_factory_for_quadrature_tag(qtag)) @@ -186,6 +188,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): make_face_to_all_faces_embedding return make_face_to_all_faces_embedding( + self._setup_actx, faces_conn, self.discr_from_dd(to_dd), self.discr_from_dd(from_dd)) @@ -224,6 +227,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): make_same_mesh_connection return make_same_mesh_connection( + self._setup_actx, self.discr_from_dd(to_dd), self.discr_from_dd(from_dd)) @@ -276,7 +280,7 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): def _quad_volume_discr(self, quadrature_tag): from meshmode.discretization import Discretization - return Discretization(self._volume_discr.cl_context, self._volume_discr.mesh, + return Discretization(self._setup_actx, self._volume_discr.mesh, self.group_factory_for_quadrature_tag(quadrature_tag)) # {{{ boundary @@ -285,9 +289,10 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): def _boundary_connection(self, boundary_tag): from meshmode.discretization.connection import make_face_restriction return make_face_restriction( - self._volume_discr, - self.group_factory_for_quadrature_tag(sym.QTAG_NONE), - boundary_tag=boundary_tag) + self._setup_actx, + self._volume_discr, + self.group_factory_for_quadrature_tag(sym.QTAG_NONE), + boundary_tag=boundary_tag) # }}} @@ -298,21 +303,24 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): from meshmode.discretization.connection import ( make_face_restriction, FACE_RESTR_INTERIOR) return make_face_restriction( - self._volume_discr, - self.group_factory_for_quadrature_tag(sym.QTAG_NONE), - FACE_RESTR_INTERIOR, + self._setup_actx, + self._volume_discr, + self.group_factory_for_quadrature_tag(sym.QTAG_NONE), + FACE_RESTR_INTERIOR, - # FIXME: This will need to change as soon as we support - # pyramids or other elements with non-identical face - # types. - per_face_groups=False) + # FIXME: This will need to change as soon as we support + # pyramids or other elements with non-identical face + # types. + per_face_groups=False) @memoize_method def opposite_face_connection(self): from meshmode.discretization.connection import \ make_opposite_face_connection - return make_opposite_face_connection(self._interior_faces_connection()) + return make_opposite_face_connection( + self._setup_actx, + self._interior_faces_connection()) # }}} @@ -323,21 +331,18 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): from meshmode.discretization.connection import ( make_face_restriction, FACE_RESTR_ALL) return make_face_restriction( - self._volume_discr, - self.group_factory_for_quadrature_tag(sym.QTAG_NONE), - FACE_RESTR_ALL, + self._setup_actx, + self._volume_discr, + self.group_factory_for_quadrature_tag(sym.QTAG_NONE), + FACE_RESTR_ALL, - # FIXME: This will need to change as soon as we support - # pyramids or other elements with non-identical face - # types. - per_face_groups=False) + # FIXME: This will need to change as soon as we support + # pyramids or other elements with non-identical face + # types. + per_face_groups=False) # }}} - @property - def cl_context(self): - return self._volume_discr.cl_context - @property def dim(self): return self._volume_discr.dim @@ -358,13 +363,11 @@ class DGDiscretizationWithBoundaries(DiscretizationBase): def mesh(self): return self._volume_discr.mesh - def empty(self, queue=None, dtype=None, extra_dims=None, allocator=None): - return self._volume_discr.empty(queue, dtype, extra_dims=extra_dims, - allocator=allocator) + def empty(self, array_context: ArrayContext, dtype=None): + return self._volume_discr.empty(array_context, dtype) - def zeros(self, queue, dtype=None, extra_dims=None, allocator=None): - return self._volume_discr.zeros(queue, dtype, extra_dims=extra_dims, - allocator=allocator) + def zeros(self, array_context: ArrayContext, dtype=None): + return self._volume_discr.zeros(array_context, dtype) def is_volume_where(self, where): from grudge import sym diff --git a/grudge/eager.py b/grudge/eager.py index e64e2fa8344a848f4d0a218cfd0286ec9f896e2a..8b05edd0cebeb4541c520ae3a31ea69b292c3ea1 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -26,31 +26,23 @@ THE SOFTWARE. import numpy as np # noqa from grudge.discretization import DGDiscretizationWithBoundaries from pytools import memoize_method -from pytools.obj_array import ( - with_object_array_or_scalar, - is_obj_array) -import pyopencl as cl +from pytools.obj_array import obj_array_vectorize import pyopencl.array as cla # noqa from grudge import sym, bind -from meshmode.mesh import BTAG_ALL, BTAG_NONE # noqa - - -def with_queue(queue, ary): - return with_object_array_or_scalar( - lambda x: x.with_queue(queue), ary) - -def without_queue(ary): - return with_queue(None, ary) +from meshmode.mesh import BTAG_ALL, BTAG_NONE # noqa +from meshmode.dof_array import freeze, DOFArray class EagerDGDiscretization(DGDiscretizationWithBoundaries): def interp(self, src, tgt, vec): - if is_obj_array(vec): - return with_object_array_or_scalar( + if (isinstance(vec, np.ndarray) + and vec.dtype.char == "O" + and not isinstance(vec, DOFArray)): + return obj_array_vectorize( lambda el: self.interp(src, tgt, el), vec) - return self.connection_from_dds(src, tgt)(vec.queue, vec) + return self.connection_from_dds(src, tgt)(vec) def nodes(self): return self._volume_discr.nodes() @@ -60,7 +52,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): return bind(self, sym.nabla(self.dim) * sym.Variable("u")) def grad(self, vec): - return self._bound_grad()(vec.queue, u=vec) + return self._bound_grad()(u=vec) def div(self, vecs): return sum( @@ -71,7 +63,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): return bind(self, sym.stiffness_t(self.dim) * sym.Variable("u")) def weak_grad(self, vec): - return self._bound_weak_grad()(vec.queue, u=vec) + return self._bound_weak_grad()(u=vec) def weak_div(self, vecs): return sum( @@ -79,22 +71,25 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): @memoize_method def normal(self, dd): - with cl.CommandQueue(self.cl_context) as queue: - surface_discr = self.discr_from_dd(dd) - return without_queue( - bind(self, sym.normal( - dd, surface_discr.ambient_dim, surface_discr.dim))(queue)) + surface_discr = self.discr_from_dd(dd) + actx = surface_discr._setup_actx + return freeze( + bind(self, + sym.normal(dd, surface_discr.ambient_dim, surface_discr.dim)) + (array_context=actx)) @memoize_method def _bound_inverse_mass(self): return bind(self, sym.InverseMassOperator()(sym.Variable("u"))) def inverse_mass(self, vec): - if is_obj_array(vec): - return with_object_array_or_scalar( + if (isinstance(vec, np.ndarray) + and vec.dtype.char == "O" + and not isinstance(vec, DOFArray)): + return obj_array_vectorize( lambda el: self.inverse_mass(el), vec) - return self._bound_inverse_mass()(vec.queue, u=vec) + return self._bound_inverse_mass()(u=vec) @memoize_method def _bound_face_mass(self): @@ -102,10 +97,34 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): return bind(self, sym.FaceMassOperator()(u)) def face_mass(self, vec): - if is_obj_array(vec): - return with_object_array_or_scalar( + if (isinstance(vec, np.ndarray) + and vec.dtype.char == "O" + and not isinstance(vec, DOFArray)): + return obj_array_vectorize( lambda el: self.face_mass(el), vec) - return self._bound_face_mass()(vec.queue, u=vec) + return self._bound_face_mass()(u=vec) + + @memoize_method + def _norm(self, p): + return bind(self, sym.norm(p, sym.var("arg"))) + + def norm(self, vec, p=2): + return self._norm(p)(arg=vec) + + +def interior_trace_pair(discr, vec): + i = discr.interp("vol", "int_faces", vec) + + if (isinstance(vec, np.ndarray) + and vec.dtype.char == "O" + and not isinstance(vec, DOFArray)): + e = obj_array_vectorize( + lambda el: discr.opposite_face_connection()(el), + i) + + from grudge.symbolic.primitives import TracePair + return TracePair("int_faces", i, e) + # vim: foldmethod=marker diff --git a/grudge/execution.py b/grudge/execution.py index a95b15e75e43917688f2c3766651e3018ff9d526..0104584ef28a8707f893b1c9d0ffd576c3453ade 100644 --- a/grudge/execution.py +++ b/grudge/execution.py @@ -22,12 +22,19 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -import six +from typing import Optional, Union, Dict +from numbers import Number import numpy as np + +from pytools import memoize_in +from pytools.obj_array import make_obj_array + import loopy as lp import pyopencl as cl import pyopencl.array # noqa -from pytools import memoize_in + +from meshmode.dof_array import DOFArray, thaw, flatten, unflatten +from meshmode.array_context import ArrayContext, make_loopy_program import grudge.symbolic.mappers as mappers from grudge import sym @@ -42,17 +49,20 @@ from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_1 # noqa: F401 MPI_TAG_SEND_TAGS = 1729 +ResultType = Union[DOFArray, Number] + + # {{{ exec mapper class ExecutionMapper(mappers.Evaluator, mappers.BoundOpMapperMixin, mappers.LocalOpReducerMixin): - def __init__(self, queue, context, bound_op): + def __init__(self, array_context, context, bound_op): super(ExecutionMapper, self).__init__(context) self.discrwb = bound_op.discrwb self.bound_op = bound_op self.function_registry = bound_op.function_registry - self.queue = queue + self.array_context = array_context # {{{ expression mappings @@ -62,13 +72,14 @@ class ExecutionMapper(mappers.Evaluator, discr = self.discrwb.discr_from_dd(expr.dd) - result = discr.empty(self.queue, allocator=self.bound_op.allocator) - result.fill(1) + result = discr.empty(self.array_context) + for grp_ary in result: + grp_ary.fill(1) return result def map_node_coordinate_component(self, expr): discr = self.discrwb.discr_from_dd(expr.dd) - return discr.nodes()[expr.axis].with_queue(self.queue) + return thaw(self.array_context, discr.nodes()[expr.axis]) def map_grudge_variable(self, expr): from numbers import Number @@ -76,8 +87,9 @@ class ExecutionMapper(mappers.Evaluator, value = self.context[expr.name] if not expr.dd.is_scalar() and isinstance(value, Number): discr = self.discrwb.discr_from_dd(expr.dd) - ary = discr.empty(self.queue) - ary.fill(value) + ary = discr.empty(self.array_context) + for grp_ary in ary: + grp_ary.fill(value) value = ary return value @@ -91,42 +103,41 @@ class ExecutionMapper(mappers.Evaluator, from numbers import Number if not dd.is_scalar() and isinstance(value, Number): discr = self.discrwb.discr_from_dd(dd) - ary = discr.empty(self.queue) - ary.fill(value) + ary = discr.empty(self.array_context) + for grp_ary in ary: + grp_ary.fill(value) value = ary return value def map_call(self, expr): args = [self.rec(p) for p in expr.parameters] - return self.function_registry[expr.function.name](self.queue, *args) + return self.function_registry[expr.function.name](self.array_context, *args) # }}} # {{{ elementwise reductions def _map_elementwise_reduction(self, op_name, field_expr, dd): - @memoize_in(self, "elementwise_%s_knl" % op_name) - def knl(): - knl = lp.make_kernel( - "{[el, idof, jdof]: 0<=el