compiler: Fix HaloScheme.omapper taking alignment into account

Author: FabioLuporini

devito/mpi/halo_scheme.py

@@ -9,6 +9,7 @@
 
 from devito.data import CENTER, CORE, LEFT, OWNED, RIGHT
 from devito.ir.support import Forward, Scope
+from devito.symbolics import IntDiv
 from devito.symbolics.manipulation import _uxreplace_registry
 from devito.tools import (
     EnrichedTuple, Reconstructable, Tag, as_tuple, filter_ordered, filter_sorted, flatten,
@@ -147,6 +148,10 @@ def __init__(self, exprs, ispace):
             self._honored[i.root] = frozenset([(ltk, rtk)])
         self._honored = frozendict(self._honored)
 
+        # Further constraints on the `omapper` derivation. At construction time
+        # there's none, but lowering passes may change this
+        self._alignment = None
+
     def __repr__(self):
         fnames = ",".join(i.name for i in set(self._mapper))
         return f"HaloScheme<{fnames}>"
@@ -162,7 +167,7 @@ def __len__(self):
     def __hash__(self):
         return hash((self._mapper.__hash__(), self.honored.__hash__()))
 
-    def _rebuild(self, fmapper=None, honored=None):
+    def _rebuild(self, fmapper=None, honored=None, alignment=None):
         """
         Rebuild a HaloScheme from the provided `fmapper` and `honored`. Reuse
         `self`'s values for the missing arguments.
@@ -176,6 +181,7 @@ def _rebuild(self, fmapper=None, honored=None):
 
         obj._mapper = frozendict(fmapper)
         obj._honored = frozendict(honored)
+        obj._alignment = alignment or self._alignment
 
         return obj
 
@@ -248,10 +254,14 @@ def is_void(self):
     @cached_property
     def omapper(self):
         """
-        Logical decomposition of the DOMAIN region into OWNED and CORE sub-regions.
+        Logical decomposition of the DOMAIN region into OWNED and CORE sub-regions,
+        "cumulative" over all DiscreteFunctions in the HaloScheme.
+
+        The computed OMapper takes into account:
 
-        This is "cumulative" over all DiscreteFunctions in the HaloScheme; it also
-        takes into account IterationSpace offsets induced by SubDomains/SubDimensions.
+            * The offsets induced by SubDomains/SubDimensions ("thickness");
+            * Any data alignment requirement of the underlying expressions
+              (`_alignment` attribute).
 
         Examples
         --------
@@ -373,28 +383,62 @@ def omapper(self):
 
                 if s is CENTER:
                     where.append((d, CORE, s))
-                    mapper[d] = (d.symbolic_min + osl,
-                                 d.symbolic_max - osr)
+
+                    mapper[d] = (
+                        d.symbolic_min + osl,
+                        d.symbolic_max - osr
+                    )
+
                     if nl != 0:
                         mapper[nl] = (Max(nl - osl, 0),)
                     if nr != 0:
                         mapper[nr] = (Max(nr - osr, 0),)
                 else:
                     where.append((d, OWNED, s))
+
                     if s is LEFT:
-                        mapper[d] = (d.symbolic_min,
-                                     Min(d.symbolic_min + osl - 1, d.symbolic_max - nr))
+                        mapper[d] = (
+                            d.symbolic_min,
+                            Min(d.symbolic_min + osl - 1, d.symbolic_max - nr)
+                        )
+
                         if nl != 0:
                             mapper[nl] = (nl,)
                             mapper[nr] = (0,)
                     else:
-                        mapper[d] = (Max(d.symbolic_max - osr + 1, d.symbolic_min + nl),
-                                     d.symbolic_max)
+                        mapper[d] = (
+                            Max(d.symbolic_max - osr + 1, d.symbolic_min + nl),
+                            d.symbolic_max
+                        )
+
                         if nr != 0:
                             mapper[nl] = (0,)
                             mapper[nr] = (nr,)
+
             processed.append((tuple(where), frozendict(mapper)))
 
+        # Apply the alignment constraints, if any
+        # First, get the fastest varying (contiguous) Dimension, which is the
+        # one that matters for alignment
+        if self._alignment:
+            fvds = {f.dimensions[-1] for f in self.fmapper}
+            if len(fvds) != 1:
+                raise HaloSchemeException(
+                    "Unexpected contiguous Dimensions found while computing the "
+                    f"`omapper`: {fvds}"
+                )
+            fvd = fvds.pop()
+
+            for i, (where, mapper) in enumerate(list(processed)):
+                try:
+                    m, M = mapper[fvd]
+                except KeyError:
+                    continue
+
+                aligned_m = IntDiv(m, self._alignment) * self._alignment
+
+                processed[i] = (where, frozendict({**mapper, fvd: (aligned_m, M)}))
+
         _, core = processed.pop(0)
         owned = processed
 

devito/passes/iet/mpi.py

@@ -11,7 +11,7 @@
 from devito.mpi.reduction_scheme import DistReduce
 from devito.mpi.routines import HaloExchangeBuilder, ReductionBuilder
 from devito.passes.iet.engine import iet_pass
-from devito.symbolics import VectorAccess
+from devito.symbolics import VectorAccess, search
 from devito.tools import generator
 from devito.types import TensorMove
 
@@ -287,11 +287,11 @@ def _mark_overlappable(iet):
         # Check legality. Comp/comm overlaps is legal only if the OWNED regions
         # can grow arbitrarily, which means all of the dependencies must be
         # carried along a non-halo Dimension
-        expressions = FindNodes(Expression).visit(hs)
-        if not expressions:
+        exprs = FindNodes(Expression).visit(hs)
+        if not exprs:
             continue
 
-        scope = Scope(i.expr for i in expressions)
+        scope = Scope([n.expr for n in exprs])
 
         for dep in scope.d_all_gen():
             if dep.function in hs.functions:
@@ -305,11 +305,28 @@ def _mark_overlappable(iet):
                     #     f[x, y] = ...
                     break
         else:
-            # All good!
+            # All good -- we can perform comp/comm overlap!
             found.append(hs)
 
+    # The underlying `exprs` might have data alignment constraints due to the
+    # presence of objects such as VectorAccess or TensorMove, which expect the
+    # starting address of the data to be aligned to a certain value. Comp/comm
+    # overlap creates multiple iteration spaces (for the core and owned
+    # regions), which might break the alignment contract if we don't play safe
+    # -- imposing these regions start at a carefully rounded-down point, at the
+    # cost of potentially performing a bit of redundant compute
+    mapper = {}
+    for hs in found:
+        exprs = [n.expr for n in FindNodes(Expression).visit(hs)]
+        objs = search(exprs, (VectorAccess, TensorMove))
+        alignment = max([i._expected_alignment for i in objs], default=None)
+
+        hsf = hs.halo_scheme._rebuild(alignment=alignment)
+        hs1 = hs._rebuild(halo_scheme=hsf)
+
+        mapper[hs] = OverlappableHaloSpot(**hs1.args)
+
     # Transform the IET replacing HaloSpots with OverlappableHaloSpots
-    mapper = {hs: OverlappableHaloSpot(**hs.args) for hs in found}
     iet = Transformer(mapper, nested=True).visit(iet)
 
     return iet

devito/symbolics/extended_sympy.py

@@ -119,6 +119,10 @@ def __new__(cls, lhs, rhs, params=None):
         elif rhs == 1 or rhs is None:
             return lhs
 
+        if is_integer(lhs) and is_integer(rhs):
+            # Both sides are plain integers -- perform the division right away
+            return lhs // rhs
+
         if not is_integer(rhs):
             # Perhaps it's a symbolic RHS -- but we wanna be sure it's of type int
             if not hasattr(rhs, 'dtype'):
@@ -890,6 +894,12 @@ class VectorAccess(Expr, Pickable, BasicWrapperMixin):
     Represent a vector access operation at high-level.
     """
 
+    _expected_alignment = 16
+    """
+    The expected alignment in bytes for the accessed vector. This must be
+    honored by the compiler for correctness.
+    """
+
     def __new__(cls, *args, **kwargs):
         return Expr.__new__(cls, *args)
 

devito/types/parallel.py

@@ -421,6 +421,12 @@ class TensorMove(Expr, Reserved, Terminal):
 
     __rargs__ = ('base', 'tid0', 'coords')
 
+    _expected_alignment = 16
+    """
+    The expected alignment in bytes for the accessed vector. This must be
+    honored by the compiler for correctness.
+    """
+
     def __new__(cls, base, tid0, coords, **kwargs):
         return super().__new__(cls, base, tid0, coords)