From ef652e8235482936a74c0674c919b48900012dab Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Fri, 9 Oct 2026 15:32:49 +0100 Subject: [PATCH 1/2] compiler: FindSymbols caching and unsorted mode --- devito/ir/iet/nodes.py | 2 +- devito/ir/iet/utils.py | 2 +- devito/ir/iet/visitors.py | 46 +++++++++++++- devito/operator/operator.py | 3 +- devito/passes/iet/engine.py | 32 ++++++++-- tests/test_iet.py | 123 ++++++++++++++++++++++++++++++++++-- 6 files changed, 192 insertions(+), 16 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index d3ff53bfd6..d3b8775e31 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -503,7 +503,7 @@ def is_initializable(self): def defines(self): return (self.output.base,) if self.is_initializable else () - @property + @cached_property def expr_symbols(self): return tuple(self.expr.free_symbols) diff --git a/devito/ir/iet/utils.py b/devito/ir/iet/utils.py index 5da6ae075c..b5f36b4799 100644 --- a/devito/ir/iet/utils.py +++ b/devito/ir/iet/utils.py @@ -121,7 +121,7 @@ def derive_parameters(iet, drop_locals=False, ordering='default'): candidates = filter_ordered(candidates) # Filter off symbols which are defined somewhere within `iet` - defines = [s.name for s in FindSymbols('defines').visit(iet)] + defines = {s.name for s in FindSymbols('defines', ordered=False).visit(iet)} parameters = [s for s in candidates if s.name not in defines] # Drop globally-visible objects diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 6186ff4e76..05307261c8 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -7,7 +7,7 @@ import ctypes from collections import OrderedDict from collections.abc import Callable, Generator, Iterable, Iterator, Sequence -from contextlib import suppress +from contextlib import contextmanager, suppress from itertools import chain, groupby from typing import Any, Generic, TypeVar @@ -1090,6 +1090,10 @@ class FindSymbols(LazyVisitor[Any, list[Any], None]): - `indexedbases`: Collect all IndexedBase objects - `defines`: Collect all defined objects - `defines-aliases`: Collect all defined objects and their aliases + ordered : bool, optional + If True (default), sort the symbols naturally by name. Otherwise, + return them in traversal order, which is cheaper when only membership + matters. """ @staticmethod @@ -1117,10 +1121,30 @@ def _defines_aliases(n): 'defines-aliases': _defines_aliases } - def __init__(self, mode: str = 'symbolics') -> None: + # Sort keys cached while an Operator is built, see `caching_sort_keys` + _sort_keys: dict[int, tuple[Any, Any]] | None = None + + @classmethod + @contextmanager + def caching_sort_keys(cls): + """ + Cache the sort keys within this context, as the same symbols are sorted + over and over while an Operator is built. The cache is keyed by identity, + since symbols may compare equal yet print differently, and retains each + symbol so its id cannot be reused. It is released on exit. + """ + prev = cls._sort_keys + cls._sort_keys = {} if prev is None else prev + try: + yield + finally: + cls._sort_keys = prev + + def __init__(self, mode: str = 'symbolics', ordered: bool = True) -> None: super().__init__() self.mode = mode + self.ordered = ordered modes = mode.split('|') if len(modes) == 1: self.rule = self.rules[mode] @@ -1128,7 +1152,23 @@ def __init__(self, mode: str = 'symbolics') -> None: self.rule = lambda n: chain(*[self.rules[mode](n) for mode in modes]) def _post_visit(self, ret): - return sorted(filter_ordered(ret, key=id), key=natural_sort_key) + found = filter_ordered(ret, key=id) + if not self.ordered: + return found + + cache = self._sort_keys + if cache is None: + return sorted(found, key=natural_sort_key) + + def key(i): + try: + return cache[id(i)][1] + except KeyError: + v = natural_sort_key(i) + cache[id(i)] = (i, v) + return v + + return sorted(found, key=key) def visit_Node(self, o: Node) -> Iterator[Any]: yield from self._visit(o.children) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index b3690caf74..50203f6828 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -186,7 +186,8 @@ def __new__(cls, expressions, **kwargs): # Lower to a JIT-compilable object with timed_region('op-compile') as r: try: - op = cls._build(expressions, **kwargs) + with FindSymbols.caching_sort_keys(): + op = cls._build(expressions, **kwargs) finally: CacheInstances.clear_caches() memoized_func.clear_build_caches() diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index fe9a78b65b..1ecdd897e9 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -723,6 +723,26 @@ def _(i, mapper, sregistry): mapper[i] = i._rebuild(name=sregistry.make_name(prefix='nthreads')) +def _contains(found): + """ + Membership test over `found` with list semantics. A set is used as a fast + path, but some symbols (e.g. `DataSymbol`s) compare structurally while + hashing by identity, so misses fall back to the list. + """ + try: + lookup = set(found) + except TypeError: + lookup = set() + + def contains(a): + with suppress(TypeError): + if a in lookup: + return True + return a in found + + return contains + + def _update_args_efunc(root, efuncs, dag): """ Re-derive the parameters of `root` and apply the changes in cascade through @@ -763,19 +783,19 @@ def _update_args_efunc(root, efuncs, dag): # 2) defined a symbol for which no definition was available yet (e.g. # via a malloc, or a Dereference) - defines = FindSymbols('defines').visit(root.body) - drop_params = [a for a in root.parameters if a in defines] + in_defines = _contains(FindSymbols('defines', ordered=False).visit(root.body)) + drop_params = [a for a in root.parameters if in_defines(a)] # 3) removed a symbol that was previously necessary (e.g., `x_size` after # linearization) - symbols = FindSymbols('basics').visit(root.body) + in_symbols = _contains(FindSymbols('basics', ordered=False).visit(root.body)) drop_params.extend(a for a in root.parameters - if (a.is_Symbol or a.is_LocalObject) and a not in symbols) + if (a.is_Symbol or a.is_LocalObject) and not in_symbols(a)) # 4) removed a function that was previously necessary - functions = FindSymbols('symbolics').visit(root.body) + in_functions = _contains(FindSymbols('symbolics', ordered=False).visit(root.body)) drop_params.extend(a for a in root.parameters - if a.is_AbstractFunction and a not in functions) + if a.is_AbstractFunction and not in_functions(a)) # Must record the index, not the param itself, since a param may be # bound to whatever arg, possibly a generic SymPy expr diff --git a/tests/test_iet.py b/tests/test_iet.py index c2d096918f..5d9af9b2de 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -1,3 +1,5 @@ +import pickle +from copy import deepcopy from ctypes import POINTER, c_void_p import cgen @@ -6,7 +8,7 @@ import sympy from devito import ( # noqa - Dimension, Eq, Function, Grid, Operator, TimeFunction, switchconfig + Constant, Dimension, Eq, Function, Grid, Operator, TimeFunction, switchconfig ) from devito.ir import SymbolRegistry from devito.ir.iet import ( @@ -16,16 +18,18 @@ retrieve_iteration_tree ) from devito.ir.iet.visitors import sorted_efuncs -from devito.passes.iet.engine import Graph, reuse_efuncs +from devito.passes.iet.engine import ( + Graph, _contains, _update_args_efunc, create_call_graph, reuse_efuncs +) from devito.passes.iet.languages.C import CDataManager from devito.symbolics import ( FLOAT, Byref, Class, FieldFromComposite, InlineIf, ListInitializer, Macro, SizeOf, String ) from devito.symbolics.extended_dtypes import c_complex -from devito.tools import CustomDtype, as_tuple, dtype_to_ctype +from devito.tools import CustomDtype, as_hashable, as_tuple, dtype_to_ctype from devito.types import ( - Array, Bundle, CustomDimension, LocalObject, Pointer, Symbol, Temp + Array, Bundle, CustomDimension, LocalObject, Object, Pointer, Symbol, Temp ) from devito.types.misc import FunctionMap @@ -154,6 +158,10 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(expected) + # Same objects when unordered + unordered = FindSymbols(mode, ordered=False).visit(call) + assert {id(i) for i in unordered} == {id(i) for i in found} + def test_find_symbols_natural_numbering(): temporaries = [Temp(name=f'r{i}') for i in range(12)] @@ -163,6 +171,113 @@ def test_find_symbols_natural_numbering(): assert [i.name for i in found] == [i.name for i in temporaries] + # Unordered results retain the traversal order + found = FindSymbols('basics', ordered=False).visit(call) + + assert [i.name for i in found] == [i.name for i in temporaries[::-1]] + + +def test_find_symbols_cached_sort_keys(): + temporaries = [Temp(name=f'r{i}') for i in range(12)] + call = Call('foo', temporaries[::-1]) + + with FindSymbols.caching_sort_keys(): + for _ in range(2): + found = FindSymbols('basics').visit(call) + assert [i.name for i in found] == [i.name for i in temporaries] + assert len(FindSymbols._sort_keys) == len(temporaries) + + # Released on exit, and never retained beyond an Operator build + assert FindSymbols._sort_keys is None + + grid = Grid(shape=(4, 4)) + u = TimeFunction(name='u', grid=grid) + op = Operator(Eq(u.forward, u + 1)) + assert FindSymbols._sort_keys is None + assert op.dimensions + assert FindSymbols._sort_keys is None + + +def test_find_symbols_cached_sort_keys_identity(): + """ + Symbols comparing and hashing equal but printing differently keep their + own sort keys. + """ + class Sym: + + def __init__(self, name): + self.name = name + + def __eq__(self, other): + return isinstance(other, Sym) + + def __hash__(self): + return 0 + + def __str__(self): + return self.name + + a, b = Sym('b'), Sym('a') + + with FindSymbols.caching_sort_keys(): + for _ in range(2): + assert FindSymbols()._post_visit([a, b]) == [b, a] + + +def _update_args(root): + efuncs = {root.name: root} + dag = create_call_graph(root.name, as_hashable(efuncs)) + return _update_args_efunc(root, efuncs, dag)[root.name] + + +def test_update_args_equal_symbols(): + """ + Symbols such as Constants and Objects compare structurally but hash by + identity, so parameter reconciliation must retain list-membership semantics. + """ + # A formal parameter equal to a used symbol is retained, along with its value + q0 = Constant(name='q', value=1.) + q1 = Constant(name='q', value=2.) + root = Callable('foo', DummyExpr(Symbol(name='a'), q1 + 1, init=True), 'void', + parameters=[q0]) + + parameters = _update_args(root).parameters + assert len(parameters) == 1 + assert parameters[0] is q0 and parameters[0].data == 1. + + # A formal parameter equal to a defined symbol is dropped + o0 = Object(name='o', dtype=c_void_p) + o1 = Object(name='o', dtype=c_void_p) + root = Callable('bar', List(body=[Definition(o1)]), 'void', parameters=[o0]) + + assert _update_args(root).parameters == () + + +def test_contains_unhashable(): + in_found = _contains([Symbol(name='a'), [1]]) + + assert in_found(Symbol(name='a')) + assert in_found([1]) + assert not in_found(Symbol(name='b')) + + +def test_expression_expr_symbols_cached(): + a = Symbol(name='a') + b = Symbol(name='b') + expr = DummyExpr(a, b + 1, init=True) + + assert 'expr_symbols' not in expr.__dict__ + assert set(expr.expr_symbols) == {a, b} + assert 'expr_symbols' in expr.__dict__ + + # Rebuilding with a new expression does not reuse the cached symbols + c = Symbol(name='c') + assert set(expr._rebuild(expr=Eq(a, c)).expr_symbols) == {a, c} + + # NOTE: `free_symbols` is a set, so only the contents are guaranteed to match + for i in (pickle.loads(pickle.dumps(expr)), deepcopy(expr)): + assert set(i.expr_symbols) == i.expr.free_symbols + def test_list_denesting(): l0 = List(header=cgen.Line('a'), body=List(header=cgen.Line('b'))) From fc8d4d37103d0782994edd6d1db4deffea3095c9 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Mon, 12 Oct 2026 12:40:42 +0100 Subject: [PATCH 2/2] compiler: Cache FindSymbols results during Operator builds --- devito/ir/iet/visitors.py | 66 ++++++++++--- devito/operator/operator.py | 2 +- tests/test_iet.py | 190 +++++++++++++++++++++++++++++++++--- 3 files changed, 232 insertions(+), 26 deletions(-) diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 05307261c8..35b38e9264 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -8,6 +8,7 @@ from collections import OrderedDict from collections.abc import Callable, Generator, Iterable, Iterator, Sequence from contextlib import contextmanager, suppress +from contextvars import ContextVar from itertools import chain, groupby from typing import Any, Generic, TypeVar @@ -1073,6 +1074,23 @@ def visit_Node(self, o, ret=None, parents=None, in_parent=False): return ret +class FindSymbolsCache: + + """ + The caches populated by FindSymbols within `FindSymbols.caching`. + """ + + def __init__(self): + # `{id(obj): (obj, sort key)}` + self.sort_keys = {} + # `{(id(node), mode, ordered): (node, symbols)}` + self.results = {} + + +# Per execution context, so concurrent threads do not share caches +_find_symbols_cache = ContextVar('find_symbols_cache', default=None) + + class FindSymbols(LazyVisitor[Any, list[Any], None]): """ @@ -1121,24 +1139,43 @@ def _defines_aliases(n): 'defines-aliases': _defines_aliases } - # Sort keys cached while an Operator is built, see `caching_sort_keys` - _sort_keys: dict[int, tuple[Any, Any]] | None = None - @classmethod @contextmanager - def caching_sort_keys(cls): + def caching(cls): """ - Cache the sort keys within this context, as the same symbols are sorted - over and over while an Operator is built. The cache is keyed by identity, - since symbols may compare equal yet print differently, and retains each - symbol so its id cannot be reused. It is released on exit. + Cache results and sort keys within this context, as the same IETs are + searched over and over while an Operator is built. The caches are keyed + by identity and retain the keyed objects, so ids cannot be reused; nested + contexts share them, and they are released on exit of the outermost one. """ - prev = cls._sort_keys - cls._sort_keys = {} if prev is None else prev + if _find_symbols_cache.get() is not None: + yield + return + + token = _find_symbols_cache.set(FindSymbolsCache()) try: yield finally: - cls._sort_keys = prev + _find_symbols_cache.reset(token) + + def visit(self, o): + cache = _find_symbols_cache.get() + # Only Nodes handled by `visit_Node` are cached, as some others (e.g. + # Operators) carry state that may change after construction + if cache is None or \ + type(self) is not FindSymbols or \ + self.lookup_method(o).__func__ is not FindSymbols.visit_Node: + return super().visit(o) + + key = (id(o), self.mode, self.ordered) + try: + found = cache.results[key][1] + except KeyError: + found = tuple(super().visit(o)) + cache.results[key] = (o, found) + + # A copy, as callers may modify the result + return list(found) def __init__(self, mode: str = 'symbolics', ordered: bool = True) -> None: super().__init__() @@ -1156,16 +1193,17 @@ def _post_visit(self, ret): if not self.ordered: return found - cache = self._sort_keys + cache = _find_symbols_cache.get() if cache is None: return sorted(found, key=natural_sort_key) + sort_keys = cache.sort_keys def key(i): try: - return cache[id(i)][1] + return sort_keys[id(i)][1] except KeyError: v = natural_sort_key(i) - cache[id(i)] = (i, v) + sort_keys[id(i)] = (i, v) return v return sorted(found, key=key) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 50203f6828..7ca971f2b0 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -186,7 +186,7 @@ def __new__(cls, expressions, **kwargs): # Lower to a JIT-compilable object with timed_region('op-compile') as r: try: - with FindSymbols.caching_sort_keys(): + with FindSymbols.caching(): op = cls._build(expressions, **kwargs) finally: CacheInstances.clear_caches() diff --git a/tests/test_iet.py b/tests/test_iet.py index 5d9af9b2de..8058138c3c 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -1,4 +1,5 @@ import pickle +import threading from copy import deepcopy from ctypes import POINTER, c_void_p @@ -13,15 +14,16 @@ from devito.ir import SymbolRegistry from devito.ir.iet import ( Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr, - ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List, - Switch, ThreadFence, Transformer, filter_iterations, make_callable, make_efunc, - retrieve_iteration_tree + ElementalFunction, EntryFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, + Lambda, List, MetaCall, Prodder, Switch, ThreadFence, Transformer, Uxreplace, + filter_iterations, make_callable, make_efunc, retrieve_iteration_tree ) -from devito.ir.iet.visitors import sorted_efuncs +from devito.ir.iet.visitors import _find_symbols_cache, sorted_efuncs from devito.passes.iet.engine import ( Graph, _contains, _update_args_efunc, create_call_graph, reuse_efuncs ) from devito.passes.iet.languages.C import CDataManager +from devito.passes.iet.languages.openmp import ThreadedProdder from devito.symbolics import ( FLOAT, Byref, Class, FieldFromComposite, InlineIf, ListInitializer, Macro, SizeOf, String @@ -181,21 +183,187 @@ def test_find_symbols_cached_sort_keys(): temporaries = [Temp(name=f'r{i}') for i in range(12)] call = Call('foo', temporaries[::-1]) - with FindSymbols.caching_sort_keys(): + with FindSymbols.caching(): for _ in range(2): found = FindSymbols('basics').visit(call) assert [i.name for i in found] == [i.name for i in temporaries] - assert len(FindSymbols._sort_keys) == len(temporaries) + assert len(_find_symbols_cache.get().sort_keys) == len(temporaries) - # Released on exit, and never retained beyond an Operator build - assert FindSymbols._sort_keys is None +def _make_iet(): + """ + An IET exercising most node types, along with the objects it is made of. + """ + grid = Grid(shape=(4, 4)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + g = Array(name='g', dimensions=(x, y), shape=(4, 4)) + # Same name, yet distinct objects + q0 = Constant(name='q', value=1.) + q1 = Constant(name='q', value=2.) + s = Symbol(name='s') + + # Appears under two parents + shared = DummyExpr(s, q1, init=True) + # An imperfect nest + nest = Iteration([shared, Iteration(DummyExpr(g[x, y], f[x, y] + q0), y, 3)], x, 3) + body = List(body=[ + nest, + Call('foo', [Call('bar', [s, g.indexed])]), + Conditional(s > 0, shared), + ThreadedProdder(Prodder('prod', [q0])) + ]) + parameters = [f, g, q0, q1] + roots = [Callable('kernel', body, 'void', parameters=parameters), + EntryFunction('entry', body, 'int', parameters=parameters)] + + return roots, {'f': f, 'g': g, 'q0': q0, 'q1': q1, 's': s, 'x': x, 'y': y} + + +def _fingerprint(found, objects): + labels = {id(v): k for k, v in objects.items()} + return [(labels.get(id(i)), type(i).__name__, str(i)) for i in found] + + +@pytest.mark.parametrize('mode', list(FindSymbols.rules) + ['basics|symbolics']) +@pytest.mark.parametrize('ordered', [True, False]) +def test_find_symbols_caching_cold(mode, ordered): + """ + Searching a freshly built IET finds the same symbols with and without caching. + + The order is compared on the same IET by `test_find_symbols_caching_warm`, as + the traversal order of equivalent IETs may differ (`free_symbols` is a set). + """ + roots0, objects0 = _make_iet() + roots1, objects1 = _make_iet() + + expected = [sorted(_fingerprint(FindSymbols(mode, ordered=ordered).visit(i), + objects0), key=str) + for i in roots0] + with FindSymbols.caching(): + found = [sorted(_fingerprint(FindSymbols(mode, ordered=ordered).visit(i), + objects1), key=str) + for i in roots1] + + assert found == expected + + +@pytest.mark.parametrize('ordered', [True, False]) +def test_find_symbols_caching_warm(ordered): + roots, _ = _make_iet() + + for root in roots: + for mode in FindSymbols.rules: + expected = [id(i) for i in FindSymbols(mode, ordered=ordered).visit(root)] + with FindSymbols.caching(): + for _ in range(2): + found = FindSymbols(mode, ordered=ordered).visit(root) + assert [id(i) for i in found] == expected + # Callers may modify the result + found.clear() + + +def test_find_symbols_caching_rebuild(): + roots, objects = _make_iet() + t = Symbol(name='t') + + with FindSymbols.caching(): + FindSymbols('basics').visit(roots[0]) + root = Uxreplace({objects['s']: t}).visit(roots[0]) + found = FindSymbols('basics').visit(root) + + assert t in found and objects['s'] not in found + assert [id(i) for i in found] == [id(i) for i in FindSymbols('basics').visit(root)] + + +def test_find_symbols_caching_operator(): + """ + Operators are never cached, as their function table may change. + """ grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid) + h = Function(name='h', grid=grid) op = Operator(Eq(u.forward, u + 1)) - assert FindSymbols._sort_keys is None + efunc = Callable('extra', Call('bar', [h]), 'void', parameters=[h]) + + with FindSymbols.caching(): + assert h not in FindSymbols().visit(op) + op._func_table['extra'] = MetaCall(efunc, True) + assert h in FindSymbols().visit(op) + del op._func_table['extra'] + assert h not in FindSymbols().visit(op) + + +def test_find_symbols_caching_subclass(): + class FindNothing(FindSymbols): + rules = {k: (lambda n: ()) for k in FindSymbols.rules} + + roots, _ = _make_iet() + + with FindSymbols.caching(): + assert FindSymbols('basics').visit(roots[0]) + assert FindNothing('basics').visit(roots[0]) == [] + assert FindSymbols('basics').visit(roots[0]) + + +def test_find_symbols_caching_context(): + assert _find_symbols_cache.get() is None + + with FindSymbols.caching(): + cache = _find_symbols_cache.get() + assert cache is not None + # Nested contexts share the caches + with FindSymbols.caching(): + assert _find_symbols_cache.get() is cache + assert _find_symbols_cache.get() is cache + assert _find_symbols_cache.get() is None + + with pytest.raises(ValueError), FindSymbols.caching(): + raise ValueError + assert _find_symbols_cache.get() is None + + # Never retained beyond an Operator build + grid = Grid(shape=(4, 4)) + u = TimeFunction(name='u', grid=grid) + op = Operator(Eq(u.forward, u + 1)) + assert _find_symbols_cache.get() is None assert op.dimensions - assert FindSymbols._sort_keys is None + assert _find_symbols_cache.get() is None + + +def test_find_symbols_caching_threads(): + """ + Threads have their own caches, even when their contexts do not nest. + """ + entered = threading.Barrier(2) + exited = threading.Barrier(2) + seen = {} + + def first(): + with FindSymbols.caching(): + seen['first'] = _find_symbols_cache.get() + entered.wait() + # Exits while `second` is still within its context + exited.wait() + seen['first-after'] = _find_symbols_cache.get() + + def second(): + with FindSymbols.caching(): + seen['second'] = _find_symbols_cache.get() + entered.wait() + exited.wait() + seen['second-after'] = _find_symbols_cache.get() + + threads = [threading.Thread(target=i) for i in (first, second)] + for i in threads: + i.start() + for i in threads: + i.join() + + assert seen['first'] is not None and seen['second'] is not None + assert seen['first'] is not seen['second'] + assert seen['first-after'] is None and seen['second-after'] is None + assert _find_symbols_cache.get() is None def test_find_symbols_cached_sort_keys_identity(): @@ -219,7 +387,7 @@ def __str__(self): a, b = Sym('b'), Sym('a') - with FindSymbols.caching_sort_keys(): + with FindSymbols.caching(): for _ in range(2): assert FindSymbols()._post_visit([a, b]) == [b, a]