1+ import gc
2+ import pickle
13import weakref
4+ from copy import copy , deepcopy
25from ctypes import byref , c_void_p
36
47import numpy as np
58import pytest
69from sympy import Expr
710
811from 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
1418from 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
0 commit comments