Optimize M2P

sumpy/expansion/__init__.py

@@ -22,7 +22,8 @@
 
 from abc import ABC, abstractmethod
 from typing import (
-        Any, ClassVar, Dict, Hashable, List, Optional, Sequence, Tuple, Type)
+        Any, ClassVar, Dict, Hashable, List, Optional, Sequence, Tuple, Type,
+        Callable)
 
 from pytools import memoize_method
 import loopy as lp
@@ -66,9 +67,9 @@ class ExpansionBase(ABC):
     .. automethod:: get_coefficient_identifiers
     .. automethod:: coefficients_from_source
     .. automethod:: coefficients_from_source_vec
-    .. automethod:: loopy_expansion_formation
     .. automethod:: evaluate
-    .. automethod:: loopy_evaluator
+    .. automethod:: loopy_expansion_formation
+    .. automethod:: loopy_evaluator_and_optimizations
 
     .. automethod:: with_kernel
     .. automethod:: copy
@@ -183,7 +184,9 @@ def evaluate(self, kernel, coeffs, bvec, rscale, sac=None):
             in *coeffs*.
         """
 
-    def loopy_evaluator(self, kernels: Sequence[Kernel]) -> lp.TranslationUnit:
+    def loopy_evaluator_and_optimizations(self, kernels: Sequence[Kernel]) \
+            -> Tuple[lp.TranslationUnit, Sequence[
+                Callable[[lp.TranslationUnit], lp.TranslationUnit]]]:
         """
         :returns: a :mod:`loopy` kernel that returns the evaluated
             target transforms of the potential given by *kernels*.

sumpy/expansion/local.py

@@ -32,6 +32,11 @@
         VolumeTaylorExpansionMixin,
         LinearPDEConformingVolumeTaylorExpansion)
 from sumpy.tools import add_to_sac, mi_increment_axis
+from sumpy.kernel import Kernel
+
+import loopy as lp
+
+from typing import Sequence, Callable, Tuple
 
 import logging
 logger = logging.getLogger(__name__)
@@ -405,6 +410,16 @@ def loopy_translate_from(self, src_expansion):
             f"A direct loopy kernel for translation from "
             f"{src_expansion} to {self} is not implemented.")
 
+    def loopy_evaluator_and_optimizations(self, kernels: Sequence[Kernel]) \
+            -> Tuple[lp.TranslationUnit, Sequence[
+                Callable[[lp.TranslationUnit], lp.TranslationUnit]]]:
+        from sumpy.expansion.loopy import (make_l2p_loopy_kernel_for_volume_taylor,
+            make_e2p_loopy_kernel)
+        try:
+            return make_l2p_loopy_kernel_for_volume_taylor(self, kernels)
+        except NotImplementedError:
+            return make_e2p_loopy_kernel(self, kernels)
+
 
 class VolumeTaylorLocalExpansion(
         VolumeTaylorExpansion,