compiler: Fix MPI alignment by using alignment_elems

Author: FabioLuporini

devito/mpi/halo_scheme.py

@@ -150,6 +150,8 @@ def __init__(self, exprs, ispace):
 
         # Further constraints on the `omapper` derivation. At construction time
         # there's none, but lowering passes may change this
+        # * `_alignment` may be a positive integer representing the alignment
+        #   requirement, in number of *elements*, of the underlying expressions
         self._alignment = None
 
     def __repr__(self):

devito/passes/iet/mpi.py

@@ -319,7 +319,7 @@ def _mark_overlappable(iet):
     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)
+        alignment = max([i._expected_alignment_elems for i in objs], default=None)
 
         hsf = hs.halo_scheme._rebuild(alignment=alignment)
         hs1 = hs._rebuild(halo_scheme=hsf)

devito/symbolics/extended_sympy.py

@@ -3,6 +3,7 @@
 """
 import re
 from contextlib import suppress
+from functools import cached_property
 
 import numpy as np
 import sympy
@@ -25,7 +26,7 @@
            'MathFunction', 'InlineIf', 'Reserved', 'ReservedWord', 'Keyword',
            'String', 'Macro', 'Class', 'MacroArgument', 'RoundUp', 'Deref',
            'Namespace', 'Rvalue', 'Null', 'SizeOf', 'rfunc', 'BasicWrapperMixin',
-           'ValueLimit', 'VectorAccess']
+           'ValueLimit', 'AlignedAccess', 'VectorAccess']
 
 
 class CondEq(sympy.Eq):
@@ -888,16 +889,47 @@ def __str__(self):
     __repr__ = __str__
 
 
-class VectorAccess(Expr, Pickable, BasicWrapperMixin):
+class AlignedAccess(Expr, Reserved, BasicWrapperMixin):
 
     """
-    Represent a vector access operation at high-level.
+    Abstract base class for an aligned access operation, that is an access to a
+    memory location that is guaranteed to be aligned to a certain byte
+    boundary.
     """
 
-    _expected_alignment = 16
+    @property
+    def _expected_alignment(self):
+        """
+        The expected alignment in bytes for the underlying LOAD/STORE operation.
+
+        To be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    @property
+    def _expected_alignment_elems(self):
+        """
+        The expected alignment in number of elements for the underlying
+        LOAD/STORE operation.
+        """
+        return self._expected_alignment // self.dtype().itemsize
+
+    @property
+    def base(self):
+        return self.args[0]
+
+    func = Reserved._rebuild
+
+    @cacheit
+    def sort_key(self, order=None):
+        # Ensure that the AlignedAccess is sorted as the base
+        return self.base.sort_key(order=order)
+
+
+class VectorAccess(AlignedAccess):
+
     """
-    The expected alignment in bytes for the accessed vector. This must be
-    honored by the compiler for correctness.
+    Represent a vector access operation at high-level.
     """
 
     def __new__(cls, *args, **kwargs):
@@ -908,21 +940,28 @@ def __str__(self):
 
     __repr__ = __str__
 
-    func = Pickable._rebuild
-
-    @property
-    def base(self):
-        return self.args[0]
+    @cached_property
+    def _expected_alignment(self):
+        """
+        The expected alignment in bytes for the underlying LOAD/STORE operation.
+        """
+        mapper = {
+            # dtype==float => lowered with float4 => 4*4=16 bytes alignment;
+            np.float32: 16,
+            # dtype==half  => lowered with float2 => 2*4=8 bytes alignment;
+            np.float16: 8
+        }
+        try:
+            return mapper[self.function.dtype]
+        except KeyError:
+            raise ValueError(
+                f"Unsupported dtype `{self.function.dtype}` for VectorAccess"
+            )
 
     @property
     def indices(self):
         return self.base.indices
 
-    @cacheit
-    def sort_key(self, order=None):
-        # Ensure that the VectorAccess is sorted as the base
-        return self.base.sort_key(order=order)
-
 
 # Some other utility objects
 Null = Macro('NULL')

devito/types/parallel.py

@@ -11,11 +11,10 @@
 from functools import cached_property
 
 import numpy as np
-from sympy import Expr
 
 from devito.exceptions import InvalidArgument
 from devito.parameters import configuration
-from devito.symbolics import Reserved, Terminal, search
+from devito.symbolics import AlignedAccess, Terminal, search
 from devito.tools import as_list, as_tuple, is_integer
 from devito.types.array import Array, ArrayObject
 from devito.types.basic import Scalar, Symbol
@@ -403,7 +402,7 @@ def __init_finalize__(self, *args, **kwargs):
         super().__init_finalize__(*args, **kwargs)
 
 
-class TensorMove(Expr, Reserved, Terminal):
+class TensorMove(AlignedAccess, Terminal):
 
     """
     Represent the LOAD/STORE of a multi-dimensional block of data from/to a higher
@@ -423,17 +422,12 @@ class TensorMove(Expr, Reserved, Terminal):
 
     _expected_alignment = 16
     """
-    The expected alignment in bytes for the accessed vector. This must be
-    honored by the compiler for correctness.
+    The expected alignment in bytes for the underlying LOAD/STORE operation.
     """
 
     def __new__(cls, base, tid0, coords, **kwargs):
         return super().__new__(cls, base, tid0, coords)
 
-    @property
-    def base(self):
-        return self.args[0]
-
     @property
     def tid0(self):
         return self.args[1]
@@ -442,10 +436,6 @@ def tid0(self):
     def coords(self):
         return self.args[2]
 
-    @property
-    def function(self):
-        return self.base.function
-
     @cached_property
     def indexed(self):
         return self.function[self.coords]
@@ -454,9 +444,5 @@ def indexed(self):
     def ndim(self):
         return self.function.ndim
 
-    func = Reserved._rebuild
-
     def _ccode(self, printer):
         return str(self)
-
-    _sympystr = _ccode