Enable arraycontext + lazy evaluation + loop fusion

.github/workflows/ci.yml

@@ -101,8 +101,8 @@ jobs:
             run: |
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
-                if [[ "$DOWNSTREAM_PROJECT" == "pytential" && "$GITHUB_HEAD_REF" == "e2p" ]]; then
-                    DOWNSTREAM_PROJECT=https://github.com/isuruf/pytential.git@e2p
+                if [[ "$DOWNSTREAM_PROJECT" == "pytential" && "$GITHUB_HEAD_REF" == "towards-array-context" ]]; then
+                   DOWNSTREAM_PROJECT=https://github.com/alexfikl/pytential.git@towards-array-context
                 fi
                 test_downstream "$DOWNSTREAM_PROJECT"
 

examples/cost-model-sumpy-fused.py

@@ -0,0 +1,248 @@
+import logging
+import os
+
+import numpy as np
+
+import matplotlib.pyplot as plt
+
+import pyopencl as cl
+
+from boxtree import TreeBuilder
+from boxtree.array_context import PyOpenCLArrayContext
+from sumpy.array_context import PytatoPyOpenCLArrayContext
+
+from boxtree.cost import FMMCostModel, make_pde_aware_translation_cost_model
+from boxtree.fmm import drive_fmm
+from boxtree.traversal import FMMTraversalBuilder
+from boxtree.tools import make_normal_particle_array as p_normal
+
+from functools import partial
+
+from sumpy.expansion.multipole import VolumeTaylorMultipoleExpansion
+from sumpy.expansion.local import VolumeTaylorLocalExpansion 
+from sumpy.kernel import LaplaceKernel
+from sumpy.expansion.m2l import NonFFTM2LTranslationClassFactory
+from sumpy.fmm import (
+    SumpyTreeIndependentDataForWrangler, 
+    SumpyExpansionWrangler
+)
+
+logging.basicConfig(level=os.environ.get("LOGLEVEL", "WARNING"))
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def demo_cost_model(plot_results=False, lazy=False):
+
+    # {{{ useful variables and actx setup
+
+    nsources_list = [5000]
+    ntargets_list = [5000]
+    #nsources_list = [100, 200, 300, 400, 500]
+    #ntargets_list = [100, 200, 300, 400, 500]
+    nparticles_per_box_list = [32, 64, 128, 256, 512]
+    #nparticles_per_box_list = [32, 64, 128]
+    dim = 2
+    dtype = np.float64
+
+    ctx = cl.create_some_context()
+    queue = cl.CommandQueue(ctx)
+    actx_boxtree = PyOpenCLArrayContext(queue, force_device_scalars=True)
+
+    traversals = []
+    traversals_dev = []
+    level_orders_list = []
+    timing_results = []
+    results = {}
+    fields = ["form_multipoles", "eval_direct", "multipole_to_local", 
+              "eval_multipoles", "form_locals", "eval_locals", 
+              "coarsen_multipoles", "refine_locals"]
+    for field in fields:
+        results[field] = []
+
+    # }}}
+
+    def fmm_level_to_order(kernel, kernel_args, tree, ilevel):
+        return 10
+
+    timings = {}
+    for nparticles_per_box in nparticles_per_box_list:
+        for nsources, ntargets in zip(nsources_list, ntargets_list):
+            logger.info(f"Testing nsources = {nsources}, ntargets = {ntargets} "
+                        f"with nparticles per box = {nparticles_per_box}")
+
+            # {{{ Generate sources, targets and target_radii
+
+            sources = p_normal(actx_boxtree, nsources, dim, dtype, seed=15)
+            targets = p_normal(actx_boxtree, ntargets, dim, dtype, seed=18)
+
+            rng = np.random.default_rng(seed=22)
+            target_radii = rng.uniform(low=0.0, high=0.05, size=ntargets)
+
+            # }}}
+
+            # {{{ Generate tree and traversal
+
+            tb = TreeBuilder(actx_boxtree)
+            tree, _ = tb(
+                actx_boxtree, sources, targets=targets, target_radii=target_radii,
+                stick_out_factor=0.15, 
+                max_particles_in_box=nparticles_per_box, debug=True
+            )
+
+            tg = FMMTraversalBuilder(actx_boxtree, well_sep_is_n_away=2)
+            trav_dev, _ = tg(actx_boxtree, tree, debug=True)
+            #trav = actx.to_numpy(trav_dev)
+            trav = trav_dev
+
+            traversals.append(trav)
+            traversals_dev.append(trav_dev)
+
+            # }}}
+
+            # {{{ snag queue from eager tree building arraycontext
+
+            if lazy:
+                queue = actx_boxtree.queue
+                actx = PytatoPyOpenCLArrayContext(queue)
+            else:
+                actx = actx_boxtree 
+
+            # }}}
+
+            # {{{ define kernel and expansion classes
+
+            kernel = LaplaceKernel(dim)
+
+            mpole_expansion_cls = VolumeTaylorMultipoleExpansion 
+            local_expansion_cls = VolumeTaylorLocalExpansion
+            m2l_factory = NonFFTM2LTranslationClassFactory()
+            m2l = m2l_factory.get_m2l_translation_class(kernel, 
+                                                        local_expansion_cls)()
+
+            # }}}
+
+            # {{{ define interface for fmm driver
+
+            tree_indep = SumpyTreeIndependentDataForWrangler(
+                actx,
+                partial(mpole_expansion_cls, kernel),
+                partial(local_expansion_cls, kernel, m2l_translation=m2l),
+                [kernel])
+
+            wrangler = SumpyExpansionWrangler(
+                tree_indep, 
+                trav,
+                np.float64,
+                fmm_level_to_order=fmm_level_to_order)
+
+            level_orders_list.append(wrangler.level_orders)
+
+            # }}}
+
+            # {{{ fmm
+
+            timing_data = {}
+            src_weights = np.random.rand(tree.nsources).astype(tree.coord_dtype)
+            drive_fmm(actx, wrangler, (src_weights,), timing_data=timing_data)
+            timing_results.append(timing_data)
+
+            # def driver(src_weights):
+            #     return drive_fmm(actx, wrangler, (src_weights,),
+            #                      timing_data=timing_data)
+            #
+            # src_weights = actx.from_numpy(src_weights)
+            # actx.compile(driver)(src_weights)
+
+            # }}}
+
+        # {{{ build cost model 
+
+        time_field_name = "process_elapsed"
+
+        cost_model = FMMCostModel(make_pde_aware_translation_cost_model)
+
+        model_results = []
+        for icase in range(len(traversals)-1):
+            traversal = traversals_dev[icase]
+            model_results.append(
+                cost_model.cost_per_stage(
+                    actx_boxtree, traversal, level_orders_list[icase],
+                    FMMCostModel.get_unit_calibration_params(),
+                )
+            )
+
+        # }}}
+
+        queue.finish()
+
+        if not timing_results:
+            return
+
+        # {{{ analyze and report cost model results
+
+        params = cost_model.estimate_calibration_params(
+            model_results, timing_results[:-1], time_field_name=time_field_name
+        )
+
+        predicted_time = cost_model.cost_per_stage(
+            actx_boxtree, traversals_dev[-1], level_orders_list[-1], params,
+        )
+
+        queue.finish()
+
+        for field in ["form_multipoles", "eval_direct", "multipole_to_local",
+                      "eval_multipoles", "form_locals", "eval_locals",
+                      "coarsen_multipoles", "refine_locals"]:
+            # measured = timing_results[-1][field]["process_elapsed"]
+            # pred_err = (
+            #         (measured - predicted_time[field])
+            #         / measured)
+            # logger.info("actual/predicted time for %s: %.3g/%.3g -> %g %% error",
+            #         field,
+            #         measured,
+            #         predicted_time[field],
+            #         abs(100*pred_err))
+
+            if nparticles_per_box != nparticles_per_box_list[0]:
+                results[field].append(predicted_time[field])
+
+        # }}}
+
+
+    if plot_results:
+        x = np.arange(len(nparticles_per_box_list) - 1)
+        width = 0.1
+        mult = 0
+
+        fig, ax = plt.subplots()
+
+        for field, timing in results.items():
+            offset = width*mult
+            bars = ax.bar(x + offset, timing, width, label=field)
+            #ax.bar_label(bars, padding=4)
+            mult += 1
+
+        ax.set_xlabel("Particles per box")
+        ax.set_xticks(x + 3*width, nparticles_per_box_list[1:])
+
+        ax.set_ylabel("Process time (s)")
+
+        ax.legend(loc="upper left", ncols=3)
+
+        #plt.show()
+        plt.savefig("./balancing.pdf")
+
+
+if __name__ == "__main__":
+
+    import argparse
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--lazy", action="store_true")
+    parser.add_argument("--plot_results", action="store_true")
+
+    args = parser.parse_args()
+
+    demo_cost_model(lazy=args.lazy, plot_results=args.plot_results)

examples/curve-pot.py

@@ -91,7 +91,8 @@ def draw_pot_figure(aspect_ratio,
         knl_kwargs = {}
 
     vol_source_knl, vol_target_knl = process_kernel(knl, what_operator)
-    p2p = P2P(actx.context,
+    p2p = P2P(
+            actx,
             source_kernels=(vol_source_knl,),
             target_kernels=(vol_target_knl,),
             exclude_self=False,
@@ -100,7 +101,8 @@ def draw_pot_figure(aspect_ratio,
     lpot_source_knl, lpot_target_knl = process_kernel(knl, what_operator_lpot)
 
     from sumpy.qbx import LayerPotential
-    lpot = LayerPotential(actx.context,
+    lpot = LayerPotential(
+            actx,
             expansion=expn_class(knl, order=order),
             source_kernels=(lpot_source_knl,),
             target_kernels=(lpot_target_knl,),
@@ -183,7 +185,8 @@ def map_to_curve(t):
 
         def apply_lpot(x):
             xovsmp = np.dot(fim, x)
-            evt, (y,) = lpot(actx.queue,
+            y, = lpot(
+                    actx,
                     sources,
                     ovsmp_sources,
                     actx.from_numpy(centers),
@@ -208,7 +211,7 @@ def apply_lpot(x):
     density = np.cos(mode_nr*2*np.pi*native_t).astype(np.complex128)
     strength = actx.from_numpy(native_curve.speed * native_weights * density)
 
-    evt, (vol_pot,) = p2p(actx.queue,
+    vol_pot, = p2p(
             targets,
             sources,
             [strength], **volpot_kwargs)
@@ -218,7 +221,8 @@ def apply_lpot(x):
     ovsmp_strength = actx.from_numpy(
         ovsmp_curve.speed * ovsmp_weights * ovsmp_density)
 
-    evt, (curve_pot,) = lpot(actx.queue,
+    curve_pot, = lpot(
+            actx,
             sources,
             ovsmp_sources,
             actx.from_numpy(centers),
@@ -305,11 +309,11 @@ def apply_lpot(x):
         plt.savefig("eigvals-ext-nsrc100-novsmp100.pdf")
         plt.clf()
 
-    # draw_pot_figure(
-    #         aspect_ratio=1, nsrc=100, novsmp=100, helmholtz_k=0,
-    #         what_operator="D", what_operator_lpot="D", force_center_side=-1)
-    # plt.savefig("eigvals-int-nsrc100-novsmp100.pdf")
-    # plt.clf()
+    #draw_pot_figure(
+    #        aspect_ratio=1, nsrc=100, novsmp=100, helmholtz_k=0,
+    #        what_operator="D", what_operator_lpot="D", force_center_side=-1)
+    #plt.savefig("eigvals-int-nsrc100-novsmp100.pdf")
+    #plt.clf()
 
     # draw_pot_figure(
     #         aspect_ratio=1, nsrc=100, novsmp=200, helmholtz_k=0,

examples/expansion-toys.py

@@ -23,7 +23,6 @@ def main():
     actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
 
     tctx = t.ToyContext(
-            actx.context,
             # LaplaceKernel(2),
             YukawaKernel(2), extra_kernel_kwargs={"lam": 5},
             # HelmholtzKernel(2), extra_kernel_kwargs={"k": 0.3},
@@ -37,22 +36,22 @@ def main():
     fp = FieldPlotter([3, 0], extent=8)
 
     if USE_MATPLOTLIB:
-        t.logplot(fp, pt_src, cmap="jet")
+        t.logplot(actx, fp, pt_src, cmap="jet")
         plt.colorbar()
         plt.show()
 
-    mexp = t.multipole_expand(pt_src, [0, 0], 5)
-    mexp2 = t.multipole_expand(mexp, [0, 0.25])  # noqa: F841
-    lexp = t.local_expand(mexp, [3, 0])
-    lexp2 = t.local_expand(lexp, [3, 1], 3)
+    mexp = t.multipole_expand(actx, pt_src, [0, 0], order=5)
+    mexp2 = t.multipole_expand(actx, mexp, [0, 0.25])  # noqa: F841
+    lexp = t.local_expand(actx, mexp, [3, 0])
+    lexp2 = t.local_expand(actx, lexp, [3, 1], order=3)
 
     # diff = mexp - pt_src
     # diff = mexp2 - pt_src
     diff = lexp2 - pt_src
 
-    print(t.l_inf(diff, 1.2, center=lexp2.center))
+    print(t.l_inf(actx, diff, 1.2, center=lexp2.center))
     if USE_MATPLOTLIB:
-        t.logplot(fp, diff, cmap="jet", vmin=-3, vmax=0)
+        t.logplot(actx, fp, diff, cmap="jet", vmin=-3, vmax=0)
         plt.colorbar()
         plt.show()
 

examples/fourier.py

@@ -21,8 +21,8 @@ def make_fourier_mode_extender(m, n, dtype):
     else:
         peak_pos_freq = (k-1)/2
 
-    num_pos_freq = peak_pos_freq + 1
-    num_neg_freq = k-num_pos_freq
+    num_pos_freq = int(peak_pos_freq + 1)
+    num_neg_freq = int(k-num_pos_freq)
 
     eye = np.eye(k)
     result[:num_pos_freq, :num_pos_freq] = eye[:num_pos_freq, :num_pos_freq]

requirements.txt

@@ -10,7 +10,7 @@ git+https://github.com/inducer/pytools.git#egg=pytools
 git+https://github.com/inducer/pymbolic.git#egg=pymbolic
 git+https://github.com/inducer/islpy.git#egg=islpy
 git+https://github.com/inducer/pyopencl.git#egg=pyopencl
-git+https://github.com/inducer/boxtree.git#egg=boxtree
+git+https://github.com/alexfikl/boxtree.git@towards-array-context#egg=boxtree
 git+https://github.com/inducer/loopy.git#egg=loopy
 git+https://github.com/inducer/arraycontext.git#egg=arraycontext
 git+https://github.com/inducer/pyfmmlib.git#egg=pyfmmlib