compiler: Memoize FindNodes

Author: FabioLuporini

devito/ir/iet/visitors.py

@@ -1172,8 +1172,13 @@ class FindNodes(LazyVisitor[Node, list[Node], None]):
     def __init__(self, match: type, mode: str = 'type') -> None:
         super().__init__()
         self.match = match
+        self.mode = mode
         self.rule = self.rules[mode]
 
+    @memoized_weak_meth(key=lambda i: (i.match, i.mode), freeze=tuple, thaw=list)
+    def visit(self, o, *args, **kwargs):
+        return super().visit(o, *args, **kwargs)
+
     def visit_Node(self, o: Node, **kwargs) -> Iterator[Node]:
         if self.rule(self.match, o):
             yield o
@@ -1194,6 +1199,11 @@ def __init__(self, match: type, start: Node, stop: Node | None = None) -> None:
         self.start = start
         self.stop = stop
 
+    def visit(self, o, *args, **kwargs):
+        # `start` and `stop` are part of this visitor's state.
+        return GenericVisitor.visit(self, o, *args, **kwargs)
+
+
     def visit_object(self, o: object, flag: bool = False) -> LazyVisit[Node, bool]:
         yield from ()
         return flag  # noqa: B901

tests/test_visitors.py

@@ -6,8 +6,8 @@
 from devito.ir.equations import DummyEq
 from devito.ir.iet import (
     Block, Call, Callable, Conditional, Expression, FindApplications, FindNodes,
-    FindSections, FindSymbols, IsPerfectIteration, Iteration, MapNodes, Transformer,
-    Uxreplace, printAST
+    FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration, MapNodes,
+    Transformer, Uxreplace, printAST
 )
 from devito.types import Array, SpaceDimension, Symbol
 
@@ -210,6 +210,15 @@ def test_find_sections(exprs, block1, block2, block3):
     assert len(found[2]) == 1
 
 
+def test_find_within_not_cached_like_findnodes(block3):
+    expr0 = FindWithin(Expression, block3.nodes[0], block3.nodes[1]).visit(block3)
+    expr1 = FindWithin(Expression, block3.nodes[1], block3.nodes[2]).visit(block3)
+
+    assert len(expr0) == 3
+    assert len(expr1) == 3
+    assert expr0 != expr1
+
+
 def test_is_perfect_iteration(block1, block2, block3, block4):
     checker = IsPerfectIteration()