diff --git a/grudge/op.py b/grudge/op.py index f38b466a09a2bf548359d8a5b674628dc92ca6fc..5e914d8ec0e394812f579fa1fd0c0ef0b610b7e7 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -50,15 +50,12 @@ THE SOFTWARE. """ -from arraycontext import ( - ArrayContext, - make_loopy_program -) +from arraycontext import ArrayContext from meshmode.transform_metadata import FirstAxisIsElementsTag from grudge.discretization import DiscretizationCollection -from pytools import memoize_in, keyed_memoize_in +from pytools import keyed_memoize_in from pytools.obj_array import obj_array_vectorize, make_obj_array from meshmode.dof_array import DOFArray @@ -763,52 +760,28 @@ def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec): dtype = vec.entry_dtype actx = vec.array_context - @memoize_in(actx, (_apply_face_mass_operator, "face_mass_knl")) - def prg(): - t_unit = make_loopy_program( - [ - "{[iel]: 0 <= iel < nelements}", - "{[f]: 0 <= f < nfaces}", - "{[idof]: 0 <= idof < nvol_nodes}", - "{[jdof]: 0 <= jdof < nface_nodes}" - ], - """ - result[iel, idof] = sum(f, sum(jdof, mat[idof, f, jdof] - * jac_surf[f, iel, jdof] - * vec[f, iel, jdof])) - """, - name="face_mass" - ) - import loopy as lp - from meshmode.transform_metadata import ( - ConcurrentElementInameTag, ConcurrentDOFInameTag) - return lp.tag_inames(t_unit, { - "iel": ConcurrentElementInameTag(), - "idof": ConcurrentDOFInameTag()}) - assert len(face_discr.groups) == len(volm_discr.groups) surf_area_elements = area_element(actx, dcoll, dd=dd) return DOFArray( actx, data=tuple( - actx.call_loopy(prg(), - mat=reference_face_mass_matrix( + actx.einsum("ifj,fej,fej->ei", + reference_face_mass_matrix( actx, face_element_group=afgrp, vol_element_group=vgrp, - dtype=dtype - ), - jac_surf=surf_ae_i.reshape( + dtype=dtype), + surf_ae_i.reshape( vgrp.mesh_el_group.nfaces, vgrp.nelements, - afgrp.nunit_dofs - ), - vec=vec_i.reshape( + -1), + vec_i.reshape( vgrp.mesh_el_group.nfaces, vgrp.nelements, - afgrp.nunit_dofs - ))["result"] + afgrp.nunit_dofs), + arg_names=("ref_face_mass_mat", "jac_surf", "vec"), + tagged=(FirstAxisIsElementsTag(),)) for vgrp, afgrp, vec_i, surf_ae_i in zip(volm_discr.groups, face_discr.groups,