compiler: Tweak findsymbols caching and add an unsorted mode

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)
 

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

devito/ir/iet/visitors.py

@@ -7,7 +7,8 @@
 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 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]):
 
     """
@@ -1090,6 +1108,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,18 +1139,74 @@ def _defines_aliases(n):
         'defines-aliases': _defines_aliases
     }
 
-    def __init__(self, mode: str = 'symbolics') -> None:
+    @classmethod
+    @contextmanager
+    def caching(cls):
+        """
+        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.
+        """
+        if _find_symbols_cache.get() is not None:
+            yield
+            return
+
+        token = _find_symbols_cache.set(FindSymbolsCache())
+        try:
+            yield
+        finally:
+            _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__()
 
         self.mode = mode
+        self.ordered = ordered
         modes = mode.split('|')
         if len(modes) == 1:
             self.rule = self.rules[mode]
         else:
             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 = _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 sort_keys[id(i)][1]
+            except KeyError:
+                v = natural_sort_key(i)
+                sort_keys[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)

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():
+                    op = cls._build(expressions, **kwargs)
             finally:
                 CacheInstances.clear_caches()
                 memoized_func.clear_build_caches()

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