From fa24fef1af53268077cbfeda69c2545330535631 Mon Sep 17 00:00:00 2001
From: Isuru Fernando <isuruf@gmail.com>
Date: Thu, 16 Feb 2023 20:38:09 -0600
Subject: [PATCH] Introduce SUMPY_NO_CACHE_KERNELS to selectively disable
 caches (#159)

---
 sumpy/__init__.py | 24 +++++++++++++++---------
 sumpy/tools.py    | 12 ++++++++----
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/sumpy/__init__.py b/sumpy/__init__.py
index b39ead2e..bfcbf6d7 100644
--- a/sumpy/__init__.py
+++ b/sumpy/__init__.py
@@ -67,12 +67,16 @@ CACHING_ENABLED = (
     "SUMPY_NO_CACHE" not in os.environ
     and "CG_NO_CACHE" not in os.environ)
 
+NO_CACHE_KERNELS = tuple(os.environ.get("SUMPY_NO_CACHE_KERNELS",
+                                        "").split(","))
 
-def set_caching_enabled(flag):
+
+def set_caching_enabled(flag, no_cache_kernels=()):
     """Set whether :mod:`loopy` is allowed to use disk caching for its various
     code generation stages.
     """
-    global CACHING_ENABLED
+    global CACHING_ENABLED, NO_CACHE_KERNELS
+    NO_CACHE_KERNELS = no_cache_kernels
     CACHING_ENABLED = flag
 
 
@@ -81,17 +85,19 @@ class CacheMode:
     disk caches.
     """
 
-    def __init__(self, new_flag):
+    def __init__(self, new_flag, new_no_cache_kernels=()):
         self.new_flag = new_flag
+        self.new_no_cache_kernels = new_no_cache_kernels
 
     def __enter__(self):
-        global CACHING_ENABLED
-        self.previous_mode = CACHING_ENABLED
-        CACHING_ENABLED = self.new_flag
+        global CACHING_ENABLED, NO_CACHE_KERNELS
+        self.previous_flag = CACHING_ENABLED
+        self.previous_kernels = NO_CACHE_KERNELS
+        set_caching_enabled(self.new_flag, self.new_no_cache_kernels)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        global CACHING_ENABLED
-        CACHING_ENABLED = self.previous_mode
-        del self.previous_mode
+        set_caching_enabled(self.previous_flag, self.previous_kernels)
+        del self.previous_flag
+        del self.previous_kernels
 
 # }}}
diff --git a/sumpy/tools.py b/sumpy/tools.py
index f9abe6b7..85b1f1c6 100644
--- a/sumpy/tools.py
+++ b/sumpy/tools.py
@@ -753,9 +753,11 @@ class OrderedSet(MutableSet):
 class KernelCacheMixin:
     @memoize_method
     def get_cached_optimized_kernel(self, **kwargs):
-        from sumpy import code_cache, CACHING_ENABLED, OPT_ENABLED
+        from sumpy import (code_cache, CACHING_ENABLED, OPT_ENABLED,
+            NO_CACHE_KERNELS)
 
-        if CACHING_ENABLED:
+        if CACHING_ENABLED and not (
+                NO_CACHE_KERNELS and self.name in NO_CACHE_KERNELS):
             import loopy.version
             from sumpy.version import KERNEL_VERSION
             cache_key = (
@@ -774,7 +776,8 @@ class KernelCacheMixin:
                 pass
 
         logger.info("%s: kernel cache miss", self.name)
-        if CACHING_ENABLED:
+        if CACHING_ENABLED and not (
+                NO_CACHE_KERNELS and self.name in NO_CACHE_KERNELS):
             logger.info("{}: kernel cache miss [key={}]".format(
                 self.name, cache_key))
 
@@ -785,7 +788,8 @@ class KernelCacheMixin:
             else:
                 knl = self.get_kernel()
 
-        if CACHING_ENABLED:
+        if CACHING_ENABLED and not (
+                NO_CACHE_KERNELS and self.name in NO_CACHE_KERNELS):
             code_cache.store_if_not_present(cache_key, knl)
 
         return knl
-- 
GitLab