Merge branch 'main' into e2p_opt

Author: isuruf

doc/conf.py

@@ -18,16 +18,18 @@
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
-    "modepy": ("https://documen.tician.de/modepy/", None),
+    "sympy": ("https://docs.sympy.org/latest/", None),
+    "matplotlib": ("https://matplotlib.org/stable/", None),
     "pyopencl": ("https://documen.tician.de/pyopencl/", None),
+    "pytools": ("https://documen.tician.de/pytools/", None),
+    "modepy": ("https://documen.tician.de/modepy/", None),
     "pymbolic": ("https://documen.tician.de/pymbolic/", None),
     "loopy": ("https://documen.tician.de/loopy/", None),
     "pytential": ("https://documen.tician.de/pytential/", None),
     "boxtree": ("https://documen.tician.de/boxtree/", None),
-    "sympy": ("https://docs.sympy.org/latest/", None),
-    "matplotlib": ("https://matplotlib.org/stable/", None),
+    "arraycontext": ("https://documen.tician.de/arraycontext/", None),
 }
 
 nitpick_ignore_regex = [
-        ["py:class", r"symengine\.(.+)"], # :cry:
-        ]
+    ["py:class", r"symengine\.(.+)"],  # :cry:
+]

doc/misc.rst

@@ -5,6 +5,9 @@ Misc Tools
 
 .. automodule:: sumpy.symbolic
 
+.. automodule:: sumpy.tools
+
+.. automodule:: sumpy.array_context
 
 Installation
 ============

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, source_kernels=(vol_source_knl,),
+    p2p = P2P(actx.context,
+            source_kernels=(vol_source_knl,),
             target_kernels=(vol_target_knl,),
             exclude_self=False,
             value_dtypes=np.complex128)
@@ -157,35 +158,40 @@ def map_to_curve(t):
     lpot_kwargs = knl_kwargs.copy()
 
     if what_operator == "D":
-        volpot_kwargs["src_derivative_dir"] = native_curve.normal
+        volpot_kwargs["src_derivative_dir"] = actx.from_numpy(native_curve.normal)
 
     if what_operator_lpot == "D":
-        lpot_kwargs["src_derivative_dir"] = ovsmp_curve.normal
+        lpot_kwargs["src_derivative_dir"] = actx.from_numpy(ovsmp_curve.normal)
 
     if what_operator_lpot == "S'":
-        lpot_kwargs["tgt_derivative_dir"] = native_curve.normal
+        lpot_kwargs["tgt_derivative_dir"] = actx.from_numpy(native_curve.normal)
 
     # }}}
 
+    targets = actx.from_numpy(fp.points)
+    sources = actx.from_numpy(native_curve.pos)
+    ovsmp_sources = actx.from_numpy(ovsmp_curve.pos)
+
     if 0:
         # {{{ build matrix
 
         from fourier import make_fourier_interp_matrix
         fim = make_fourier_interp_matrix(novsmp, nsrc)
+
         from sumpy.tools import build_matrix
         from scipy.sparse.linalg import LinearOperator
 
         def apply_lpot(x):
             xovsmp = np.dot(fim, x)
             evt, (y,) = lpot(actx.queue,
-                    native_curve.pos,
-                    ovsmp_curve.pos,
-                    centers,
-                    [xovsmp * ovsmp_curve.speed * ovsmp_weights],
-                    expansion_radii=np.ones(centers.shape[1]),
+                    sources,
+                    ovsmp_sources,
+                    actx.from_numpy(centers),
+                    [actx.from_numpy(xovsmp * ovsmp_curve.speed * ovsmp_weights)],
+                    expansion_radii=actx.from_numpy(np.ones(centers.shape[1])),
                     **lpot_kwargs)
 
-            return y
+            return actx.to_numpy(y)
 
         op = LinearOperator((nsrc, nsrc), apply_lpot)
         mat = build_matrix(op, dtype=np.complex128)
@@ -200,19 +206,26 @@ def apply_lpot(x):
 
     mode_nr = 0
     density = np.cos(mode_nr*2*np.pi*native_t).astype(np.complex128)
-    ovsmp_density = np.cos(mode_nr*2*np.pi*ovsmp_t).astype(np.complex128)
+    strength = actx.from_numpy(native_curve.speed * native_weights * density)
+
     evt, (vol_pot,) = p2p(actx.queue,
-            fp.points,
-            native_curve.pos,
-            [native_curve.speed*native_weights*density], **volpot_kwargs)
+            targets,
+            sources,
+            [strength], **volpot_kwargs)
+    vol_pot = actx.to_numpy(vol_pot)
+
+    ovsmp_density = np.cos(mode_nr*2*np.pi*ovsmp_t).astype(np.complex128)
+    ovsmp_strength = actx.from_numpy(
+        ovsmp_curve.speed * ovsmp_weights * ovsmp_density)
 
     evt, (curve_pot,) = lpot(actx.queue,
-            native_curve.pos,
-            ovsmp_curve.pos,
-            centers,
-            [ovsmp_density * ovsmp_curve.speed * ovsmp_weights],
-            expansion_radii=np.ones(centers.shape[1]),
+            sources,
+            ovsmp_sources,
+            actx.from_numpy(centers),
+            [ovsmp_strength],
+            expansion_radii=actx.from_numpy(np.ones(centers.shape[1])),
             **lpot_kwargs)
+    curve_pot = actx.to_numpy(curve_pot)
 
     # }}}
 

sumpy/array_context.py

@@ -26,6 +26,9 @@
         register_pytest_array_context_factory)
 
 __doc__ = """
+Array Context
+-------------
+
 .. autoclass:: PyOpenCLArrayContext
 """
 

sumpy/e2e.py

@@ -29,6 +29,7 @@
 
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 from sumpy.tools import KernelCacheMixin, to_complex_dtype
+from sumpy.codegen import register_optimization_preambles
 from pytools import memoize_method
 
 import logging
@@ -144,7 +145,8 @@ def get_kernel(self):
     def get_optimized_kernel(self):
         # FIXME
         knl = self.get_kernel()
-        knl = lp.split_iname(knl, "itgt_box", 16, outer_tag="g.0")
+        knl = lp.split_iname(knl, "itgt_box", 64, outer_tag="g.0", inner_tag="l.0")
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -278,7 +280,8 @@ def get_kernel(self):
     def get_optimized_kernel(self):
         # FIXME
         knl = self.get_kernel()
-        knl = lp.split_iname(knl, "itgt_box", 16, outer_tag="g.0")
+        knl = lp.split_iname(knl, "itgt_box", 64, outer_tag="g.0", inner_tag="l.0")
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -518,6 +521,7 @@ def get_optimized_kernel(self, result_dtype):
         knl = self.get_kernel(result_dtype)
         knl = self.tgt_expansion.m2l_translation.optimize_loopy_kernel(
                 knl, self.tgt_expansion, self.src_expansion)
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -627,6 +631,7 @@ def get_optimized_kernel(self, result_dtype):
         knl = self.get_kernel(result_dtype)
         knl = lp.tag_inames(knl, "idim*:unr")
         knl = lp.tag_inames(knl, {"itr_class": "g.0"})
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -732,6 +737,7 @@ def get_optimized_kernel(self, result_dtype):
         _, optimizations = self.get_inner_knl_and_optimizations(result_dtype)
         for optimization in optimizations:
             knl = optimization(knl)
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):
@@ -831,6 +837,7 @@ def get_optimized_kernel(self, result_dtype):
         for optimization in optimizations:
             knl = optimization(knl)
         knl = lp.add_inames_for_unused_hw_axes(knl)
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):

sumpy/e2p.py

@@ -29,6 +29,7 @@
 import pymbolic.primitives as prim
 
 from sumpy.tools import KernelCacheMixin, gather_loopy_arguments
+from sumpy.codegen import register_optimization_preambles
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 
 
@@ -220,6 +221,7 @@ def get_optimized_kernel(self, max_ntargets_in_one_box):
         knl = lp.set_temporary_address_space(knl, "coeffs", lp.AddressSpace.LOCAL)
         knl = lp.set_options(knl,
             enforce_variable_access_ordered="no_check", write_code=False)
+        knl = register_optimization_preambles(knl, self.device)
 
         for transform in optimizations:
             knl = transform(knl)
@@ -400,6 +402,7 @@ def get_optimized_kernel(self, max_ntargets_in_one_box):
                 enforce_variable_access_ordered="no_check", write_code=False)
         for transform in optimizations:
             knl = transform(knl)
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):

sumpy/expansion/__init__.py

@@ -228,7 +228,7 @@ def __len__(self):
 
     def __eq__(self, other):
         return (
-                type(self) == type(other)
+                type(self) is type(other)
                 and self.kernel == other.kernel
                 and self.order == other.order
                 and self.use_rscale == other.use_rscale)

sumpy/expansion/local.py

@@ -75,7 +75,7 @@ def update_persistent_hash(self, key_hash, key_builder):
 
     def __eq__(self, other):
         return (
-            type(self) == type(other)
+            type(self) is type(other)
             and self.kernel == other.kernel
             and self.order == other.order
             and self.use_rscale == other.use_rscale

sumpy/kernel.py

@@ -102,9 +102,9 @@ def name(self):
     def __eq__(self, other):
         if id(self) == id(other):
             return True
-        if not type(self) == KernelArgument:
+        if type(self) is not KernelArgument:
             return NotImplemented
-        if not type(other) == KernelArgument:
+        if type(other) is not KernelArgument:
             return NotImplemented
         return self.loopy_arg == other.loopy_arg
 

sumpy/p2e.py

@@ -25,6 +25,7 @@
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 
 from sumpy.tools import KernelCacheMixin, KernelComputation
+from sumpy.codegen import register_optimization_preambles
 
 import logging
 logger = logging.getLogger(__name__)
@@ -118,6 +119,7 @@ def get_optimized_kernel(self, sources_is_obj_array, centers_is_obj_array):
         knl = self._allow_redundant_execution_of_knl_scaling(knl)
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):