Updated sumpy and boxtree to arraycontext branches

doc/conf.py

@@ -18,6 +18,10 @@
         "DOFDescriptorLike": "pytential.symbolic.dof_desc.DOFDescriptorLike",
         }
 
+nitpick_ignore_regex = [
+        ["py:class", r"boxtree\.array_context\.(.+)"],
+        ]
+
 intersphinx_mapping = {
     "arraycontext": ("https://documen.tician.de/arraycontext", None),
     "boxtree": ("https://documen.tician.de/boxtree", None),

doc/qbx.rst

@@ -8,6 +8,11 @@ QBX internals
     This page documents :mod:`pytential` internals and is not typically
     needed in end-user applications.
 
+Array Context
+-------------
+
+.. automodule:: pytential.array_context
+
 Refinement
 ----------
 

examples/helmholtz-dirichlet.py

@@ -1,13 +1,13 @@
 import numpy as np
 import numpy.linalg as la
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 
@@ -80,7 +80,7 @@
             fmm_order=fmm_order
             )
 
     from sumpy.visualization import FieldPlotter
     fplot = FieldPlotter(np.zeros(2), extent=5, npoints=500)
     targets = actx.from_numpy(fplot.points)
 
@@ -94,7 +94,7 @@
 
     # {{{ describe bvp
 
     from sumpy.kernel import HelmholtzKernel, LaplaceKernel
     kernel = HelmholtzKernel(2)
 
     sigma_sym = sym.var("sigma")

examples/laplace-dirichlet-3d.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 
@@ -72,7 +72,7 @@
             fmm_order=fmm_order,
             )
 
     from sumpy.visualization import FieldPlotter
     fplot = FieldPlotter(np.zeros(3), extent=20, npoints=50)
     targets = actx.from_numpy(fplot.points)
 
@@ -86,7 +86,7 @@
 
     # {{{ describe bvp
 
     from sumpy.kernel import LaplaceKernel
     kernel = LaplaceKernel(3)
 
     sigma_sym = sym.var("sigma")

examples/laplace-dirichlet-simple.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 
@@ -59,7 +59,7 @@
             fmm_order=fmm_order
             )
 
     from sumpy.visualization import FieldPlotter
     fplot = FieldPlotter(np.zeros(2), extent=5, npoints=500)
     targets = actx.from_numpy(fplot.points)
 
@@ -73,7 +73,7 @@
 
     # {{{ describe bvp
 
     from sumpy.kernel import LaplaceKernel
     kernel = LaplaceKernel(2)
 
     sigma_sym = sym.var("sigma")

examples/scaling-study.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 

examples/cost.py → experiments/cost.py

@@ -138,8 +138,7 @@ def calibrate_cost_model(ctx):
 
         for _ in range(RUNS):
             timing_data = {}
-            bound_op.eval({"sigma": sigma}, array_context=actx,
-                    timing_data=timing_data)
+            bound_op.eval({"sigma": sigma}, array_context=actx)
 
             model_results.append(modeled_cost)
             timing_results.append(timing_data)
@@ -175,8 +174,7 @@ def test_cost_model(ctx, calibration_params):
         temp_timing_results = []
         for _ in range(RUNS):
             timing_data = {}
-            bound_op.eval({"sigma": sigma},
-                    array_context=actx, timing_data=timing_data)
+            bound_op.eval({"sigma": sigma}, array_context=actx)
             temp_timing_results.append(one(timing_data.values()))
 
         timing_result = {}

pytential/array_context.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+
+__copyright__ = "Copyright (C) 2022 Alexandru Fikl"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+from arraycontext.pytest import (
+    _PytestPyOpenCLArrayContextFactoryWithClass,
+    register_pytest_array_context_factory,
+)
+from sumpy.array_context import (  # noqa: F401
+    PyOpenCLArrayContext as SumpyPyOpenCLArrayContext,
+    make_loopy_program,
+)
+
+
+__doc__ = """
+.. autoclass:: PyOpenCLArrayContext
+"""
+
+
+# {{{ PyOpenCLArrayContext
+
+class PyOpenCLArrayContext(SumpyPyOpenCLArrayContext):
+    def transform_loopy_program(self, t_unit):
+        kernel = t_unit.default_entrypoint
+        options = kernel.options
+
+        if not options.return_dict or not options.no_numpy:
+            raise ValueError(
+                "loopy kernels passed to 'call_loopy' must have 'return_dict' "
+                "and 'no_numpy' options set. Did you use 'make_loopy_program' "
+                f"to create the kernel '{kernel.name}'?")
+
+        # FIXME: this probably needs some proper logic
+        from meshmode.array_context import _transform_loopy_inner
+        transformed_t_unit = _transform_loopy_inner(t_unit)
+
+        if transformed_t_unit is not None:
+            return transformed_t_unit
+
+        return t_unit
+
+# }}}
+
+
+# {{{ pytest
+
+def _acf():
+    import pyopencl as cl
+    ctx = cl.create_some_context()
+    queue = cl.CommandQueue(ctx)
+
+    return PyOpenCLArrayContext(queue, force_device_scalars=True)
+
+
+class PytestPyOpenCLArrayContextFactory(
+        _PytestPyOpenCLArrayContextFactoryWithClass):
+    actx_class = PyOpenCLArrayContext
+
+    def __call__(self):
+        # NOTE: prevent any cache explosions during testing!
+        from sympy.core.cache import clear_cache
+        clear_cache()
+
+        return super().__call__()
+
+
+register_pytest_array_context_factory(
+    "pytential.pyopencl",
+    PytestPyOpenCLArrayContextFactory)
+
+# }}}

pytential/linalg/proxy.py

@@ -120,28 +120,29 @@
         from pytential.qbx.utils import tree_code_container
         tcc = tree_code_container(lpot_source._setup_actx)
 
-        tree, _ = tcc.build_tree()(actx.queue,
+        tree, _ = tcc.build_tree()(actx,
                 particles=flatten(
                     actx.thaw(discr.nodes()), actx, leaf_class=DOFArray
                     ),
                 max_particles_in_box=max_particles_in_box,
                 kind=tree_kind)
 
         from boxtree import box_flags_enum
-        tree = tree.get(actx.queue)
-        # FIXME: maybe this should use IS_LEAF once available?
+
+        tree = actx.to_numpy(tree)
+        # FIXME maybe this should use IS_LEAF once available?
         assert tree.box_flags is not None
         leaf_boxes, = (
                 tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0
                 ).nonzero()
 
         indices = np.empty(len(leaf_boxes), dtype=object)
         starts: onp.Array1D[np.integer] | None = None
 
         for i, ibox in enumerate(leaf_boxes):
             box_start = tree.box_source_starts[ibox]
             box_end = box_start + tree.box_source_counts_cumul[ibox]
             indices[i] = tree.user_source_ids[box_start:box_end]
     else:
         if discr.ambient_dim != 2 and discr.dim == 1:
             raise ValueError("only curves are supported for 'tree_kind=None'")
@@ -766,12 +767,11 @@
     assert isinstance(setup_actx, PyOpenCLArrayContext)
 
     tcc = tree_code_container(setup_actx)
-    tree, _ = tcc.build_tree()(actx.queue, targets,
-            max_particles_in_box=max_particles_in_box)
-    query, _ = tcc.build_area_query()(actx.queue, tree, pxy.centers, pxy.radii)
+    tree, _ = tcc.build_tree()(actx, targets, max_particles_in_box=max_particles_in_box)
+    query, _ = tcc.build_area_query()(actx, tree, pxy.centers, pxy.radii)
 
-    tree = tree.get(actx.queue)
-    query = query.get(actx.queue)
+    tree = actx.to_numpy(tree)
+    query = actx.to_numpy(query)
 
     # }}}