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,11 @@ 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 in `params` (for higher-order primitives).
+# 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,135 @@ 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 type(v) is tuple:
+        return tuple(weakrefs_to_sentinel(v1, acc) for v1 in v)
+      elif type(v) is 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 type(v) is tuple:
+        return tuple(sentinel_to_referrents(v1, it, key_id) for v1 in v)
+      elif type(v) is 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(*orig_args, **orig_kwargs):
+      acc_weakrefs: list[Any] = []
+      args, kwargs = weakrefs_to_sentinel((orig_args, orig_kwargs),
+                                          acc_weakrefs)
+      nr_weakrefs = len(acc_weakrefs)
+      if nr_weakrefs == 0:
+        return cached_call(_multi_weakref_placeholder,
+                           *orig_args, **orig_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()