Make use of loopy.TranslationUnit.executor

This avoids long-lived references to CL kernels held by loopy caches

Author: inducer

sumpy/e2e.py

@@ -82,7 +82,7 @@ def __init__(self, ctx, src_expansion, tgt_expansion,
                     SourceTransformationRemover()(
                         TargetTransformationRemover()(tgt_expansion.kernel)))
 
-        self.ctx = ctx
+        self.context = ctx
         self.src_expansion = src_expansion
         self.tgt_expansion = tgt_expansion
         self.name = name or self.default_name
@@ -297,7 +297,7 @@ def __call__(self, queue, **kwargs):
         src_rscale = centers.dtype.type(kwargs.pop("src_rscale"))
         tgt_rscale = centers.dtype.type(kwargs.pop("tgt_rscale"))
 
-        knl = self.get_cached_optimized_kernel()
+        knl = self.get_cached_kernel_executor()
 
         return knl(queue,
                 centers=centers,
@@ -537,7 +537,7 @@ def __call__(self, queue, **kwargs):
         tgt_rscale = centers.dtype.type(kwargs.pop("tgt_rscale"))
         src_expansions = kwargs.pop("src_expansions")
 
-        knl = self.get_cached_optimized_kernel(result_dtype=src_expansions.dtype)
+        knl = self.get_cached_kernel_executor(result_dtype=src_expansions.dtype)
 
         return knl(queue,
                 src_expansions=src_expansions,
@@ -647,7 +647,7 @@ def __call__(self, queue, **kwargs):
                 "m2l_translation_classes_dependent_data")
         result_dtype = m2l_translation_classes_dependent_data.dtype
 
-        knl = self.get_cached_optimized_kernel(result_dtype=result_dtype)
+        knl = self.get_cached_kernel_executor(result_dtype=result_dtype)
 
         return knl(queue,
                 src_rscale=src_rscale,
@@ -741,7 +741,7 @@ def __call__(self, queue, **kwargs):
         """
         preprocessed_src_expansions = kwargs.pop("preprocessed_src_expansions")
         result_dtype = preprocessed_src_expansions.dtype
-        knl = self.get_cached_optimized_kernel(result_dtype=result_dtype)
+        knl = self.get_cached_kernel_executor(result_dtype=result_dtype)
 
         return knl(queue,
                 preprocessed_src_expansions=preprocessed_src_expansions, **kwargs)
@@ -840,7 +840,7 @@ def __call__(self, queue, **kwargs):
         """
         tgt_expansions = kwargs.pop("tgt_expansions")
         result_dtype = tgt_expansions.dtype
-        knl = self.get_cached_optimized_kernel(result_dtype=result_dtype)
+        knl = self.get_cached_kernel_executor(result_dtype=result_dtype)
 
         return knl(queue, tgt_expansions=tgt_expansions, **kwargs)
 
@@ -950,7 +950,7 @@ def __call__(self, queue, **kwargs):
         :arg tgt_rscale:
         :arg centers:
         """
-        knl = self.get_cached_optimized_kernel()
+        knl = self.get_cached_kernel_executor()
 
         centers = kwargs.pop("centers")
         # "1" may be passed for rscale, which won't have its type
@@ -1054,7 +1054,7 @@ def __call__(self, queue, **kwargs):
         :arg tgt_rscale:
         :arg centers:
         """
-        knl = self.get_cached_optimized_kernel()
+        knl = self.get_cached_kernel_executor()
 
         centers = kwargs.pop("centers")
         # "1" may be passed for rscale, which won't have its type

sumpy/e2p.py

@@ -68,7 +68,7 @@ def __init__(self, ctx, expansion, kernels,
         for knl in kernels:
             assert txr(knl) == expansion.kernel
 
-        self.ctx = ctx
+        self.context = ctx
         self.expansion = expansion
         self.kernels = kernels
         self.name = name or self.default_name
@@ -210,7 +210,7 @@ def __call__(self, queue, **kwargs):
         :arg centers:
         :arg targets:
         """
-        knl = self.get_cached_optimized_kernel()
+        knl = self.get_cached_kernel_executor()
 
         centers = kwargs.pop("centers")
         # "1" may be passed for rscale, which won't have its type
@@ -327,7 +327,7 @@ def get_optimized_kernel(self):
         return knl
 
     def __call__(self, queue, **kwargs):
-        knl = self.get_cached_optimized_kernel()
+        knl = self.get_cached_kernel_executor()
 
         centers = kwargs.pop("centers")
         # "1" may be passed for rscale, which won't have its type

sumpy/p2e.py

@@ -124,7 +124,7 @@ def __call__(self, queue, **kwargs):
         from sumpy.tools import is_obj_array_like
         sources = kwargs.pop("sources")
         centers = kwargs.pop("centers")
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 sources_is_obj_array=is_obj_array_like(sources),
                 centers_is_obj_array=is_obj_array_like(centers))
 

sumpy/p2p.py

@@ -256,7 +256,7 @@ def get_kernel(self):
         return loopy_knl
 
     def __call__(self, queue, targets, sources, strength, **kwargs):
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 targets_is_obj_array=is_obj_array_like(targets),
                 sources_is_obj_array=is_obj_array_like(sources))
 
@@ -318,7 +318,7 @@ def get_kernel(self):
         return loopy_knl
 
     def __call__(self, queue, targets, sources, **kwargs):
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 targets_is_obj_array=is_obj_array_like(targets),
                 sources_is_obj_array=is_obj_array_like(sources))
 
@@ -429,7 +429,7 @@ def __call__(self, queue, targets, sources, tgtindices, srcindices, **kwargs):
         :returns: a one-dimensional array of interactions, for each index pair
             in (*srcindices*, *tgtindices*)
         """
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 targets_is_obj_array=is_obj_array_like(targets),
                 sources_is_obj_array=is_obj_array_like(sources))
 
@@ -731,7 +731,7 @@ def __call__(self, queue, **kwargs):
         else:
             dtype_size = None
 
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 max_nsources_in_one_box=max_nsources_in_one_box,
                 max_ntargets_in_one_box=max_ntargets_in_one_box,
                 dtype_size=dtype_size,

sumpy/qbx.py

@@ -288,7 +288,7 @@ def __call__(self, queue, targets, sources, centers, strengths, expansion_radii,
             already multiplied in.
         """
 
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 targets_is_obj_array=is_obj_array_like(targets),
                 sources_is_obj_array=is_obj_array_like(sources),
                 centers_is_obj_array=is_obj_array_like(centers))
@@ -359,7 +359,7 @@ def get_kernel(self):
         return loopy_knl
 
     def __call__(self, queue, targets, sources, centers, expansion_radii, **kwargs):
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 targets_is_obj_array=is_obj_array_like(targets),
                 sources_is_obj_array=is_obj_array_like(sources),
                 centers_is_obj_array=is_obj_array_like(centers))
@@ -479,7 +479,7 @@ def __call__(self, queue, targets, sources, centers, expansion_radii,
             in (*srcindices*, *tgtindices*)
         """
 
-        knl = self.get_cached_optimized_kernel(
+        knl = self.get_cached_kernel_executor(
                 targets_is_obj_array=is_obj_array_like(targets),
                 sources_is_obj_array=is_obj_array_like(sources),
                 centers_is_obj_array=is_obj_array_like(centers))

sumpy/tools.py

@@ -381,8 +381,17 @@ def __eq__(self, other):
 
 
 class KernelCacheMixin:
-    @memoize_method
     def get_cached_optimized_kernel(self, **kwargs):
+        from warnings import warn
+        warn("get_cached_optimized_kernel is deprecated. "
+             "Use get_cached_kernel_executor instead. "
+             "This will stop working in October 2023.",
+             DeprecationWarning, stacklevel=2)
+
+        return self.get_cached_kernel_executor(**kwargs)
+
+    @memoize_method
+    def get_cached_kernel_executor(self, **kwargs) -> lp.ExecutorBase:
         from sumpy import (code_cache, CACHING_ENABLED, OPT_ENABLED,
             NO_CACHE_KERNELS)
 
@@ -401,7 +410,7 @@ def get_cached_optimized_kernel(self, **kwargs):
                 result = code_cache[cache_key]
                 logger.debug("{}: kernel cache hit [key={}]".format(
                     self.name, cache_key))
-                return result
+                return result.executor(self.context)
             except KeyError:
                 pass
 
@@ -422,7 +431,7 @@ def get_cached_optimized_kernel(self, **kwargs):
                 NO_CACHE_KERNELS and self.name in NO_CACHE_KERNELS):
             code_cache.store_if_not_present(cache_key, knl)
 
-        return knl
+        return knl.executor(self.context)
 
     @staticmethod
     def _allow_redundant_execution_of_knl_scaling(knl):