Skip to content

Commit eae0039

Browse files
committed
compiler: Remove IET visitor result memoization
Cached results and cache partitions can keep temporary IET roots alive despite weak cache keys. Remove shared result memoization from FindNodes, FindSymbols and FindApplications, along with the unused helper and the redundant FindWithin.visit override. Add regression coverage for root and child retention, scope query targets, returned back-references, repeated visits, subclasses and copied Operators. For a space-order-16 3D TTI Born Operator, six paired runs changed median construction and code generation from 21.18s to 21.24s with identical C.
1 parent 1069c63 commit eae0039

5 files changed

Lines changed: 119 additions & 108 deletions

File tree

‎devito/ir/iet/visitors.py‎

Lines changed: 1 addition & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
from devito.symbolics.extended_dtypes import NoDeclStruct
2929
from devito.tools import (
3030
GenericVisitor, as_tuple, c_restrict_void_p, filter_ordered, filter_sorted, flatten,
31-
is_external_ctype, memoized_weak_meth, natural_sort_key
31+
is_external_ctype, natural_sort_key
3232
)
3333
from devito.types import (
3434
ArrayObject, CompositeObject, DeviceMap, Dimension, IndexedData, Pointer
@@ -1127,10 +1127,6 @@ def __init__(self, mode: str = 'symbolics') -> None:
11271127
else:
11281128
self.rule = lambda n: chain(*[self.rules[mode](n) for mode in modes])
11291129

1130-
@memoized_weak_meth(key=lambda i: i.mode, freeze=tuple, thaw=list)
1131-
def visit(self, o, *args, **kwargs):
1132-
return super().visit(o, *args, **kwargs)
1133-
11341130
def _post_visit(self, ret):
11351131
return sorted(filter_ordered(ret, key=id), key=natural_sort_key)
11361132

@@ -1180,10 +1176,6 @@ def __init__(self, match: type, mode: str = 'type') -> None:
11801176
self.mode = mode
11811177
self.rule = self.rules[mode]
11821178

1183-
@memoized_weak_meth(key=lambda i: (i.match, i.mode), freeze=tuple, thaw=list)
1184-
def visit(self, o, *args, **kwargs):
1185-
return super().visit(o, *args, **kwargs)
1186-
11871179
def visit_Node(self, o: Node, **kwargs) -> Iterator[Node]:
11881180
if self.rule(self.match, o):
11891181
yield o
@@ -1204,10 +1196,6 @@ def __init__(self, match: type, start: Node, stop: Node | None = None) -> None:
12041196
self.start = start
12051197
self.stop = stop
12061198

1207-
def visit(self, o, *args, **kwargs):
1208-
# `start` and `stop` are part of this visitor's state.
1209-
return GenericVisitor.visit(self, o, *args, **kwargs)
1210-
12111199
def visit_object(self, o: object, flag: bool = False) -> LazyVisit[Node, bool]:
12121200
yield from ()
12131201
return flag # noqa: B901
@@ -1258,10 +1246,6 @@ def __init__(self, cls: type[ApplicationType] = Application):
12581246
self.cls = cls
12591247
self.match = lambda i: isinstance(i, cls) and not isinstance(i, Basic)
12601248

1261-
@memoized_weak_meth(key=lambda i: i.cls, freeze=frozenset, thaw=set)
1262-
def visit(self, o, *args, **kwargs):
1263-
return super().visit(o, *args, **kwargs)
1264-
12651249
def _post_visit(self, ret):
12661250
return set(ret)
12671251

‎devito/tools/memoization.py‎

Lines changed: 0 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,13 @@
22
from functools import lru_cache, partial, wraps
33
from itertools import tee
44
from typing import TypeVar
5-
from weakref import WeakKeyDictionary
65

76
__all__ = [
87
'CacheInstances',
98
'cached_hash',
109
'memoized_func',
1110
'memoized_generator',
1211
'memoized_meth',
13-
'memoized_weak_meth',
1412
'reuse_if_unchanged'
1513
]
1614

@@ -191,54 +189,6 @@ def __call__(self, *args, **kwargs):
191189
return result
192190

193191

194-
def memoized_weak_meth(*, key=None, freeze=None, thaw=None):
195-
"""
196-
Cache a method result against its first argument using weak references.
197-
198-
This is useful for visitors operating on temporary IR roots: the cache can
199-
be shared across short-lived visitor instances without keeping those roots
200-
alive. Only calls without extra arguments are cached; all other calls fall
201-
back to the wrapped method.
202-
203-
Parameters
204-
----------
205-
key : callable, optional
206-
A callable receiving ``self`` and returning a hashable cache partition.
207-
freeze : callable, optional
208-
Convert the method result before storing it in the cache.
209-
thaw : callable, optional
210-
Convert the cached value before returning it to the caller.
211-
"""
212-
def decorator(func):
213-
caches = {}
214-
215-
@wraps(func)
216-
def wrapper(self, o, *args, **kwargs):
217-
if args or kwargs:
218-
return func(self, o, *args, **kwargs)
219-
220-
try:
221-
partition = key(self) if key is not None else None
222-
cache = caches.setdefault(partition, WeakKeyDictionary())
223-
ret = cache[o]
224-
except KeyError:
225-
ret = func(self, o)
226-
if freeze is not None:
227-
ret = freeze(ret)
228-
cache[o] = ret
229-
except TypeError:
230-
return func(self, o)
231-
232-
if thaw is not None:
233-
return thaw(ret)
234-
235-
return ret
236-
237-
return wrapper
238-
239-
return decorator
240-
241-
242192
# Describes the type of a subclass of CacheInstances
243193
InstanceType = TypeVar('InstanceType', bound='CacheInstances', covariant=True)
244194

‎tests/test_caching.py‎

Lines changed: 81 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,20 @@
1+
import gc
2+
import pickle
13
import weakref
4+
from copy import copy, deepcopy
25
from ctypes import byref, c_void_p
36

47
import numpy as np
58
import pytest
69
from sympy import Expr
710

811
from devito import (
9-
ConditionalDimension, Constant, DefaultDimension, Dimension, Eq, Function, Grid,
12+
ConditionalDimension, Constant, DefaultDimension, Dimension, Eq, Function, Grid, Min,
1013
Operator, SparseFunction, SparseTimeFunction, SubDimension, TensorFunction,
1114
TensorTimeFunction, TimeFunction, VectorFunction, VectorTimeFunction, _SymbolCache,
1215
clear_cache, solve, switchconfig
1316
)
17+
from devito.ir.iet import Call, FindApplications, FindNodes, FindSymbols, List, Node
1418
from devito.types import (
1519
DeviceID, LocalObject, NPThreads, NThreadsBase, Object, Scalar, Symbol, ThreadID
1620
)
@@ -790,6 +794,79 @@ class TestMemoryLeaks:
790794
Tests ensuring there are no memory leaks.
791795
"""
792796

797+
@pytest.mark.parametrize('nested', [False, True])
798+
def test_findnodes_leakage(self, nested):
799+
"""A traversal containing its root must not keep the tree alive."""
800+
def visit_temporary_tree():
801+
tree = List(body=[Call('foo')] if nested else [])
802+
nodes = [tree, *tree.body]
803+
references = [weakref.ref(i) for i in nodes]
804+
assert FindNodes(Node).visit(tree) == nodes
805+
return references
806+
807+
references = visit_temporary_tree()
808+
gc.collect()
809+
clear_cache()
810+
811+
assert all(i() is None for i in references)
812+
813+
@pytest.mark.parametrize('kind', ['list', 'operator', 'jitted-operator'])
814+
@pytest.mark.parametrize('copier', [
815+
copy, deepcopy, pytest.param(lambda o: pickle.loads(pickle.dumps(o)), id='pickle')
816+
])
817+
def test_copied_iet_leakage(self, kind, copier):
818+
"""An unvisited copy must not retain the original traversal's root."""
819+
def copy_temporary_tree():
820+
if kind != 'list':
821+
f = Function(name='f', grid=Grid(shape=(3, 3)))
822+
tree = Operator(Eq(f, f + 1))
823+
if kind == 'jitted-operator':
824+
tree.apply()
825+
else:
826+
tree = List()
827+
assert FindNodes(Node).visit(tree)[0] is tree
828+
return copier(tree), weakref.ref(tree)
829+
830+
copied, reference = copy_temporary_tree()
831+
clear_cache()
832+
assert reference() is None
833+
assert FindNodes(Node).visit(copied)[0] is copied
834+
835+
@pytest.mark.parametrize('match_root', [False, True])
836+
def test_findnodes_scope_leakage(self, match_root):
837+
"""A scope query must not keep its target or containing tree alive."""
838+
def visit_temporary_tree():
839+
child = Call('foo')
840+
tree = List(body=[child])
841+
match = tree if match_root else child
842+
expected = [] if match_root else [tree]
843+
assert FindNodes(match, mode='scope').visit(tree) == expected
844+
return weakref.ref(tree), weakref.ref(child)
845+
846+
references = visit_temporary_tree()
847+
clear_cache()
848+
assert all(i() is None for i in references)
849+
850+
@pytest.mark.parametrize('visitor', [FindSymbols, FindApplications])
851+
def test_visitor_leakage_backref(self, visitor):
852+
"""A returned object's reference to the root must not cause a leak."""
853+
def visit_temporary_tree():
854+
if visitor is FindSymbols:
855+
result = Object(name='context', dtype=c_void_p)
856+
tree = Call('foo', arguments=[result])
857+
result.value = tree
858+
else:
859+
result = Min(Symbol(name='s'), 1)
860+
tree = Call('foo', arguments=[result])
861+
result._test_owner = tree
862+
863+
assert result in visitor().visit(tree)
864+
return weakref.ref(tree), weakref.ref(result)
865+
866+
references = visit_temporary_tree()
867+
clear_cache()
868+
assert all(i() is None for i in references)
869+
793870
def test_operator_leakage_function(self):
794871
"""
795872
Test to ensure that Operator creation does not cause memory leaks for
@@ -806,6 +883,9 @@ def test_operator_leakage_function(self):
806883
# Create operator and delete everything again
807884
op = Operator(Eq(f, 2 * g))
808885
w_op = weakref.ref(op)
886+
FindNodes(Node).visit(op)
887+
FindSymbols().visit(op)
888+
FindApplications().visit(op)
809889
del op
810890
del f
811891
del g

‎tests/test_tools.py‎

Lines changed: 1 addition & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from devito import Eq, Operator, switchenv
99
from devito.tools import (
1010
CacheInstances, DefaultFrozenDict, UnboundedMultiTuple, UnboundTuple, ctypes_to_cstr,
11-
filter_ordered, memoized_meth, memoized_weak_meth, toposort, transitive_closure
11+
filter_ordered, memoized_meth, toposort, transitive_closure
1212
)
1313
from devito.types.basic import Symbol
1414

@@ -88,42 +88,6 @@ def f(self, x=None):
8888
assert obj.calls == 4
8989

9090

91-
def test_memoized_weak_meth():
92-
93-
class Root:
94-
pass
95-
96-
class Obj:
97-
98-
def __init__(self, mode):
99-
self.mode = mode
100-
self.calls = 0
101-
102-
@memoized_weak_meth(key=lambda i: i.mode, freeze=tuple, thaw=list)
103-
def f(self, root):
104-
self.calls += 1
105-
return [self.mode]
106-
107-
root = Root()
108-
obj0 = Obj('a')
109-
obj1 = Obj('a')
110-
obj2 = Obj('b')
111-
112-
ret = obj0.f(root)
113-
ret.append('mutated')
114-
115-
assert obj1.f(root) == ['a']
116-
assert obj0.calls == 1
117-
assert obj1.calls == 0
118-
119-
assert obj2.f(root) == ['b']
120-
assert obj2.calls == 1
121-
122-
assert obj0.f([]) == ['a']
123-
assert obj0.f([]) == ['a']
124-
assert obj0.calls == 3
125-
126-
12791
def test_default_frozen_dict():
12892
mapper = DefaultFrozenDict({'a': 'b'}, default='c')
12993

‎tests/test_visitors.py‎

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@
88
from devito.ir.equations import DummyEq
99
from devito.ir.iet import (
1010
Block, Call, Callable, Conditional, Definition, Expression, FindApplications,
11-
FindNodes, FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration,
12-
MapNodes, Transformer, Uxreplace, printAST
11+
FindNodes, FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration, List,
12+
MapNodes, Node, Transformer, Uxreplace, printAST
1313
)
1414
from devito.symbolics import ListInitializer
1515
from devito.types import Array, LocalObject, SpaceDimension, Symbol
@@ -213,7 +213,40 @@ def test_find_sections(exprs, block1, block2, block3):
213213
assert len(found[2]) == 1
214214

215215

216-
def test_find_within_not_cached_like_findnodes(block3):
216+
@pytest.mark.parametrize('match', [Node, List, Call])
217+
def test_find_nodes_repeated(match):
218+
call0 = Call('foo')
219+
call1 = Call('bar')
220+
inner = List(body=[call1])
221+
tree = List(body=[call0, inner])
222+
expected = [i for i in [tree, call0, inner, call1] if isinstance(i, match)]
223+
224+
finder = FindNodes(match)
225+
result = finder.visit(tree)
226+
assert result == expected
227+
result.clear()
228+
assert finder.visit(tree) == expected
229+
assert FindNodes(match).visit(tree) == expected
230+
231+
232+
def test_find_nodes_subclass():
233+
234+
class CallsOnly(FindNodes):
235+
236+
def visit_Node(self, o, **kwargs):
237+
if isinstance(o, Call):
238+
yield o
239+
for child in o.children:
240+
yield from self._visit(child, **kwargs)
241+
242+
call = Call('foo')
243+
tree = List(body=[call])
244+
assert FindNodes(Node).visit(tree) == [tree, call]
245+
assert CallsOnly(Node).visit(tree) == [call]
246+
assert FindNodes(Node).visit(tree) == [tree, call]
247+
248+
249+
def test_find_within_bounds(block3):
217250
expr0 = FindWithin(Expression, block3.nodes[0], block3.nodes[1]).visit(block3)
218251
expr1 = FindWithin(Expression, block3.nodes[1], block3.nodes[2]).visit(block3)
219252

0 commit comments

Comments
 (0)