Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add fast operator evaluation for tensor-product discretizations #362

Draft
wants to merge 65 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
65 commits
Select commit Hold shift + click to select a range
8bd3969
start updating grad
a-alveyblanc Jul 25, 2024
0afe732
add TODOs
a-alveyblanc Jul 25, 2024
9c4c150
add strong form tensor product gradient
a-alveyblanc Jul 26, 2024
c264cf9
Merge branch 'main' of https://github.com/inducer/grudge into tensor-…
a-alveyblanc Aug 31, 2024
e105a64
start working on weak-overint case for TP elements
a-alveyblanc Sep 5, 2024
a3f7f7e
start brainstorming better ways to handle weak overintegration
a-alveyblanc Sep 6, 2024
830c2e0
small changes
a-alveyblanc Sep 6, 2024
102f6e0
do not support overintegration (yet)
a-alveyblanc Sep 6, 2024
597fcc9
fix gradient tests
a-alveyblanc Sep 6, 2024
790ab13
fix tensor product gradient
a-alveyblanc Sep 6, 2024
91e3121
do not compute stiffness matrix; apply mass to all axes and diff op t…
a-alveyblanc Sep 6, 2024
9ff324b
arg name change
a-alveyblanc Sep 6, 2024
7b50e32
add fast operator eval for weak and strong divergence
a-alveyblanc Sep 6, 2024
230c1bd
add tensor product inverse mass
a-alveyblanc Sep 6, 2024
30f984e
tag mass operator as such
a-alveyblanc Sep 6, 2024
92b54a2
add fast operator eval for mass operator application
a-alveyblanc Sep 7, 2024
34b1cd6
toward overintegration support
a-alveyblanc Sep 7, 2024
b8e312e
minor changes to adjust for changes in pytato
a-alveyblanc Sep 8, 2024
a79a603
adjust operators to accept 1D tensor product discretizations
a-alveyblanc Sep 8, 2024
51884a7
Merge branch 'main' of https://github.com/inducer/grudge into tensor-…
a-alveyblanc Sep 9, 2024
c8d0a92
checkpoint before adding wadg + overintegration + fast operator evalu…
a-alveyblanc Sep 12, 2024
ad0d2e4
tag axes of grad result
a-alveyblanc Sep 12, 2024
ca22480
add WADG + overintegration for simplices
a-alveyblanc Sep 12, 2024
092a6cc
add WADG + overintegration for simplices
a-alveyblanc Sep 12, 2024
bea91ad
start fixing up fast operator eval + overintegration + wadg
a-alveyblanc Sep 13, 2024
8182571
small changes
a-alveyblanc Sep 16, 2024
4cacc6b
overintegration + fast operator evaluation
a-alveyblanc Sep 19, 2024
73ba29a
first round of clean-ups
a-alveyblanc Sep 19, 2024
0a431d4
start updating face mass for TP
a-alveyblanc Sep 23, 2024
3dfaeec
undo face mass changes; add rough version of bilinear form evaluator
a-alveyblanc Oct 4, 2024
851a1f9
rename bilinear forms file; add rough draft of SeparableBilinearForm
a-alveyblanc Oct 4, 2024
91c87df
minor change: fix typing on generic dispatching function
a-alveyblanc Oct 4, 2024
351ae62
changes after review; bilinear forms now only internal
a-alveyblanc Oct 8, 2024
d174c40
get things working with quadrature; improve test_mass_operator_inverse
a-alveyblanc Oct 13, 2024
21a5855
remove unused variable
a-alveyblanc Oct 13, 2024
4e8b5f8
remove large refined mesh files
a-alveyblanc Oct 13, 2024
16f7f36
add default quadrature for computing bilinear forms
a-alveyblanc Oct 22, 2024
4d8d468
get all tests passing with quadrature rules + numpy array context
a-alveyblanc Oct 27, 2024
5b53482
fix merge conflicts
a-alveyblanc Oct 27, 2024
0a91985
some ruff fixes
a-alveyblanc Oct 28, 2024
53abda5
fix failing MPI wave op test
a-alveyblanc Oct 29, 2024
943f9ca
add redundant mass/inverse mass mappers
a-alveyblanc Oct 31, 2024
fa5b24e
remove 2x refined gh-339 mesh
a-alveyblanc Nov 1, 2024
5c74b6e
toward transformations
a-alveyblanc Nov 9, 2024
8898863
resolve merge conflicts
a-alveyblanc Nov 9, 2024
19f7a49
some dag rewriter work; restrict fast operator eval to non-overintegr…
a-alveyblanc Nov 13, 2024
2dfd7ff
rewrite operators to be more predictable in the DAG
a-alveyblanc Nov 13, 2024
9359246
remove tagging for now
a-alveyblanc Nov 13, 2024
da6dff7
bypass WADG if base and quad discretizations are the same
a-alveyblanc Nov 14, 2024
3f88f04
initial algebraic dag xforms for tp
a-alveyblanc Nov 16, 2024
5bd67ed
algebraic transforms v0.1
a-alveyblanc Nov 17, 2024
5a95360
basic (more like primitive) parallelization scheme; add more dag tran…
a-alveyblanc Nov 17, 2024
8c5be5a
ruff fixes
a-alveyblanc Nov 17, 2024
aad468a
changes to make ruff and pylint happy
a-alveyblanc Nov 17, 2024
7f8fb0e
ruff fix
a-alveyblanc Nov 17, 2024
654dd36
remove ghost nodes left over after TP DAG xforms
a-alveyblanc Nov 18, 2024
7819616
update some docs; fix dim from num faces computation
a-alveyblanc Nov 19, 2024
6a84f74
op refactor updates + transform updates
a-alveyblanc Dec 16, 2024
f6fd35d
move matrices to their own file and update operators accordingly
a-alveyblanc Dec 20, 2024
5b23379
all tests passing except mpi tests
a-alveyblanc Dec 24, 2024
3b69323
all tests passing; disable xforms for now
a-alveyblanc Dec 24, 2024
65cfe02
attempt to resolve missing axis tags; minor transform changes
a-alveyblanc Dec 30, 2024
d46246a
track down final culprit for untagged axes
a-alveyblanc Dec 31, 2024
a70cdb3
add name hints and propert tags to matrices
a-alveyblanc Jan 3, 2025
567637c
re-implementation of tp transforms
a-alveyblanc Jan 5, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
some ruff fixes
  • Loading branch information
a-alveyblanc committed Oct 28, 2024
commit 0a9198582fb0969b9420255366b9d41737f99f64
18 changes: 5 additions & 13 deletions examples/euler/acoustic_pulse.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,29 +25,21 @@

import logging

from meshmode.discretization.poly_element import (
InterpolatoryEdgeClusteredGroupFactory,
QuadratureGroupFactory
)
import numpy as np

import pyopencl as cl
import pyopencl.tools as cl_tools
from arraycontext import ArrayContext, NumpyArrayContext
from meshmode.mesh import (
BTAG_ALL,
SimplexElementGroup,
TensorProductElementGroup
from meshmode.discretization.poly_element import (
InterpolatoryEdgeClusteredGroupFactory,
QuadratureGroupFactory,
)
from meshmode.mesh import BTAG_ALL, SimplexElementGroup, TensorProductElementGroup
from pytools.obj_array import make_obj_array

import grudge.op as op
from grudge.array_context import PytatoPyOpenCLArrayContext
from grudge.models.euler import (
ConservedEulerField,
EulerOperator,
InviscidWallBC
)
from grudge.models.euler import ConservedEulerField, EulerOperator, InviscidWallBC
from grudge.shortcuts import rk4_step


Expand Down
4 changes: 2 additions & 2 deletions grudge/array_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,15 +46,15 @@
)
from warnings import warn

from grudge.transform.metadata import OutputIsTensorProductDOFArrayOrdered

from meshmode.array_context import (
PyOpenCLArrayContext as _PyOpenCLArrayContextBase,
PytatoPyOpenCLArrayContext as _PytatoPyOpenCLArrayContextBase,
)
from pytools import to_identifier
from pytools.tag import Tag

from grudge.transform.metadata import OutputIsTensorProductDOFArrayOrdered


logger = logging.getLogger(__name__)

Expand Down
48 changes: 19 additions & 29 deletions grudge/bilinear_forms.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations


__copyright__ = """
Copyright (C) 2024 Addison Alvey-Blanco
Copyright (C) 2024 University of Illinois Board of Trustees
Expand All @@ -26,56 +27,45 @@
"""


from arraycontext import Array, ArrayContext, ArrayOrContainer, tag_axes

from collections.abc import Callable

from dataclasses import dataclass
from typing import Optional

from grudge.discretization import DiscretizationCollection
from grudge.dof_desc import (
DISCR_TAG_BASE,
DD_VOLUME_ALL,
DOFDesc
)
from grudge.geometry import (
area_element,
inverse_surface_metric_derivative_mat
)
from grudge.tools import rec_map_subarrays
from grudge.transform.metadata import (
ReferenceTensorProductMassOperatorTag,
TensorProductDOFAxisTag,
TensorProductOperatorAxisTag
)
import numpy as np
import numpy.linalg as la

import modepy as mp
from arraycontext import Array, ArrayContext, ArrayOrContainer, tag_axes
from meshmode.discretization import (
ElementGroupBase,
InterpolatoryElementGroupBase,
NodalElementGroupBase
NodalElementGroupBase,
)
from meshmode.discretization.poly_element import (
PolynomialSimplexElementGroupBase,
SimplexElementGroupBase,
TensorProductElementGroupBase
TensorProductElementGroupBase,
)
from meshmode.dof_array import DOFArray
from meshmode.transform_metadata import (
DiscretizationDOFAxisTag,
DiscretizationElementAxisTag
DiscretizationElementAxisTag,
)

import modepy as mp
from modepy.spaces import TensorProductSpace
from modepy.tools import (
reshape_array_for_tensor_product_space as fold,
unreshape_array_for_tensor_product_space as unfold
unreshape_array_for_tensor_product_space as unfold,
)

import numpy as np
import numpy.linalg as la

from typing import Optional
from grudge.discretization import DiscretizationCollection
from grudge.dof_desc import DD_VOLUME_ALL, DISCR_TAG_BASE, DOFDesc
from grudge.geometry import area_element, inverse_surface_metric_derivative_mat
from grudge.tools import rec_map_subarrays
from grudge.transform.metadata import (
ReferenceTensorProductMassOperatorTag,
TensorProductDOFAxisTag,
TensorProductOperatorAxisTag,
)


# {{{ BilinearForm base class
Expand Down
64 changes: 28 additions & 36 deletions grudge/op.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@

from __future__ import annotations


__copyright__ = """
Copyright (C) 2021 Andreas Kloeckner
Copyright (C) 2021 University of Illinois Board of Trustees
Expand Down Expand Up @@ -75,45 +76,39 @@
import numpy as np

import modepy as mp
from modepy import inverse_mass_matrix
from modepy.tools import (
reshape_array_for_tensor_product_space as fold,
unreshape_array_for_tensor_product_space as unfold
)

from arraycontext import (
ArrayContext,
ArrayOrContainer,
map_array_container,
tag_axes
)
from meshmode.dof_array import DOFArray
from arraycontext import ArrayContext, ArrayOrContainer, map_array_container, tag_axes
from meshmode.discretization import (
Discretization,
InterpolatoryElementGroupBase,
NodalElementGroupBase
NodalElementGroupBase,
)
from meshmode.discretization.poly_element import (
SimplexElementGroupBase,
TensorProductElementGroupBase,
SimplexElementGroupBase
)
from meshmode.dof_array import DOFArray
from meshmode.transform_metadata import (
DiscretizationAmbientDimAxisTag,
DiscretizationDOFAxisTag,
DiscretizationElementAxisTag,
DiscretizationFaceAxisTag,
FirstAxisIsElementsTag,
)
from modepy import inverse_mass_matrix
from modepy.tools import (
reshape_array_for_tensor_product_space as fold,
unreshape_array_for_tensor_product_space as unfold,
)
from pytools import keyed_memoize_in
from pytools.obj_array import make_obj_array

import grudge.dof_desc as dof_desc
from grudge.bilinear_forms import (
_NonTensorProductBilinearForm,
_TensorProductBilinearForm,
apply_bilinear_form,
single_axis_contraction,
_TensorProductBilinearForm,
_NonTensorProductBilinearForm
)
import grudge.dof_desc as dof_desc
from grudge.discretization import DiscretizationCollection
from grudge.dof_desc import (
DD_VOLUME_ALL,
Expand All @@ -124,7 +119,7 @@
VolumeDomainTag,
as_dofdesc,
)
from grudge.geometry import inverse_surface_metric_derivative_mat, area_element
from grudge.geometry import area_element, inverse_surface_metric_derivative_mat
from grudge.interpolation import interp
from grudge.projection import project
from grudge.reductions import (
Expand Down Expand Up @@ -154,9 +149,9 @@
)
from grudge.transform.metadata import (
OutputIsTensorProductDOFArrayOrdered,
ReferenceTensorProductMassOperatorTag as MassMatrix1D,
ReferenceTensorProductMassInverseOperatorTag as InverseMassMatrix1D,
TensorProductOperatorAxisTag
ReferenceTensorProductMassOperatorTag as MassMatrix1D,
TensorProductOperatorAxisTag,
)


Expand Down Expand Up @@ -373,8 +368,7 @@ def compute_tensor_product_divergence(actx, out_grp, in_grp, vec, ijm):

div += ref_deriv * ijm[func_axis, rst_axis]

return tag_axes(actx, { 0: DiscretizationElementAxisTag() }, div)

return tag_axes(actx, {0: DiscretizationElementAxisTag()}, div)

per_group_divs = []
for out_grp, in_grp, vec_i, ijm_i in zip(
Expand Down Expand Up @@ -427,7 +421,7 @@ def get_ref_derivative_mats(
return actx.freeze(
tag_axes(
actx,
{ i: TensorProductOperatorAxisTag() for i in range(2) },
{i: TensorProductOperatorAxisTag() for i in range(2)},
actx.from_numpy(diff_mat)))

return actx.freeze(
Expand Down Expand Up @@ -663,10 +657,8 @@ def _weak_scalar_grad(dcoll, dd_in, vec, *args,

def _weak_scalar_div(dcoll, dd_in, vecs, *args,
use_tensor_product_fast_eval=True):
from arraycontext import (
get_container_context_recursively,
serialize_container
)
from arraycontext import get_container_context_recursively, serialize_container

from grudge.geometry import inverse_surface_metric_derivative_mat

assert isinstance(vecs, np.ndarray)
Expand Down Expand Up @@ -862,7 +854,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args,
# {{{ Mass operator

def mass(dcoll: DiscretizationCollection, *args,
use_tensor_product_fast_eval = True) -> ArrayOrContainer:
use_tensor_product_fast_eval=True) -> ArrayOrContainer:
r"""Return the action of the DG mass matrix on a vector (or vectors)
of :class:`~meshmode.dof_array.DOFArray`\ s, *vec*. In the case of
*vec* being an :class:`~arraycontext.ArrayContainer`,
Expand Down Expand Up @@ -923,7 +915,7 @@ def get_ref_inv_mass_mat(grp, use_tensor_product_fast_eval):
InverseMassMatrix1D(),
tag_axes(
actx,
{ i : TensorProductOperatorAxisTag() for i in range(2) },
{i: TensorProductOperatorAxisTag() for i in range(2)},
actx.from_numpy(
inverse_mass_matrix(basis_1d, nodes_1d)))))

Expand Down Expand Up @@ -1120,8 +1112,7 @@ def get_ref_face_mass_mat(face_grp, vol_grp):

import modepy as mp
from meshmode.discretization import ElementGroupWithBasis
from meshmode.discretization.poly_element import \
QuadratureSimplexElementGroup
from meshmode.discretization.poly_element import QuadratureSimplexElementGroup

n = vol_grp.order
m = face_grp.order
Expand Down Expand Up @@ -1307,8 +1298,8 @@ def get_ref_mass_mat(out_grp, in_grp):
MassMatrix1D(),
tag_axes(
actx,
{ i : TensorProductOperatorAxisTag()
for i in range(2) },
{i: TensorProductOperatorAxisTag()
for i in range(2)},
actx.from_numpy(mp.mass_matrix(basis_1d, nodes_1d))))

else:
Expand All @@ -1332,8 +1323,8 @@ def get_ref_mass_mat(out_grp, in_grp):
MassMatrix1D(),
tag_axes(
actx,
{ i : TensorProductOperatorAxisTag()
for i in range(2) },
{i: TensorProductOperatorAxisTag()
for i in range(2)},
actx.from_numpy(mass_matrix))))

elif isinstance(out_grp, SimplexElementGroupBase):
Expand All @@ -1354,6 +1345,7 @@ def get_ref_mass_mat(out_grp, in_grp):

return get_ref_mass_mat(out_element_group, in_element_group)


def _apply_mass_operator(
dcoll: DiscretizationCollection, dd_out, dd_in, vec):
if not isinstance(vec, DOFArray):
Expand Down
3 changes: 2 additions & 1 deletion grudge/transform/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,10 @@
THE SOFTWARE.
"""

from pytools.tag import Tag, tag_dataclass
from pytato.transform.metadata import AxisIgnoredForPropagationTag

from meshmode.transform_metadata import DiscretizationDOFAxisTag
from pytools.tag import Tag, tag_dataclass


class OutputIsTensorProductDOFArrayOrdered(Tag):
Expand Down
8 changes: 6 additions & 2 deletions test/test_mpi_communication.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@
from pytools.obj_array import flat_obj_array

from grudge import dof_desc, op
from grudge.array_context import MPIPyOpenCLArrayContext, MPIPytatoArrayContext
from grudge.array_context import (
MPIPyOpenCLArrayContext, MPIPytatoArrayContext, MPINumpyArrayContext
)
from grudge.discretization import make_discretization_collection
from grudge.shortcuts import rk4_step

Expand All @@ -49,7 +51,7 @@ class SimpleTag:

# {{{ mpi test infrastructure

DISTRIBUTED_ACTXS = [MPIPyOpenCLArrayContext, MPIPytatoArrayContext]
DISTRIBUTED_ACTXS = [MPINumpyArrayContext, MPIPytatoArrayContext]


def run_test_with_mpi(num_ranks, f, *args):
Expand Down Expand Up @@ -87,6 +89,8 @@ def run_test_with_mpi_inner():
actx = actx_class(comm, queue, mpi_base_tag=15000)
elif actx_class is MPIPyOpenCLArrayContext:
actx = actx_class(comm, queue, force_device_scalars=True)
elif actx_class is MPINumpyArrayContext:
actx = actx_class(comm)
else:
raise ValueError("unknown actx_class")

Expand Down
6 changes: 1 addition & 5 deletions test/test_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,7 @@
InterpolatoryEdgeClusteredGroupFactory,
QuadratureGroupFactory,
)
from meshmode.mesh import (
SimplexElementGroup,
TensorProductElementGroup,
BTAG_ALL
)
from meshmode.mesh import BTAG_ALL, SimplexElementGroup, TensorProductElementGroup
from pytools.obj_array import make_obj_array

from grudge import geometry, op
Expand Down