Fix cache leaks for pe._cached_abstract_eval

The function `pe._cached_abstract_eval` uses `util.cache`,
while most other cached functions in JAX use `weakref_lru_cache`.
The main difference is that `util.cache` keeps strong references to the function
arguments.

The modified test `lax_control_flow_test::test_cond_memory_leak` (added
a jit for one of the branches) is failing without this fix.
This is because the Jaxpr including the closed-over constant
leaks due to the strong references kept by the `util.cache` used in
`pe._cached_abstract_eval`.

We cannot use directly `weakref_lru_cache` because it keeps
weak references only to the first positional argument. We add here
a variant `multi_weakref_lru_cache` that uses weak references
for all the positional and keyword arguments for which
`util.is_weakref_cache_key_type` is true.
This is set for `Jaxpr`, `ClosedJaxpr` and `Callable`.
Eventually, we can decide to generalize the existing
`weakref_lru_cache` to have this behavior.

Author: gnecula

jax/_src/core.py

@@ -51,7 +51,8 @@
 from jax._src.util import (safe_zip, safe_map, curry, tuple_insert,
                            tuple_delete, cache,
                            HashableFunction, HashableWrapper, weakref_lru_cache,
-                           partition_list, StrictABCMeta, foreach)
+                           partition_list, StrictABCMeta, foreach,
+                           weakref_cache_key_types)
 import jax._src.pretty_printer as pp
 from jax._src.named_sharding import NamedSharding
 from jax._src.lib import jax_jit
@@ -196,6 +197,7 @@ def replace(self, **kwargs):
       raise ValueError(f"Unknown keyword arguments: {kwargs}")
     return jaxpr
 
+weakref_cache_key_types.add(Jaxpr)
 
 def join_effects(*effects: Effects) -> Effects:
   return set().union(*effects) if effects else no_effects
@@ -288,6 +290,9 @@ def pretty_print(self, *, source_info=False, print_shapes=True,
   def _repr_pretty_(self, p, cycle):
     return p.text(self.pretty_print(use_color=True))
 
+weakref_cache_key_types.add(ClosedJaxpr)
+
+
 @curry
 def jaxpr_as_fun(closed_jaxpr: ClosedJaxpr, *args):
   # TODO(dougalm): remove this hack when we add contexts to jaxpr.

jax/_src/dispatch.py

@@ -93,6 +93,10 @@ def apply_primitive(prim, *args, **params):
     lib.jax_jit.swap_thread_local_state_disable_jit(prev)
   return outs
 
+# TODO(necula): this cache will contain strong references to all
+# Jaxprs on which it is used. This is not immediately fixable by using
+# util.multi_weakref_lru_cache, because the `params` (including the Jaxpr)
+# are closed over in the `prim_fun` lambda. Leaving this fix for a later PR.
 @util.cache()
 def xla_primitive_callable(prim: core.Primitive, **params):
   util.test_event("xla_primitive_callable_cache_miss")

jax/_src/interpreters/partial_eval.py

@@ -47,8 +47,9 @@
 from jax._src.tree_util import PyTreeDef, treedef_tuple, register_static
 from jax._src.util import (unzip2, safe_zip, safe_map, toposort, split_list,
                            merge_lists, partition_list, OrderedSet,
-                           as_hashable_function, weakref_lru_cache, subs_list,
-                           HashableFunction, foreach, cache)
+                           as_hashable_function, weakref_lru_cache,
+                           multi_weakref_lru_cache, subs_list,
+                           HashableFunction, foreach)
 
 
 map, unsafe_map = safe_map, map
@@ -1627,7 +1628,7 @@ def move_outvars_to_back(jaxpr: ClosedJaxpr, to_move: Sequence[bool]) -> ClosedJ
   return _move_outvars_to_back(jaxpr, tuple(to_move))
 
 @weakref_lru_cache
-def _move_outvars_to_back(jaxpr, to_move):
+def _move_outvars_to_back(jaxpr: core.ClosedJaxpr, to_move):
   new_outvars = ([e for e, m in zip(jaxpr.jaxpr.outvars, to_move) if not m] +
                  [e for e, m in zip(jaxpr.jaxpr.outvars, to_move) if     m])
   return jaxpr.replace(jaxpr=jaxpr.jaxpr.replace(outvars=new_outvars))
@@ -1879,7 +1880,7 @@ def vars(atom: Atom) -> list[Var]:
   return constvars, constvals
 
 
-@cache()
+@multi_weakref_lru_cache
 def _cached_abstract_eval(primitive: core.Primitive, *aval_qdds, **params):
   return primitive.abstract_eval(*aval_qdds, **params)
 
@@ -2729,7 +2730,7 @@ def _linearize_of_pmap_hack(f: lu.WrappedFun, jaxpr, consts) -> tuple[Jaxpr, lis
 
 
 @weakref_lru_cache
-def lower_jaxpr(hi_jaxpr):
+def lower_jaxpr(hi_jaxpr: core.ClosedJaxpr):
   lo_avals = [lo_ty for aval in hi_jaxpr.in_aval_qdds for lo_ty in aval.lo_ty()]
   f = lu.wrap_init(partial(lower_traceable, hi_jaxpr),
                    debug_info=hi_jaxpr.jaxpr.debug_info)

jax/_src/util.py

@@ -16,13 +16,14 @@
 
 import abc
 from collections.abc import Callable, Iterable, Iterator, Sequence
+import dataclasses
 import functools
 from functools import partial
 import itertools as it
 import logging
 import math
 import operator
-from typing import (Any, Generic, SupportsIndex, TypeVar, overload, TYPE_CHECKING, cast)
+from typing import (Any, Generic, SupportsIndex, Type, TypeVar, overload, TYPE_CHECKING, cast)
 import weakref
 
 import numpy as np
@@ -331,8 +332,8 @@ def weakref_lru_cache(call: Callable, maxsize=2048,
   Least recently used cache decorator with weakref support.
 
   The cache will take a weakref to the first argument of the wrapped function
-  and strong refs to all subsequent operations. In all other respects it should
-  behave similar to `functools.lru_cache`.
+  and strong refs to all other arguments. In all other respects it should
+  behave similar to `functools.lru_cache`. The cache is thread local.
   """
   cached_call = _weakref_lru_cache.weakref_lru_cache(
       config.trace_context if trace_context_in_key else _ignore, call, maxsize
@@ -341,6 +342,132 @@ def weakref_lru_cache(call: Callable, maxsize=2048,
   return cached_call
 
 
+@dataclasses.dataclass(frozen=True, slots=True, weakref_slot=True)
+class MultiWeakRefCacheKey:
+  weakrefs: tuple[weakref.ref, ...]  # Used only when len(weakrefs) >= 2
+
+
+class MultiWeakRefPlaceholder:
+  # Stands for an arg/kwarg that was replaced with a weakref
+  pass
+_multi_weakref_placeholder = MultiWeakRefPlaceholder()
+
+# The types of arguments for which `multi_weakref_lru_cache` should keep
+# weak references.
+weakref_cache_key_types: set[Type] = set()
+def is_weakref_cache_key_type(v):
+  return callable(v) or (type(v) in weakref_cache_key_types)
+
+
+def multi_weakref_lru_cache(
+      call: Callable, *,
+      maxsize=2048,
+      trace_context_in_key: bool = True):
+    """
+    Least recently used cache decorator with weakref support.
+
+    Similar to `weakref_lru_cache`, except that it keeps weak references
+    to all positional and keyword arguments for which
+    `is_weakref_cache_key_type()` is true, and strong references to
+    other arguments. The cache entry is removed if any of the weakref
+    arguments dies.
+    """
+    # Keep strong references to the MultiWeakRefCacheKeys that resulted in
+    # cache misses, and are cache keys. Indexed by id. Only keys with all
+    # included weakrefs live are present.
+    id_to_key: dict[int, MultiWeakRefCacheKey] = {}
+    # For each `wr: weakref.ref` present in `key: MultiWeakRefCacheKey` we have
+    # `id(key) in weakref_to_key_ids[wr]`.
+    weakref_to_key_ids: dict[weakref.ref, set[int]] = {}
+
+    def remove_weakref(wr: weakref.ref):
+      key_ids = weakref_to_key_ids.get(wr, set())
+      for key_id in key_ids:
+        try:
+          del id_to_key[key_id]
+        except KeyError:
+          pass
+      try:
+        del weakref_to_key_ids[wr]
+      except KeyError:
+        pass
+
+    def weakrefs_to_sentinel(v, acc: list[Any]):
+      if isinstance(v, tuple):
+        return tuple(weakrefs_to_sentinel(v1, acc) for v1 in v)
+      elif isinstance(v, dict):
+        return {k: weakrefs_to_sentinel(v1, acc) for k, v1 in v.items()}
+      elif is_weakref_cache_key_type(v):
+        acc.append(v)
+        return _multi_weakref_placeholder
+      else:
+        return v
+
+    def sentinel_to_referrents(v,
+                               it: Iterator[weakref.ref],
+                               key_id: int | None):
+      # key_id is not None iff we use a MultiWeakRefCacheKey (>= 2 weakrefs)
+      if isinstance(v, tuple):
+        return tuple(sentinel_to_referrents(v1, it, key_id) for v1 in v)
+      elif isinstance(v, dict):
+        return {k: sentinel_to_referrents(v1, it, key_id)
+                for k, v1 in v.items()}
+      elif v is _multi_weakref_placeholder:
+        wr = next(it)
+        if key_id is not None:
+          weakref_to_key_ids.setdefault(wr, set()).add(key_id)
+        return wr()
+      else:
+        return v
+
+    def cache_miss(key: MultiWeakRefCacheKey | MultiWeakRefPlaceholder | Any,
+                   *args, **kwargs):
+      if isinstance(key, MultiWeakRefCacheKey):  # had at least 2 weakrefs
+        # We know `key` is in `cached_call` cache, so store strong references
+        key_id = id(key)
+        id_to_key[key_id] = key
+        orig_args, orig_kwargs = sentinel_to_referrents(
+            (args, kwargs), iter(key.weakrefs), key_id)
+      elif key is _multi_weakref_placeholder:  # had 0 weakrefs
+        orig_args = args
+        orig_kwargs = kwargs
+      else:  # had 1 weakref, we had put it first as the `key`
+        orig_args, orig_kwargs = sentinel_to_referrents(
+            (args, kwargs), iter([weakref.ref(key)]), None)
+      return call(*orig_args, **orig_kwargs)
+
+
+    cached_call = _weakref_lru_cache.weakref_lru_cache(
+        config.trace_context if trace_context_in_key else _ignore,
+        cache_miss, maxsize
+    )
+    register_cache(cached_call, str(call))
+
+    @functools.wraps(call)
+    def wrapper(*args, **kwargs):
+      acc_weakrefs: list[Any] = []
+      args, kwargs = weakrefs_to_sentinel((args, kwargs), acc_weakrefs)
+      nr_weakrefs = len(acc_weakrefs)
+      if nr_weakrefs == 0:
+        return cached_call(_multi_weakref_placeholder, *args, **kwargs)
+      elif nr_weakrefs == 1:
+        # Put the single weakref first, and skip the MultiWeakRefCacheKey
+        return cached_call(acc_weakrefs[0], *args, **kwargs)
+      else:
+        value_to_weakref = {v: weakref.ref(v, remove_weakref)
+                            for v in set(acc_weakrefs)}
+        key = MultiWeakRefCacheKey(weakrefs=tuple(value_to_weakref[v]
+                                                  for v in acc_weakrefs))
+        return cached_call(key, *args, **kwargs)
+
+    wrapper.cache_info = cached_call.cache_info
+    wrapper.cache_clear = cached_call.cache_clear
+    wrapper.cache_keys = cached_call.cache_keys
+    wrapper._multi_weakref_id_to_key = id_to_key  # stays alive as long as wrapper
+    wrapper._multi_weakref_to_key_ids = weakref_to_key_ids
+    return wrapper
+
+
 class Unhashable:
   __slots__ = ["val"]
 

tests/lax_control_flow_test.py

@@ -15,6 +15,7 @@
 
 import collections
 import contextlib
+import gc
 from functools import partial
 import itertools
 import math
@@ -3162,18 +3163,21 @@ def test_cond_casting(self):
   @jtu.thread_unsafe_test()  # live_arrays count isn't thread-safe
   def test_cond_memory_leak(self):
     # https://github.com/jax-ml/jax/issues/12719
-
     def leak():
       data = jax.device_put(np.zeros((1024), dtype=np.float32) + 1)
       def g():
         return jax.lax.cond(
               True,
-              lambda: data[0],  # noqa: F821
+              jax.jit(lambda: data[0]),  # noqa: F821
               lambda: data[1],  # noqa: F821
           )
+      # _ = g()  # TODO(necula): enable this, requires fixing leaks in the
+      # caching of dispatch.xla_primitive_callable.
       jg = jax.jit(g)
       _ = jg().block_until_ready()
+      jg.clear_cache()
       del g, jg, data, _
+      gc.collect()
 
     nbufs = lambda: len(jax.live_arrays())
     base = nbufs()

tests/util_test.py

@@ -11,11 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import dataclasses
+import random
 from functools import partial
+import gc
 import operator
+import threading
 
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 import jax
 from jax import api_util
 from jax._src import linear_util as lu
@@ -155,6 +158,106 @@ def reference_loop_generator(x):
     for _ in range(4097):
       reference_loop_generator(lambda x: x)
 
+  @parameterized.named_parameters(
+      dict(weakref_count=weakref_count,
+           testcase_name=f"_{weakref_count=}")
+      for weakref_count in [0, 1, 3])
+  def test_multi_weak_ref_cache(self, *, weakref_count=0):
+
+    class Key:  # hashed by id
+      def __init__(self, x):
+        self.x = x
+
+    if weakref_count > 0:
+      util.weakref_cache_key_types.add(Key)
+
+    @partial(util.multi_weakref_lru_cache, trace_context_in_key=False)
+    def myfun(a, k1, *, k2, k3):
+      return f"{a=}, {k1=}, {k2=}, {k3=}"
+
+    def check_invariant(expected_live_keys: int):
+      self.assertLen(myfun._multi_weakref_id_to_key, expected_live_keys if weakref_count > 1 else 0)
+      for key_id, key in myfun._multi_weakref_id_to_key.items():
+        for wr in key.weakrefs:
+          self.assertIn(wr, myfun._multi_weakref_to_key_ids)
+          self.assertIn(key_id, myfun._multi_weakref_to_key_ids[wr])
+
+    k1 = Key(1)
+    k3 = (k1, k1) if weakref_count > 1 else 4
+    util.clear_all_caches()
+    r1 = myfun(2, k1, k2=3, k3=k3)  # miss
+    c1 = myfun.cache_info()
+    self.assertEqual((0, 1, 1), (c1.hits, c1.misses, c1.currsize))
+    check_invariant(1)
+
+    for i in range(10):
+      r2 = myfun(2, k1, k2=3, k3=k3)  # all hits
+      self.assertIs(r1, r2)
+      c2 = myfun.cache_info()
+      self.assertEqual((1 + i, 1, 1), (c2.hits, c2.misses, c2.currsize))
+      check_invariant(1)
+
+    del k1, k3  # expect that the cache entries are removed (if weakref_count > 0)
+    gc.collect()
+    c3 = myfun.cache_info()
+    self.assertEqual(c3.currsize, 0 if weakref_count > 0 else 1)
+    check_invariant(0)
+
+    k1_2 = Key(2)
+    k3_2 = (Key(3), Key(3)) if weakref_count > 1 else (3, 3)
+    r4 = myfun(2, k1_2, k2=3, k3=k3_2)  # miss
+    c4 = myfun.cache_info()
+    self.assertEqual((10, 2, (1 if weakref_count > 0 else 2)), (c4.hits, c4.misses, c4.currsize))
+    check_invariant(1)
+
+    if weakref_count > 1:
+      del k3_2  # clear the cache entry
+      gc.collect()
+      c5 = myfun.cache_info()
+      self.assertEqual((10, 2, 0), (c5.hits, c5.misses, c5.currsize))
+      check_invariant(0)
+
+      k3_3 = (Key(3), Key(3))
+      r6 = myfun(2, k1_2, k2=3, k3=k3_3)  # miss because Key hashed by it
+      self.assertIsNot(r4, r6)
+      c6 = myfun.cache_info()
+      self.assertEqual((10, 3, 1), (c6.hits, c6.misses, c6.currsize))
+      check_invariant(1)
+
+    del k1_2
+    gc.collect()
+    c7 = myfun.cache_info()
+    self.assertEqual(0 if weakref_count > 0 else 2, c7.currsize )
+    check_invariant(0)
+
+  def test_multi_weakref_lru_cache_threads(self):
+    num_workers = 5
+    num_live_keys_per_worker = 16
+    size_key_space = 32
+    @dataclasses.dataclass(frozen=True)
+    class WRKey:
+      f: int
+
+    util.weakref_cache_key_types.add(WRKey)
+
+    @partial(util.multi_weakref_lru_cache, maxsize=size_key_space // 2)
+    def myfun(k: WRKey):
+      return None
+
+    def Worker():
+      keys = [None] * num_live_keys_per_worker  # These are the live keys for this worker
+      for i in range(1000):
+        key_idx = random.randint(0, num_live_keys_per_worker - 1)
+        key = WRKey(random.randint(0, size_key_space))
+        myfun(key)
+        keys[key_idx] = key  # Kill some previous key and keep this live
+
+    workers = [threading.Thread(target=Worker()) for _ in range(num_workers)]
+    for t in workers:
+      t.start()
+    for t in workers:
+      t.join()
+
 
 class SafeMapTest(jtu.JaxTestCase):