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 all commits
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
42 changes: 26 additions & 16 deletions examples/advection/weak.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,11 @@
import pyopencl as cl
import pyopencl.tools as cl_tools
from arraycontext import flatten
from meshmode.mesh import BTAG_ALL
from meshmode.mesh import BTAG_ALL, TensorProductElementGroup

import grudge.dof_desc as dof_desc
import grudge.op as op
from grudge.array_context import PyOpenCLArrayContext
from grudge.array_context import NumpyArrayContext, PytatoPyOpenCLArrayContext


logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -96,24 +96,27 @@ def __call__(self, evt, basename, overwrite=True):
# }}}


def main(ctx_factory, dim=2, order=4, visualize=False):
def main(ctx_factory, dim=1, order=2, lazy=False,
visualize=False, group_cls=TensorProductElementGroup):
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
actx = PyOpenCLArrayContext(
queue,
allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)),
force_device_scalars=True,
)

if lazy is False:
actx = NumpyArrayContext()
else:
actx = PytatoPyOpenCLArrayContext(
queue,
allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)),)

# {{{ parameters

# domain [-d/2, d/2]^dim
d = 1.0
# number of points in each dimension
npoints = 20
npoints = 10

# final time
final_time = 1.0
final_time = 0.5

# velocity field
c = np.array([0.5] * dim)
Expand All @@ -129,7 +132,8 @@ def main(ctx_factory, dim=2, order=4, visualize=False):
from meshmode.mesh.generation import generate_box_mesh
mesh = generate_box_mesh(
[np.linspace(-d/2, d/2, npoints) for _ in range(dim)],
order=order)
order=order,
group_cls=group_cls)

from grudge.discretization import make_discretization_collection

Expand Down Expand Up @@ -163,7 +167,10 @@ def u_analytic(x, t=0):
def rhs(t, u):
return adv_operator.operator(t, u)

dt = actx.to_numpy(adv_operator.estimate_rk4_timestep(actx, dcoll, fields=u))
rhs_compiled = actx.compile(rhs)

# dt = actx.to_numpy(adv_operator.estimate_rk4_timestep(actx, dcoll, fields=u))
dt = 0.01

logger.info("Timestep size: %g", dt)

Expand All @@ -172,7 +179,7 @@ def rhs(t, u):
# {{{ time stepping

from grudge.shortcuts import set_up_rk4
dt_stepper = set_up_rk4("u", float(dt), u, rhs)
dt_stepper = set_up_rk4("u", float(dt), u, rhs_compiled)
plot = Plotter(actx, dcoll, order, visualize=visualize,
ylim=[-1.1, 1.1])

Expand Down Expand Up @@ -200,13 +207,16 @@ def rhs(t, u):
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--dim", default=2, type=int)
parser.add_argument("--order", default=4, type=int)
parser.add_argument("--dim", default=1, type=int)
parser.add_argument("--order", default=2, type=int)
parser.add_argument("--visualize", action="store_true")
parser.add_argument("--lazy", action="store_true")
parser.add_argument("--tp-elements", action="store_true")
args = parser.parse_args()

logging.basicConfig(level=logging.INFO)
main(cl.create_some_context,
dim=args.dim,
order=args.order,
visualize=args.visualize)
visualize=args.visualize,
lazy=args.lazy)
63 changes: 39 additions & 24 deletions examples/euler/acoustic_pulse.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,16 @@

import pyopencl as cl
import pyopencl.tools as cl_tools
from arraycontext import ArrayContext
from meshmode.mesh import BTAG_ALL
from arraycontext import ArrayContext, NumpyArrayContext
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 PyOpenCLArrayContext, PytatoPyOpenCLArrayContext
from grudge.array_context import PytatoPyOpenCLArrayContext
from grudge.models.euler import ConservedEulerField, EulerOperator, InviscidWallBC
from grudge.shortcuts import rk4_step

Expand Down Expand Up @@ -106,7 +110,8 @@ def run_acoustic_pulse(actx,
final_time=1,
resolution=16,
overintegration=False,
visualize=False):
visualize=False,
tensor_product_elements=False):

# eos-related parameters
gamma = 1.4
Expand All @@ -115,18 +120,19 @@ def run_acoustic_pulse(actx,

from meshmode.mesh.generation import generate_regular_rect_mesh

if tensor_product_elements:
group_cls = TensorProductElementGroup
else:
group_cls = SimplexElementGroup

dim = 2
box_ll = -0.5
box_ur = 0.5
mesh = generate_regular_rect_mesh(
a=(box_ll,)*dim,
b=(box_ur,)*dim,
nelements_per_axis=(resolution,)*dim)

from meshmode.discretization.poly_element import (
QuadratureSimplexGroupFactory,
default_simplex_group_factory,
)
nelements_per_axis=(resolution,)*dim,
group_cls=group_cls)

from grudge.discretization import make_discretization_collection
from grudge.dof_desc import DISCR_TAG_BASE, DISCR_TAG_QUAD
Expand All @@ -141,9 +147,8 @@ def run_acoustic_pulse(actx,
dcoll = make_discretization_collection(
actx, mesh,
discr_tag_to_group_factory={
DISCR_TAG_BASE: default_simplex_group_factory(
base_dim=mesh.dim, order=order),
DISCR_TAG_QUAD: QuadratureSimplexGroupFactory(2*order)
DISCR_TAG_BASE: InterpolatoryEdgeClusteredGroupFactory(order=order),
DISCR_TAG_QUAD: QuadratureGroupFactory(2*order)
}
)

Expand Down Expand Up @@ -182,12 +187,20 @@ def rhs(t, q):

# {{{ time stepping

import time

step = 0
t = 0.0
elapsed = 0.0
while t < final_time:
if step % 10 == 0:
norm_q = actx.to_numpy(op.norm(dcoll, fields, 2))
logger.info("[%04d] t = %.5f |q| = %.5e", step, t, norm_q)
if step != 0:
logger.info("[%04d] t = %.5f |q| = %.5e time per step = %.5f",
step, t, norm_q, elapsed / step)
else:
logger.info("[%04d] t = %.5f |q| = %.5e time per step = %.5f",
step, t, norm_q, 0)
if visualize:
vis.write_vtk_file(
f"{exp_name}-{step:04d}.vtu",
Expand All @@ -199,16 +212,19 @@ def rhs(t, q):
)
assert norm_q < 5

start = time.time()
fields = actx.thaw(actx.freeze(fields))
fields = rk4_step(fields, t, dt, compiled_rhs)
elapsed += time.time() - start
t += dt
step += 1

# }}}


def main(ctx_factory, order=3, final_time=1, resolution=16,
overintegration=False, visualize=False, lazy=False):
overintegration=False, visualize=False, lazy=False,
tensor_product_elements=False):
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)

Expand All @@ -218,29 +234,27 @@ def main(ctx_factory, order=3, final_time=1, resolution=16,
allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)),
)
else:
actx = PyOpenCLArrayContext(
queue,
allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)),
force_device_scalars=True,
)
actx = NumpyArrayContext()

run_acoustic_pulse(
actx,
order=order,
resolution=resolution,
overintegration=overintegration,
final_time=final_time,
visualize=visualize
visualize=visualize,
tensor_product_elements=tensor_product_elements
)


if __name__ == "__main__":
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--order", default=3, type=int)
parser.add_argument("--tpe", action="store_true")
parser.add_argument("--order", default=2, type=int)
parser.add_argument("--tfinal", default=0.1, type=float)
parser.add_argument("--resolution", default=16, type=int)
parser.add_argument("--resolution", default=4, type=int)
parser.add_argument("--oi", action="store_true",
help="use overintegration")
parser.add_argument("--visualize", action="store_true",
Expand All @@ -256,4 +270,5 @@ def main(ctx_factory, order=3, final_time=1, resolution=16,
resolution=args.resolution,
overintegration=args.oi,
visualize=args.visualize,
lazy=args.lazy)
lazy=args.lazy,
tensor_product_elements=args.tpe)
Loading