pytato_array_context changes (#216)

* adds meshmode.Pytato_Array_Context

Co-authored-by: Matthias Diener <mdiener@illinois.edu>

* simple-dg: take --lazy to use PytatoPyOpenCLArrayContext

Co-authored-by: Matthias Diener <mdiener@illinois.edu>

* avoid naming collision b/w variable and function name

Co-authored-by: Matthias Diener <mdiener@illinois.edu>

* test with arraycontext.git@pytato

* add lazy simple-dg test

* correct path

* add pytato to requirements.txt

* fix test

* PytatoArrayContext is available on main

* [not working] Cache codegen result in freeze()

* Revert "[not working] Cache codegen result in freeze()"

This reverts commit dd97621.

* Draft: two-actx simple-dg

* fix the 2-actx

* remove unused code

Co-authored-by: Kaushik Kulkarni <kaushikcfd@gmail.com>
Co-authored-by: Andreas Kloeckner <inform@tiker.net>

.github/workflows/ci.yml

@@ -93,6 +93,11 @@ jobs:
                 curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-py-project-and-run-examples.sh
                 . ./build-py-project-and-run-examples.sh
 
+
+                # Test lazy eval
+                rm -f *.vtu
+                python simple-dg.py --lazy
+
     docs:
         name: Documentation
         runs-on: ubuntu-latest

examples/simple-dg.py

@@ -32,7 +32,8 @@
 
 from meshmode.mesh import BTAG_ALL, BTAG_NONE  # noqa
 from meshmode.dof_array import DOFArray, flat_norm
-from meshmode.array_context import PyOpenCLArrayContext
+from meshmode.array_context import (PyOpenCLArrayContext,
+                                    PytatoPyOpenCLArrayContext)
 from arraycontext import (
         freeze, thaw,
         make_loopy_program,
@@ -471,12 +472,16 @@ def array_context(self):
 
 
 @log_process(logger)
-def main():
+def main(lazy=False):
     logging.basicConfig(level=logging.INFO)
 
     cl_ctx = cl.create_some_context()
     queue = cl.CommandQueue(cl_ctx)
-    actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
+    actx_outer = PyOpenCLArrayContext(queue, force_device_scalars=True)
+    if lazy:
+        actx_rhs = PytatoPyOpenCLArrayContext(queue)
+    else:
+        actx_rhs = actx_outer
 
     nel_1d = 16
     from meshmode.mesh.generation import generate_regular_rect_mesh
@@ -492,31 +497,37 @@ def main():
 
     logger.info("%d elements", mesh.nelements)
 
-    discr = DGDiscretization(actx, mesh, order=order)
+    discr = DGDiscretization(actx_outer, mesh, order=order)
 
     fields = WaveState(
-            u=bump(actx, discr),
-            v=make_obj_array([discr.zeros(actx) for i in range(discr.dim)]),
+            u=bump(actx_outer, discr),
+            v=make_obj_array([discr.zeros(actx_outer) for i in range(discr.dim)]),
             )
 
     from meshmode.discretization.visualization import make_visualizer
-    vis = make_visualizer(actx, discr.volume_discr)
+    vis = make_visualizer(actx_outer, discr.volume_discr)
 
     def rhs(t, q):
-        return wave_operator(actx, discr, c=1, q=q)
+        return wave_operator(actx_rhs, discr, c=1, q=q)
 
-    t = 0
+    compiled_rhs = actx_rhs.compile(rhs)
+
+    def rhs_wrapper(t, q):
+        r = compiled_rhs(t, thaw(freeze(q, actx_outer), actx_rhs))
+        return thaw(freeze(r, actx_rhs), actx_outer)
+
+    t = np.float64(0)
     t_final = 3
     istep = 0
     while t < t_final:
-        fields = rk4_step(fields, t, dt, rhs)
+        fields = rk4_step(fields, t, dt, rhs_wrapper)
 
         if istep % 10 == 0:
             # FIXME: Maybe an integral function to go with the
             # DOFArray would be nice?
             assert len(fields.u) == 1
             logger.info("[%05d] t %.5e / %.5e norm %.5e",
-                    istep, t, t_final, actx.to_numpy(flat_norm(fields.u, 2)))
+                    istep, t, t_final, actx_outer.to_numpy(flat_norm(fields.u, 2)))
             vis.write_vtk_file("fld-wave-min-%04d.vtu" % istep, [
                 ("q", fields),
                 ])
@@ -528,6 +539,11 @@ def rhs(t, q):
 
 
 if __name__ == "__main__":
-    main()
+    import argparse
+    parser = argparse.ArgumentParser(description="Wave Equation Solver")
+    parser.add_argument("--lazy", action="store_true",
+                        help="switch to a lazy computation mode")
+    args = parser.parse_args()
+    main(lazy=args.lazy)
 
 # vim: foldmethod=marker

meshmode/array_context.py

@@ -27,6 +27,7 @@
 import sys
 from warnings import warn
 from arraycontext import PyOpenCLArrayContext as PyOpenCLArrayContextBase
+from arraycontext import PytatoPyOpenCLArrayContext as PytatoPyOpenCLArrayContextBase
 from arraycontext.pytest import (
         _PytestPyOpenCLArrayContextFactoryWithClass,
         register_pytest_array_context_factory)
@@ -303,4 +304,9 @@ def _import_names():
 # }}}
 
 
+class PytatoPyOpenCLArrayContext(PytatoPyOpenCLArrayContextBase):
+    def transform_loopy_program(self, t_unit):
+        # FIXME: Do not parallelize for now.
+        return t_unit
+
 # vim: foldmethod=marker

meshmode/discretization/__init__.py

@@ -583,7 +583,7 @@ def prg():
                     result[iel, idof] = \
                         sum(j, resampling_mat[idof, j] * nodes[iel, j])
                     """,
-                name="nodes")
+                name="lp_nodes")
 
             return lp.tag_inames(t_unit, {
                 "iel": ConcurrentElementInameTag(),

requirements.txt

@@ -6,6 +6,7 @@ git+https://github.com/inducer/pyvisfile.git#egg=pyvisfile
 git+https://github.com/inducer/modepy.git#egg=modepy
 git+https://github.com/inducer/pyopencl.git#egg=pyopencl
 git+https://github.com/inducer/islpy.git#egg=islpy
+git+https://github.com/inducer/pytato.git#egg=pytato
 
 # required by pytential, which is in turn needed for some tests
 git+https://github.com/inducer/pymbolic.git#egg=pymbolic