misc: Weakly cache Scope instances

Author: enwask

devito/ir/clusters/cluster.py

@@ -183,7 +183,7 @@ def dist_dimensions(self):
 
     @cached_property
     def scope(self):
-        return Scope(self.exprs)
+        return Scope.maybe_cached(as_tuple(self.exprs))
 
     @cached_property
     def functions(self):
@@ -473,7 +473,7 @@ def exprs(self):
 
     @cached_property
     def scope(self):
-        return Scope(exprs=self.exprs)
+        return Scope.maybe_cached(exprs=as_tuple(self.exprs))
 
     @cached_property
     def ispace(self):

devito/ir/clusters/visitors.py

@@ -134,7 +134,7 @@ def _fetch_scope(self, clusters):
         exprs = flatten(c.exprs for c in as_tuple(clusters))
         key = tuple(exprs)
         if key not in self.state.scopes:
-            self.state.scopes[key] = Scope(exprs)
+            self.state.scopes[key] = Scope.maybe_cached(as_tuple(exprs))
         return self.state.scopes[key]
 
     def _fetch_properties(self, clusters, prefix):

devito/ir/support/basic.py

@@ -852,6 +852,15 @@ def __init__(self, exprs, rules=None):
         self.rules = as_tuple(rules)
         assert all(callable(i) for i in self.rules)
 
+    @classmethod
+    @weak_instance_cache
+    def maybe_cached(cls, exprs, rules=None) -> 'Scope':
+        """
+        Constructs a new Scope or retrieves one from the cache if it already/still
+        exists.
+        """
+        return cls(exprs, rules)
+
     @memoized_generator
     def writes_gen(self):
         """

devito/mpi/halo_scheme.py

@@ -519,7 +519,7 @@ def classify(exprs, ispace):
     # performed before (reads) or after (writes) the OWNED region is computed
     loc_indices_from_reads = configuration['mpi'] not in ('dual',)
 
-    scope = Scope(exprs)
+    scope = Scope.maybe_cached(as_tuple(exprs))
 
     mapper = {}
     for f, r in scope.reads.items():

devito/passes/clusters/blocking.py

@@ -279,7 +279,7 @@ def callback(self, clusters, prefix):
         if len(clusters) > 1:
             # Heuristic: same as above if it induces dynamic bounds
             exprs = flatten(c.exprs for c in as_tuple(clusters))
-            scope = Scope(exprs)
+            scope = Scope.maybe_cached(as_tuple(exprs))
             if any(i.is_lex_non_stmt for i in scope.d_all_gen()):
                 return clusters
         else:

devito/passes/clusters/cse.py

@@ -124,7 +124,7 @@ def _cse(maybe_exprs, make, min_cost=1, mode='basic'):
         maybe_exprs = as_list(maybe_exprs)
         if all(e.is_Equality for e in maybe_exprs):
             exprs = maybe_exprs
-            scope = Scope(maybe_exprs)
+            scope = Scope.maybe_cached(tuple(maybe_exprs))
         else:
             exprs = [Eq(make(e), e) for e in maybe_exprs]
             scope = Scope([])

devito/passes/clusters/misc.py

@@ -355,7 +355,7 @@ def is_cross(source, sink):
             for n1, cg1 in enumerate(cgroups[n+1:], start=n+1):
 
                 # A Scope to compute all cross-ClusterGroup anti-dependences
-                scope = Scope(exprs=cg0.exprs + cg1.exprs, rules=is_cross)
+                scope = Scope(exprs=tuple(cg0.exprs + cg1.exprs), rules=is_cross)
 
                 # Anti-dependences along `prefix` break the execution flow
                 # (intuitively, "the loop nests are to be kept separated")
@@ -444,7 +444,7 @@ def callback(self, clusters, prefix):
             return clusters
 
         # Analyze and abort if fissioning would break a dependence
-        scope = Scope(flatten(c.exprs for c in clusters))
+        scope = Scope.maybe_cached(tuple(flatten(c.exprs for c in clusters)))
         if any(d._defines & dep.cause or dep.is_reduce(d) for dep in scope.d_all_gen()):
             return clusters
 

devito/passes/iet/mpi.py

@@ -41,10 +41,14 @@ def _drop_reduction_halospots(iet):
 
     # If all HaloSpot reads pertain to reductions, then the HaloSpot is useless
     for hs, expressions in MapNodes(HaloSpot, Expression).visit(iet).items():
-        scope = Scope(i.expr for i in expressions)
-        for f, v in scope.reads.items():
-            if f in hs.fmapper and all(i.is_reduction for i in v):
-                mapper[hs].add(f)
+        scope = Scope.maybe_cached(tuple(i.expr for i in expressions))
+        for k, v in hs.fmapper.items():
+            f = v.bundle or k
+            if f not in scope.reads:
+                continue
+            v = scope.reads[f]
+            if all(i.is_reduction for i in v):
+                mapper[hs].add(k)
 
     # Transform the IET introducing the "reduced" HaloSpots
     mapper = {hs: hs._rebuild(halo_scheme=hs.halo_scheme.drop(mapper[hs]))
@@ -79,7 +83,7 @@ def _hoist_redundant_from_conditionals(iet):
 
     mapper = HaloSpotMapper()
     for it, halo_spots in iter_mapper.items():
-        scope = Scope(e.expr for e in FindNodes(Expression).visit(it))
+        scope = Scope.maybe_cached(tuple(e.expr for e in FindNodes(Expression).visit(it)))
 
         for hs0 in halo_spots:
             conditions = cond_mapper[hs0]
@@ -279,7 +283,7 @@ def _mark_overlappable(iet):
         if not expressions:
             continue
 
-        scope = Scope(i.expr for i in expressions)
+        scope = Scope.maybe_cached(tuple(i.expr for i in expressions))
 
         # Comp/comm overlaps is legal only if the OWNED regions can grow
         # arbitrarly, which means all of the dependences must be carried
@@ -459,7 +463,7 @@ def _derive_scope(it, hs0, hs1):
     and ends at the HaloSpot `hs1`.
     """
     expressions = FindWithin(Expression, hs0, stop=hs1).visit(it)
-    return Scope(e.expr for e in expressions)
+    return Scope.maybe_cached(tuple(e.expr for e in expressions))
 
 
 def _check_control_flow(hs0, hs1, cond_mapper):