Merge pull request #2629 from devitocodes/fix-mpi-hoisting-2

compiler: Revamp MPI hoisting and merging

Author: mloubout

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

@@ -382,32 +382,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 +1412,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