diff --git a/sumpy/e2e.py b/sumpy/e2e.py index bcb02f126a0040329947104e9f9fe268a1efd38a..9c96ac569bb7212cd35cc98b853f986a563fdb70 100644 --- a/sumpy/e2e.py +++ b/sumpy/e2e.py @@ -20,6 +20,8 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +from abc import ABC, abstractmethod + import numpy as np import loopy as lp import sumpy.symbolic as sym @@ -48,9 +50,7 @@ Expansion-to-expansion # {{{ translation base class -class E2EBase(KernelCacheMixin): - default_name = "e2e" - +class E2EBase(KernelCacheMixin, ABC): def __init__(self, ctx, src_expansion, tgt_expansion, name=None, device=None): """ @@ -94,6 +94,11 @@ class E2EBase(KernelCacheMixin): self.dim = src_expansion.dim + @property + @abstractmethod + def default_name(self): + pass + @memoize_method def get_translation_loopy_insns(self): from sumpy.symbolic import make_sym_vector @@ -132,8 +137,9 @@ class E2EBase(KernelCacheMixin): self.tgt_expansion, ) + @abstractmethod def get_kernel(self): - raise NotImplementedError + pass def get_optimized_kernel(self): # FIXME @@ -152,13 +158,15 @@ class E2EFromCSR(E2EBase): list. """ - default_name = "e2e_from_csr" - def __init__(self, ctx, src_expansion, tgt_expansion, name=None, device=None): super().__init__(ctx, src_expansion, tgt_expansion, name=name, device=device) + @property + def default_name(self): + return "e2e_from_csr" + def get_translation_loopy_insns(self): from sumpy.symbolic import make_sym_vector dvec = make_sym_vector("d", self.dim) @@ -302,7 +310,9 @@ class M2LUsingTranslationClassesDependentData(E2EFromCSR): list using M2L translation classes dependent data """ - default_name = "m2l_using_translation_classes_dependent_data" + @property + def default_name(self): + return "m2l_using_translation_classes_dependent_data" def get_translation_loopy_insns(self, result_dtype): from sumpy.symbolic import make_sym_vector @@ -538,7 +548,10 @@ class M2LGenerateTranslationClassesDependentData(E2EBase): """Implements precomputing the M2L kernel dependent data which are translation classes dependent derivatives. """ - default_name = "m2l_generate_translation_classes_dependent_data" + + @property + def default_name(self): + return "m2l_generate_translation_classes_dependent_data" def get_kernel(self, result_dtype): m2l_translation = self.tgt_expansion.m2l_translation @@ -649,7 +662,9 @@ class M2LGenerateTranslationClassesDependentData(E2EBase): class M2LPreprocessMultipole(E2EBase): """Computes the preprocessed multipole expansion for accelerated M2L""" - default_name = "m2l_preprocess_multipole" + @property + def default_name(self): + return "m2l_preprocess_multipole" def get_kernel(self, result_dtype): m2l_translation = self.tgt_expansion.m2l_translation @@ -731,7 +746,9 @@ class M2LPreprocessMultipole(E2EBase): class M2LPostprocessLocal(E2EBase): """Postprocesses locals expansions for accelerated M2L""" - default_name = "m2l_postprocess_local" + @property + def default_name(self): + return "m2l_postprocess_local" def get_kernel(self, result_dtype): m2l_translation = self.tgt_expansion.m2l_translation @@ -816,7 +833,9 @@ class M2LPostprocessLocal(E2EBase): # {{{ translation from a box's children class E2EFromChildren(E2EBase): - default_name = "e2e_from_children" + @property + def default_name(self): + return "e2e_from_children" def get_kernel(self): ncoeffs_src = len(self.src_expansion) @@ -933,7 +952,9 @@ class E2EFromChildren(E2EBase): # {{{ translation from a box's parent class E2EFromParent(E2EBase): - default_name = "e2e_from_parent" + @property + def default_name(self): + return "e2e_from_parent" def get_kernel(self): ncoeffs_src = len(self.src_expansion) diff --git a/sumpy/e2p.py b/sumpy/e2p.py index 34bdd5f520cf077301b8f0d487bac87f346f8737..611b2718d6e22f8b3741cc4c80f6d912f2be8d1e 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -20,6 +20,8 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +from abc import ABC, abstractmethod + import numpy as np import loopy as lp import sumpy.symbolic as sym @@ -42,9 +44,7 @@ Expansion-to-particle # {{{ E2P base class -class E2PBase(KernelCacheMixin): - default_name = "e2p" - +class E2PBase(KernelCacheMixin, ABC): def __init__(self, ctx, expansion, kernels, name=None, device=None): """ @@ -77,6 +77,11 @@ class E2PBase(KernelCacheMixin): self.dim = expansion.dim + @property + @abstractmethod + def default_name(self): + pass + def get_loopy_insns_and_result_names(self): from sumpy.symbolic import make_sym_vector bvec = make_sym_vector("b", self.dim) @@ -132,7 +137,9 @@ class E2PBase(KernelCacheMixin): # {{{ E2P to single box (L2P, likely) class E2PFromSingleBox(E2PBase): - default_name = "e2p_from_single_box" + @property + def default_name(self): + return "e2p_from_single_box" def get_kernel(self): ncoeffs = len(self.expansion) @@ -234,7 +241,9 @@ class E2PFromSingleBox(E2PBase): # {{{ E2P from CSR-like interaction list class E2PFromCSR(E2PBase): - default_name = "e2p_from_csr" + @property + def default_name(self): + return "e2p_from_csr" def get_kernel(self): ncoeffs = len(self.expansion) diff --git a/sumpy/p2e.py b/sumpy/p2e.py index 21551aefadcdcba33c11f07374ea954c4882a2c0..9532f6ecf4c16270236651d149e384f5fb2951b6 100644 --- a/sumpy/p2e.py +++ b/sumpy/p2e.py @@ -49,8 +49,6 @@ class P2EBase(KernelCacheMixin, KernelComputation): .. automethod:: __init__ """ - default_name = "p2e" - def __init__(self, ctx, expansion, kernels=None, name=None, device=None, strength_usage=None): """ @@ -124,9 +122,6 @@ class P2EBase(KernelCacheMixin, KernelComputation): return (type(self).__name__, self.name, self.expansion, tuple(self.source_kernels), tuple(self.strength_usage)) - def get_kernel(self): - raise NotImplementedError - def get_optimized_kernel(self, sources_is_obj_array, centers_is_obj_array): knl = self.get_kernel() @@ -165,7 +160,9 @@ class P2EFromSingleBox(P2EBase): .. automethod:: __call__ """ - default_name = "p2e_from_single_box" + @property + def default_name(self): + return "p2e_from_single_box" def get_kernel(self): ncoeffs = len(self.expansion) @@ -267,7 +264,9 @@ class P2EFromCSR(P2EBase): .. automethod:: __call__ """ - default_name = "p2e_from_csr" + @property + def default_name(self): + return "p2e_from_csr" def get_kernel(self): ncoeffs = len(self.expansion) diff --git a/sumpy/p2p.py b/sumpy/p2p.py index abd5df0658a02a5da7be8762e4156c60419be19c..c24fa8d9154eeda184f0f394df212e51eab4c3e4 100644 --- a/sumpy/p2p.py +++ b/sumpy/p2p.py @@ -51,8 +51,6 @@ Particle-to-particle # {{{ p2p base class class P2PBase(KernelCacheMixin, KernelComputation): - default_name = "p2p" - def __init__(self, ctx, target_kernels, exclude_self, strength_usage=None, value_dtypes=None, name=None, device=None, source_kernels=None): """ @@ -175,9 +173,6 @@ class P2PBase(KernelCacheMixin, KernelComputation): if self.exclude_self else []) + gather_loopy_source_arguments(self.source_kernels)) - def get_kernel(self): - raise NotImplementedError - def get_optimized_kernel(self, targets_is_obj_array, sources_is_obj_array): # FIXME knl = self.get_kernel() @@ -203,7 +198,9 @@ class P2PBase(KernelCacheMixin, KernelComputation): class P2P(P2PBase): """Direct applier for P2P interactions.""" - default_name = "p2p_apply" + @property + def default_name(self): + return "p2p_apply" def get_kernel(self): loopy_insns, result_names = self.get_loopy_insns_and_result_names() @@ -268,7 +265,9 @@ class P2P(P2PBase): class P2PMatrixGenerator(P2PBase): """Generator for P2P interaction matrix entries.""" - default_name = "p2p_matrix" + @property + def default_name(self): + return "p2p_matrix" def get_strength_or_not(self, isrc, kernel_idx): return 1 @@ -333,7 +332,9 @@ class P2PMatrixSubsetGenerator(P2PBase): .. automethod:: __call__ """ - default_name = "p2p_subset" + @property + def default_name(self): + return "p2p_subset" def get_strength_or_not(self, isrc, kernel_idx): return 1 @@ -438,7 +439,9 @@ class P2PMatrixSubsetGenerator(P2PBase): # {{{ P2P from CSR-like interaction list class P2PFromCSR(P2PBase): - default_name = "p2p_from_csr" + @property + def default_name(self): + return "p2p_from_csr" def get_kernel(self, max_nsources_in_one_box, max_ntargets_in_one_box, gpu=False, nsplit=32): diff --git a/sumpy/qbx.py b/sumpy/qbx.py index f5e505706ce2f80d375a030537a122eafb853aa5..a3042ef91fcfe323ea38713e2b0b6a322237c753 100644 --- a/sumpy/qbx.py +++ b/sumpy/qbx.py @@ -67,8 +67,6 @@ def stringify_expn_index(i): # {{{ base class class LayerPotentialBase(KernelCacheMixin, KernelComputation): - default_name = "qbx" - def __init__(self, ctx, expansion, strength_usage=None, value_dtypes=None, name=None, device=None, source_kernels=None, target_kernels=None): @@ -229,7 +227,9 @@ class LayerPotential(LayerPotentialBase): .. automethod:: __call__ """ - default_name = "qbx_apply" + @property + def default_name(self): + return "qbx_apply" @memoize_method def get_kernel(self): @@ -307,7 +307,9 @@ class LayerPotential(LayerPotentialBase): class LayerPotentialMatrixGenerator(LayerPotentialBase): """Generator for layer potential matrix entries.""" - default_name = "qbx_matrix" + @property + def default_name(self): + return "qbx_matrix" def get_strength_or_not(self, isrc, kernel_idx): return 1 @@ -376,7 +378,9 @@ class LayerPotentialMatrixSubsetGenerator(LayerPotentialBase): .. automethod:: __call__ """ - default_name = "qbx_subset" + @property + def default_name(self): + return "qbx_subset" def get_strength_or_not(self, isrc, kernel_idx): return 1 diff --git a/sumpy/tools.py b/sumpy/tools.py index 6505fa370e6b8d06565382ff365b7abc0f8aaace..f1ac7c1e0be0ec51ebd6a27e28def5a75e3bf172 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -43,8 +43,9 @@ import warnings import os import sys import enum -import platform -from collections import defaultdict, namedtuple +from abc import ABC, abstractmethod +from dataclasses import dataclass + from pymbolic.mapper import WalkMapper import pymbolic @@ -54,11 +55,14 @@ import pyopencl as cl import pyopencl.array as cla import loopy as lp -from typing import Dict, Tuple, Any, ClassVar +from typing import Dict, Tuple, Any, List, Optional, TYPE_CHECKING import logging logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from sumpy.kernel import Kernel + # {{{ multi_index helpers @@ -452,7 +456,9 @@ def diff_derivative_coeff_dict(derivative_coeff_dict: DerivativeCoeffDict, *derivative_coeff_dict* using the variable given by **variable_idx** and return a new derivative transformation dictionary. """ + from collections import defaultdict new_derivative_coeff_dict = defaultdict(lambda: 0) + for mi, coeff in derivative_coeff_dict.items(): # In the case where we have x * u.diff(x), the result should # be x.diff(x) + x * u.diff(x, x) @@ -576,18 +582,32 @@ class ScalingAssignmentTag(Tag): pass -class KernelComputation: - """Common input processing for kernel computations.""" +class KernelComputation(ABC): + """Common input processing for kernel computations. + + .. attribute:: name + .. attribute:: target_kernels + .. attribute:: source_kernels + .. attribute:: strength_usage - default_name: ClassVar[str] = "unknown" + .. automethod:: get_kernel + """ - def __init__(self, ctx, target_kernels, source_kernels, strength_usage, - value_dtypes, name, device=None): + def __init__(self, ctx: Any, + target_kernels: List["Kernel"], + source_kernels: List["Kernel"], + strength_usage: Optional[List[int]] = None, + value_dtypes: Optional[List["np.dtype"]] = None, + name: Optional[str] = None, + device: Optional[Any] = None) -> None: """ - :arg kernels: list of :class:`sumpy.kernel.Kernel` instances - :class:`sumpy.kernel.TargetDerivative` wrappers should be + :arg target_kernels: list of :class:`~sumpy.kernel.Kernel` instances, + with :class:`sumpy.kernel.DirectionalTargetDerivative` as the outermost kernel wrappers, if present. - :arg strength_usage: A list of integers indicating which expression + :arg source_kernels: list of :class:`~sumpy.kernel.Kernel` instances + with :class:`~sumpy.kernel.DirectionalSourceDerivative` as the + outermost kernel wrappers, if present. + :arg strength_usage: list of integers indicating which expression uses which density. This implicitly specifies the number of density arrays that need to be passed. Default: all kernels use the same density. @@ -634,6 +654,11 @@ class KernelComputation: self.name = name or self.default_name + @property + @abstractmethod + def default_name(self): + pass + def get_kernel_scaling_assignments(self): from sumpy.symbolic import SympyToPymbolicMapper sympy_conv = SympyToPymbolicMapper() @@ -648,6 +673,10 @@ class KernelComputation: for i, (kernel, dtype) in enumerate( zip(self.target_kernels, self.value_dtypes))] + @abstractmethod + def get_kernel(self): + pass + # }}} @@ -945,7 +974,10 @@ def to_complex_dtype(dtype): raise RuntimeError(f"Unknown dtype: {dtype}") -ProfileGetter = namedtuple("ProfileGetter", "start, end") +@dataclass(frozen=True) +class ProfileGetter: + start: int + end: int def get_native_event(evt): @@ -1173,6 +1205,7 @@ def _get_fft_backend(queue) -> FFTBackend: "Falling back to slower implementation.") return FFTBackend.loopy + import platform if (sys.platform == "darwin" and platform.machine() == "x86_64" and queue.context.devices[0].platform.name