get_loopy_evaluator -> loopy_evaluator_and_optimizations

Author: isuruf

sumpy/e2p.py

@@ -85,15 +85,15 @@ def default_name(self):
         pass
 
     @memoize_method
-    def get_cached_loopy_knl_and_optimizations(self):
-        return self.expansion.loopy_evaluator(self.kernels)
+    def get_loopy_evaluator_and_optimizations(self):
+        return self.expansion.loopy_evaluator_and_optimizations(self.kernels)
 
     def get_cache_key(self):
         return (type(self).__name__, self.expansion, tuple(self.kernels))
 
     def add_loopy_eval_callable(
             self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit:
-        inner_knl, _ = self.get_cached_loopy_knl_and_optimizations()
+        inner_knl, _ = self.get_loopy_evaluator_and_optimizations()
         loopy_knl = lp.merge([loopy_knl, inner_knl])
         loopy_knl = lp.inline_callable_kernel(loopy_knl, "e2p")
         for kernel in self.kernels:
@@ -208,7 +208,7 @@ def get_kernel(self, max_ntargets_in_one_box):
         return loopy_knl
 
     def get_optimized_kernel(self, max_ntargets_in_one_box):
-        inner_knl, optimizations = self.get_cached_loopy_knl_and_optimizations()
+        inner_knl, optimizations = self.get_loopy_evaluator_and_optimizations()
         knl = self.get_kernel(max_ntargets_in_one_box=max_ntargets_in_one_box)
         knl = lp.tag_inames(knl, {"itgt_box": "g.0"})
         knl = lp.split_iname(knl, "itgt_offset", 256, inner_tag="l.0")
@@ -378,7 +378,7 @@ def get_kernel(self, max_ntargets_in_one_box):
         return loopy_knl
 
     def get_optimized_kernel(self, max_ntargets_in_one_box):
-        _, optimizations = self.get_cached_loopy_knl_and_optimizations()
+        _, optimizations = self.get_loopy_evaluator_and_optimizations()
         knl = self.get_kernel(max_ntargets_in_one_box=max_ntargets_in_one_box)
         knl = lp.tag_inames(knl, {"itgt_box": "g.0", "dummy": "l.0"})
         knl = lp.unprivatize_temporaries_with_inames(knl,

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

@@ -36,7 +36,7 @@
 
 import loopy as lp
 
-from typing import Sequence
+from typing import Sequence, Callable, Tuple
 
 import logging
 logger = logging.getLogger(__name__)
@@ -410,7 +410,9 @@ def loopy_translate_from(self, src_expansion):
             f"A direct loopy kernel for translation from "
             f"{src_expansion} to {self} is not implemented.")
 
-    def get_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]]]:
         from sumpy.expansion.loopy import (make_l2p_loopy_kernel_for_volume_taylor,
             make_e2p_loopy_kernel)
         try:

sumpy/expansion/loopy.py

@@ -20,7 +20,7 @@
 THE SOFTWARE.
 """
 
-from typing import Sequence
+from typing import Sequence, Callable, Tuple
 import pymbolic
 import pymbolic.primitives as prim
 import loopy as lp
@@ -36,8 +36,9 @@
 logger = logging.getLogger(__name__)
 
 
-def make_e2p_loopy_kernel(
-        expansion: ExpansionBase, kernels: Sequence[Kernel]) -> lp.TranslationUnit:
+def make_e2p_loopy_kernel(expansion: ExpansionBase, kernels: Sequence[Kernel]) \
+        -> Tuple[lp.TranslationUnit, Sequence[
+            Callable[[lp.TranslationUnit], lp.TranslationUnit]]]:
     """
     This is a helper function to create a loopy kernel for multipole/local
     evaluation. This function uses symbolic expressions given by the expansion class,