Merge branch 'main' into drop-39

Author: FabioLuporini

devito/core/gpu.py

@@ -126,16 +126,29 @@ def _rcompile_wrapper(cls, **kwargs0):
 
         def wrapper(expressions, mode='default', options=None, **kwargs1):
             kwargs = {**kwargs0, **kwargs1}
+            options = options or {}
 
             if mode == 'host':
-                options = options or {}
                 target = {
                     'platform': 'cpu64',
                     'language': 'C' if options0['par-disabled'] else 'openmp',
                     'compiler': 'custom'
                 }
             else:
-                options = {**options0, **(options or {})}
+                # Always use the default `par-tile` for recursive compilation
+                # unless the caller explicitly overrides it so that if the user
+                # supplies a multi par-tile there is no need to worry about the
+                # small kernels typically generated by recursive compilation
+                par_tile0 = options0['par-tile']
+                par_tile = options.get('par-tile')
+                if par_tile0 and par_tile:
+                    options = {**options0, **options, 'par-tile': par_tile}
+                elif par_tile0:
+                    par_tile = ParTile(par_tile0.default, default=par_tile0.default)
+                    options = {**options0, **options, 'par-tile': par_tile}
+                else:
+                    options = {**options0, **options}
+
                 target = None
 
             return rcompile(expressions, kwargs, options, target=target)

devito/ir/iet/nodes.py

@@ -1516,8 +1516,8 @@ def arguments(self):
         return self.halo_scheme.arguments
 
     @property
-    def is_empty(self):
-        return len(self.halo_scheme) == 0
+    def is_void(self):
+        return self.halo_scheme.is_void
 
     @property
     def body(self):

devito/ir/iet/visitors.py

@@ -28,9 +28,10 @@
                           IndexedData, DeviceMap)
 
 
-__all__ = ['FindApplications', 'FindNodes', 'FindSections', 'FindSymbols',
-           'MapExprStmts', 'MapHaloSpots', 'MapNodes', 'IsPerfectIteration',
-           'printAST', 'CGen', 'CInterface', 'Transformer', 'Uxreplace']
+__all__ = ['FindApplications', 'FindNodes', 'FindWithin', 'FindSections',
+           'FindSymbols', 'MapExprStmts', 'MapHaloSpots', 'MapNodes',
+           'IsPerfectIteration', 'printAST', 'CGen', 'CInterface', 'Transformer',
+           'Uxreplace']
 
 
 class Visitor(GenericVisitor):
@@ -1112,6 +1113,49 @@ def visit_Node(self, o, ret=None):
         return ret
 
 
+class FindWithin(FindNodes):
+
+    @classmethod
+    def default_retval(cls):
+        return [], False
+
+    """
+    Like FindNodes, but given an additional parameter `within=(start, stop)`,
+    it starts collecting matching nodes only after `start` is found, and stops
+    collecting matching nodes after `stop` is found.
+    """
+
+    def __init__(self, match, start, stop=None):
+        super().__init__(match)
+        self.start = start
+        self.stop = stop
+
+    def visit(self, o, ret=None):
+        found, _ = self._visit(o, ret=ret)
+        return found
+
+    def visit_Node(self, o, ret=None):
+        if ret is None:
+            ret = self.default_retval()
+        found, flag = ret
+
+        if o is self.start:
+            flag = True
+
+        if flag and self.rule(self.match, o):
+            found.append(o)
+        for i in o.children:
+            found, newflag = self._visit(i, ret=(found, flag))
+            if flag and not newflag:
+                return found, newflag
+            flag = newflag
+
+        if o is self.stop:
+            flag = False
+
+        return found, flag
+
+
 class FindApplications(Visitor):
 
     """

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
@@ -382,32 +384,36 @@ def distance(self, other):
                     if disjoint_test(self[n], other[n], sai, sit):
                         return Vector(S.ImaginaryUnit)
 
+            # Compute the distance along the current IterationInterval
             if self.function._mem_shared:
                 # Special case: the distance between two regular, thread-shared
-                # objects fallbacks to zero, as any other value would be nonsensical
+                # objects falls back to zero, as any other value would be
+                # nonsensical
+                ret.append(S.Zero)
+            elif degenerating_dimensions(sai, oai):
+                # Special case: `sai` and `oai` may be different symbolic objects
+                # but they can be proved to systematically generate the same value
                 ret.append(S.Zero)
-
             elif sai and oai and sai._defines & sit.dim._defines:
-                # E.g., `self=R<f,[t + 1, x]>`, `self.itintervals=(time, x)`, `ai=t`
+                # E.g., `self=R<f,[t + 1, x]>`, `self.itintervals=(time, x)`,
+                # and `ai=t`
                 if sit.direction is Backward:
                     ret.append(other[n] - self[n])
                 else:
                     ret.append(self[n] - other[n])
-
             elif not sai and not oai:
                 # E.g., `self=R<a,[3]>` and `other=W<a,[4]>`
                 if self[n] - other[n] == 0:
                     ret.append(S.Zero)
                 else:
                     break
-
             elif sai in self.ispace and oai in other.ispace:
                 # E.g., `self=R<f,[x, y]>`, `sai=time`,
                 #       `self.itintervals=(time, x, y)`, `n=0`
                 continue
-
             else:
-                # E.g., `self=R<u,[t+1, ii_src_0+1, ii_src_1+2]>`, `fi=p_src`, `n=1`
+                # E.g., `self=R<u,[t+1, ii_src_0+1, ii_src_1+2]>`, `fi=p_src`,
+                # and `n=1`
                 return vinf(ret)
 
         n = len(ret)
@@ -1408,3 +1414,19 @@ def disjoint_test(e0, e1, d, it):
     i1 = sympy.Interval(min(p10, p11), max(p10, p11))
 
     return not bool(i0.intersect(i1))
+
+
+def degenerating_dimensions(d0, d1):
+    """
+    True if `d0` and `d1` are Dimensions that are possibly symbolically
+    different, but they can be proved to systematically degenerate to the
+    same value, False otherwise.
+    """
+    # Case 1: ModuloDimensions of size 1
+    try:
+        if d0.is_Modulo and d1.is_Modulo and d0.modulo == d1.modulo == 1:
+            return True
+    except AttributeError:
+        pass
+
+    return False

devito/mpi/halo_scheme.py

@@ -43,6 +43,10 @@ def __hash__(self):
         return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims,
                      self.bundle))
 
+    @cached_property
+    def loc_values(self):
+        return frozenset(self.loc_indices.values())
+
     def union(self, other):
         """
         Return a new HaloSchemeEntry that is the union of this and `other`.
@@ -384,6 +388,10 @@ def owned_size(self):
                     mapper[d] = (max(maxl, s.left), max(maxr, s.right))
         return mapper
 
+    @cached_property
+    def functions(self):
+        return frozenset(self.fmapper)
+
     @cached_property
     def dimensions(self):
         retval = set()
@@ -413,6 +421,38 @@ def loc_values(self):
     def arguments(self):
         return self.dimensions | set(flatten(self.honored.values()))
 
+    def issubset(self, other):
+        """
+        Check if `self` is a subset of `other`.
+        """
+        if not isinstance(other, HaloScheme):
+            return False
+
+        if not all(f in other.fmapper for f in self.fmapper):
+            return False
+
+        for f, hse0 in self.fmapper.items():
+            hse1 = other.fmapper[f]
+
+            # Clearly, `hse0`'s halos must be a subset of `hse1`'s halos...
+            if not hse0.halos.issubset(hse1.halos) or \
+               hse0.bundle is not hse1.bundle:
+                return False
+
+            # But now, to be a subset, `hse0`'s must be expecting such halos
+            # at a time index that is less than or equal to that of `hse1`
+            if hse0.loc_dirs != hse1.loc_dirs:
+                return False
+
+            loc_dirs = hse0.loc_dirs
+            raw_loc_indices = {d: (hse0.loc_indices[d], hse1.loc_indices[d])
+                               for d in hse0.loc_indices}
+            projected_loc_indices, _ = process_loc_indices(raw_loc_indices, loc_dirs)
+            if projected_loc_indices != hse1.loc_indices:
+                return False
+
+        return True
+
     def project(self, functions):
         """
         Create a new HaloScheme that only retains the HaloSchemeEntries corresponding

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:
@@ -244,6 +258,10 @@ def _toposort(exprs):
     """
     Ensure the expression list is topologically sorted.
     """
+    if not any(isinstance(e.lhs, CTemp) for e in exprs):
+        # No CSE temps, no need to topological sort
+        return exprs
+
     dag = DAG(exprs)
 
     for e0 in exprs:
@@ -255,9 +273,14 @@ def _toposort(exprs):
                 dag.add_edge(e0, e1, force_add=True)
 
     def choose_element(queue, scheduled):
-        # Try to honor temporary names as much as possible
-        first = sorted(queue, key=lambda i: str(i.lhs)).pop(0)
-        queue.remove(first)
+        tmps = [i for i in queue if isinstance(i.lhs, CTemp)]
+        if tmps:
+            # Try to honor temporary names as much as possible
+            first = sorted(tmps, key=lambda i: i.lhs.name).pop(0)
+            queue.remove(first)
+        else:
+            first = sorted(queue, key=lambda i: exprs.index(i)).pop(0)
+            queue.remove(first)
         return first
 
     processed = dag.topological_sort(choose_element)

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)