Merge pull request #2624 from devitocodes/compiler-speed

compiler: Address various compiler hotspots with operators containing large expression counts

Author: mloubout

devito/ir/support/basic.py

@@ -11,7 +11,8 @@
                               q_constant, q_comp_acc, q_affine, q_routine, search,
                               uxreplace)
 from devito.tools import (Tag, as_mapper, as_tuple, is_integer, filter_sorted,
-                          flatten, memoized_meth, memoized_generator)
+                          flatten, memoized_meth, memoized_generator, smart_gt,
+                          smart_lt)
 from devito.types import (ComponentAccess, Dimension, DimensionTuple, Fence,
                           CriticalRegion, Function, Symbol, Temp, TempArray,
                           TBArray)
@@ -364,11 +365,12 @@ def distance(self, other):
                 # trip count. E.g. it ranges from 0 to 3; `other` performs a
                 # constant access at 4
                 for v in (self[n], other[n]):
-                    try:
-                        if bool(v < sit.symbolic_min or v > sit.symbolic_max):
-                            return Vector(S.ImaginaryUnit)
-                    except TypeError:
-                        pass
+                    # Note: Uses smart_ comparisons avoid evaluating expensive
+                    # symbolic Lt or Gt operations,
+                    # Note: Boolean is split to make the conditional short
+                    # circuit more frequently for mild speedup.
+                    if smart_lt(v, sit.symbolic_min) or smart_gt(v, sit.symbolic_max):
+                        return Vector(S.ImaginaryUnit)
 
                 # Case 2: `sit` is an IterationInterval over a local SubDimension
                 # and `other` performs a constant access

devito/passes/clusters/cse.py

@@ -1,4 +1,4 @@
-from collections import defaultdict
+from collections import defaultdict, Counter
 from functools import cached_property, singledispatch
 
 import numpy as np
@@ -13,6 +13,7 @@
 from devito.finite_differences.differentiable import IndexDerivative
 from devito.ir import Cluster, Scope, cluster_pass
 from devito.symbolics import estimate_cost, q_leaf, q_terminal
+from devito.symbolics.search import search
 from devito.symbolics.manipulation import _uxreplace
 from devito.tools import DAG, as_list, as_tuple, frozendict, extract_dtype
 from devito.types import Eq, Symbol, Temp
@@ -25,9 +26,15 @@ class CTemp(Temp):
     """
     A cluster-level Temp, similar to Temp, ensured to have different priority
     """
+
     ordering_of_classes.insert(ordering_of_classes.index('Temp') + 1, 'CTemp')
 
 
+def retrieve_ctemps(exprs, mode='all'):
+    """Shorthand to retrieve the CTemps in `exprs`"""
+    return search(exprs, lambda expr: isinstance(expr, CTemp), mode, 'dfs')
+
+
 @cluster_pass
 def cse(cluster, sregistry=None, options=None, **kwargs):
     """
@@ -225,8 +232,15 @@ def _compact(exprs, exclude):
 
     mapper = {e.lhs: e.rhs for e in candidates if q_leaf(e.rhs)}
 
-    mapper.update({e.lhs: e.rhs for e in candidates
-                   if sum([i.rhs.count(e.lhs) for i in exprs]) == 1})
+    # Find all the CTemps in expression right-hand-sides without removing duplicates
+    ctemps = retrieve_ctemps(e.rhs for e in exprs)
+
+    # If there are ctemps in the expressions, then add any that only appear once to
+    # the mapper
+    if ctemps:
+        ctemp_count = Counter(ctemps)
+        mapper.update({e.lhs: e.rhs for e in candidates
+                       if ctemp_count[e.lhs] == 1})
 
     processed = []
     for e in exprs:

devito/passes/clusters/misc.py

@@ -352,8 +352,7 @@ def is_cross(source, sink):
                 v = len(cg0.exprs)
                 return t0 < v <= t1 or t1 < v <= t0
 
-            for cg1 in cgroups[n+1:]:
-                n1 = cgroups.index(cg1)
+            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)

devito/symbolics/search.py

@@ -155,7 +155,7 @@ def retrieve_functions(exprs, mode='all', deep=False):
 
 
 def retrieve_symbols(exprs, mode='all'):
-    """Shorthand to retrieve the Scalar in ``exprs``."""
+    """Shorthand to retrieve the Scalar in `exprs`."""
     return search(exprs, q_symbol, mode, 'dfs')
 
 

devito/tools/utils.py

@@ -1,6 +1,6 @@
 from collections import OrderedDict
 from collections.abc import Iterable
-from functools import reduce
+from functools import reduce, wraps
 from itertools import chain, combinations, groupby, product, zip_longest
 from operator import attrgetter, mul
 import types
@@ -12,7 +12,8 @@
            'roundm', 'powerset', 'invert', 'flatten', 'single_or', 'filter_ordered',
            'as_mapper', 'filter_sorted', 'pprint', 'sweep', 'all_equal', 'as_list',
            'indices_to_slices', 'indices_to_sections', 'transitive_closure',
-           'humanbytes', 'contains_val', 'sorted_priority', 'as_set', 'is_number']
+           'humanbytes', 'contains_val', 'sorted_priority', 'as_set', 'is_number',
+           'smart_lt', 'smart_gt']
 
 
 def prod(iterable, initial=1):
@@ -346,3 +347,35 @@ def key(i):
         return (v, str(type(i)))
 
     return sorted(items, key=key, reverse=True)
+
+
+def avoid_symbolic(default=None):
+    """
+    Decorator to avoid calling a function where doing so will result in symbolic
+    computation being performed. For use if symbolic computation may be slow. In
+    the case that an arg is symbolic, just give up and return a default value.
+    """
+    def _avoid_symbolic(func):
+        @wraps(func)
+        def wrapper(*args):
+            if any(isinstance(expr, sympy.Basic) for expr in args):
+                # An argument is symbolic, so give up and assume default
+                return default
+
+            return func(*args)
+
+        return wrapper
+
+    return _avoid_symbolic
+
+
+@avoid_symbolic(default=False)
+def smart_lt(a, b):
+    """An Lt that gives up and returns False if supplied a symbolic argument"""
+    return a < b
+
+
+@avoid_symbolic(default=False)
+def smart_gt(a, b):
+    """A Gt that gives up and returns False if supplied a symbolic argument"""
+    return a > b