compiler: fix various corner case of multi buffering

Author: mloubout

devito/ir/equations/equation.py

@@ -237,6 +237,14 @@ def __new__(cls, *args, **kwargs):
                     cond = d.relation(cond, GuardFactor(d))
                 conditionals[d] = cond
 
+        # Replace the ConditionalDimensions in `expr`
+        for d, cond in conditionals.items():
+            # Replace dimension with index
+            index = d.index
+            index = index - relational_min(cond, d.parent)
+            shift = relational_shift(cond, d.parent)
+            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
+
         # Merge conditionals when possible. E.g if we have an implicit_dim
         # and there is a dimension with the same parent, we ca merged
         # its condition
@@ -247,19 +255,13 @@ def __new__(cls, *args, **kwargs):
                 if cd.parent == d.parent and cd != d:
                     cond = conditionals.pop(d)
                     mode = cd.relation and d.relation
-                    conditionals[cd] = mode(cond, conditionals[cd])
+                    if issubclass(mode, sympy.Or):
+                        conditionals[d] = cond
+                        conditionals.pop(cd)
+                    else:
+                        conditionals[cd] = mode(cond, conditionals[cd])
                     break
 
-        conditionals = frozendict(conditionals)
-
-        # Replace the ConditionalDimensions in `expr`
-        for d, cond in conditionals.items():
-            # Replace dimension with index
-            index = d.index
-            index = index - relational_min(cond, d.parent)
-            shift = relational_shift(cond, d.parent)
-            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
-
         # Lower all Differentiable operations into SymPy operations
         rhs = diff2sympy(expr.rhs)
 

devito/passes/clusters/asynchrony.py

@@ -193,7 +193,7 @@ def memcpy_prefetch(clusters, key0, sregistry):
         if c.properties.is_prefetchable(d._defines):
             _actions_from_update_memcpy(c, d, clusters, actions, sregistry)
         elif d.is_Custom and is_integer(c.ispace[d].size):
-            _actions_from_init(c, d, actions)
+            _actions_from_init(c, d, clusters, actions)
 
     # Attach the computed Actions
     processed = []
@@ -214,7 +214,7 @@ def memcpy_prefetch(clusters, key0, sregistry):
     return processed
 
 
-def _actions_from_init(c, d, actions):
+def _actions_from_init(c, d, clusters, actions):
     e = c.exprs[0]
     function = e.rhs.function
     target = e.lhs.function

devito/passes/clusters/buffering.py

@@ -1,9 +1,9 @@
 from collections import defaultdict, namedtuple
 from functools import cached_property
-from itertools import chain
+from itertools import chain, groupby
 
 import numpy as np
-from sympy import S, simplify
+from sympy import Mod, S, simplify
 
 from devito.exceptions import CompilationError
 from devito.ir import (
@@ -116,7 +116,7 @@ def key(f):
     # Then we inject them into the Clusters. This involves creating the
     # initializing Clusters, and replacing the buffered Functions with the buffers
     clusters = InjectBuffers(mapper, sregistry, options).process(clusters)
-
+    print(clusters)
     return clusters
 
 
@@ -142,14 +142,18 @@ def callback(self, clusters, prefix):
             return clusters
         d = prefix[-1].dim
 
-        key = lambda f, *args: f in self.mapper
+        def key(f, *args):
+            for (ff, _) in self.mapper:
+                if f == ff:
+                    return True
+            return False
         bfmap = map_buffered_functions(clusters, key)
 
         # A BufferDescriptor is a simple data structure storing additional
         # information about a buffer, harvested from the subset of `clusters`
         # that access it
-        descriptors = {b: BufferDescriptor(f, b, bfmap[f])
-                       for f, b in self.mapper.items()
+        descriptors = {b: BufferDescriptor(f, b, bfmap[f], g)
+                       for (f, g), b in self.mapper.items()
                        if f in bfmap}
 
         # Are we inside the right `d`?
@@ -184,6 +188,8 @@ def callback(self, clusters, prefix):
                     continue
                 if c not in v.firstread:
                     continue
+                if not c.guards.get(d) == v.guards.get(d):
+                    continue
 
                 idxf = v.last_idx[c]
                 idxb = mds[(v.xd, idxf)]
@@ -203,7 +209,7 @@ def callback(self, clusters, prefix):
                     guards = c.guards
 
                 properties = c.properties.sequentialize(d)
-                if not isinstance(d, BufferDimension):
+                if not isinstance(d, BufferDimension) and c.guards[d].has(Mod):
                     properties = properties.prefetchable(d)
                 # `c` may be a HaloTouch Cluster, so with no vision of the `bdims`
                 properties = properties.parallelize(v.bdims).affine(v.bdims)
@@ -227,6 +233,8 @@ def callback(self, clusters, prefix):
                     continue
                 if c not in v.lastwrite:
                     continue
+                if not c.guards.get(d) == v.guards.get(d):
+                    continue
 
                 idxf = v.last_idx[c]
                 idxb = mds[(v.xd, idxf)]
@@ -358,15 +366,16 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
     xds = {}
     mapper = {}
     for f, clusters in bfmap.items():
-        exprs = flatten(c.exprs for c in clusters)
+        for k, ck in groupby(clusters, key=lambda c: c.guards):
+            exprs = flatten(c.exprs for c in ck)
 
-        bdims = key(f, exprs)
+            bdims = key(f, exprs)
 
-        dims = [d for d in f.dimensions if d not in bdims]
-        if len(dims) != 1:
-            raise CompilationError(f"Unsupported multi-dimensional `buffering` "
-                                   f"required by `{f}`")
-        dim = dims.pop()
+            dims = [d for d in f.dimensions if d not in bdims]
+            if len(dims) != 1:
+                raise CompilationError(f"Unsupported multi-dimensional `buffering` "
+                                    f"required by `{f}`")
+            dim = dims.pop()
 
         if is_buffering(exprs):
             # Multi-level buffering
@@ -391,25 +400,25 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
                 else:
                     size = async_degree
 
-            # A special CustomDimension to use in place of `dim` in the buffer
-            try:
-                xd = xds[(dim, size)]
-            except KeyError:
-                name = sregistry.make_name(prefix='db')
-                xd = xds[(dim, size)] = BufferDimension(name, 0, size-1, size, dim)
-            extra_kwargs = {}
-
-        # The buffer dimensions
-        dimensions = list(f.dimensions)
-        assert dim in f.dimensions
-        dimensions[dimensions.index(dim)] = xd
-
-        # Finally create the actual buffer
-        cls = callback or Array
-        name = sregistry.make_name(prefix=f'{f.name}b')
-        mapper[f] = cls(name=name, dimensions=dimensions, dtype=f.dtype,
-                        grid=f.grid, halo=f.halo,
-                        space='mapped', mapped=f, f=f, **extra_kwargs)
+                # A special CustomDimension to use in place of `dim` in the buffer
+                try:
+                    xd = xds[(dim, size)]
+                except KeyError:
+                    name = sregistry.make_name(prefix='db')
+                    xd = xds[(dim, size)] = BufferDimension(name, 0, size-1, size, dim)
+                extra_kwargs = {}
+
+            # The buffer dimensions
+            dimensions = list(f.dimensions)
+            assert dim in f.dimensions
+            dimensions[dimensions.index(dim)] = xd
+
+            # Finally create the actual buffer
+            cls = callback or Array
+            name = sregistry.make_name(prefix=f'{f.name}b')
+            mapper[(f, k)] = cls(name=name, dimensions=dimensions, dtype=f.dtype,
+                                 grid=f.grid, halo=f.halo,
+                                 space='mapped', mapped=f, f=f, **extra_kwargs)
 
     return mapper
 
@@ -429,10 +438,11 @@ def map_buffered_functions(clusters, key):
 
 class BufferDescriptor:
 
-    def __init__(self, f, b, clusters):
+    def __init__(self, f, b, clusters, guards):
         self.f = f
         self.b = b
         self.clusters = clusters
+        self.guards = guards
 
         self.xd, = b.find(BufferDimension)
         self.bdims = tuple(d for d in b.dimensions if d is not self.xd)
@@ -673,8 +683,9 @@ def make_mds(descriptors, prefix, sregistry):
         # same strategy is also applied in clusters/algorithms/Stepper
         key = lambda i: -np.inf if i - p == 0 else (i - p)  # noqa: B023
         indices = sorted(v.indices, key=key)
+        v_mds = None
 
-        for i in indices:
+        for k, i in enumerate(indices):
             k = (v.xd, i)
             if k in mds:
                 continue

devito/symbolics/extended_sympy.py

@@ -19,7 +19,6 @@
 )
 from devito.types import Symbol
 from devito.types.basic import Basic
-from devito.types.relational import Ge
 
 __all__ = ['CondEq', 'CondNe', 'BitwiseNot', 'BitwiseXor', 'BitwiseAnd',  # noqa
            'LeftShift', 'RightShift', 'IntDiv', 'Terminal', 'CallFromPointer',
@@ -49,11 +48,6 @@ def canonical(self):
     def negated(self):
         return CondNe(*self.args, evaluate=False)
 
-    @property
-    def _as_min(self):
-        from devito.symbolics.extended_dtypes import INT
-        return INT(Ge(*self.args))
-
 
 class CondNe(sympy.Ne):
 

devito/types/relational.py

@@ -320,4 +320,5 @@ def _(expr, s):
 def _(expr, s):
     if isinstance(expr.lhs, sympy.Mod):
         return 0
-    return expr._as_min
+    from devito.symbolics.extended_dtypes import INT
+    return INT(Ge(*expr.args))

tests/test_buffering.py

@@ -818,3 +818,45 @@ def test_multi_cond_v1():
     for i in range(nt-1):
         assert np.allclose(f.data[i], i*2)
     assert np.allclose(f.data[nt-1], ntmod - 2)
+
+
+@pytest.mark.parametrize("factor", [1, 2, 3])
+def test_buffering_multi_cond(factor):
+    grid = Grid((16, 16))
+
+    nt = 5
+    ntmod = (nt - 1) * factor + 1
+
+    ct0 = ConditionalDimension(name="ct0", parent=grid.time_dim, factor=factor,
+                               relation=Or)
+    f = TimeFunction(grid=grid, name='f', time_order=0, space_order=0,
+                     time_dim=ct0, save=nt)
+    T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
+
+    eqs = []
+    eqs.append(Eq(T, grid.time_dim))
+
+    # conditional dimension for the last sample in the operator
+    ctend = ConditionalDimension(name="ctend", parent=grid.time_dim,
+                                 condition=CondEq(grid.time_dim, ntmod - 2),
+                                 relation=Or)
+
+    eqs.append(Eq(f, T))  # this to save times from 0 to nt - 2
+    # this to save the last time sample nt - 1
+    eqs.append(Eq(f.forward, T+1, implicit_dims=ctend))
+
+    # run operator with serialization
+    op = Operator(eqs, opt='buffering')
+    op.apply(time_m=0, time_M=ntmod-2)
+
+    # Now run backward as well with buffering
+
+    f_all = TimeFunction(grid=grid, name='f_all', time_order=0,
+                         space_order=0, time_dim=ct0, save=nt)
+
+    eq_all = [Eq(f_all, f)]
+    eq_all.append(Eq(f_all.forward, f.forward, implicit_dims=ctend))
+    op_all = Operator(eq_all, opt='buffering')
+    op_all.apply(time_m=0, time_M=ntmod-2)
+
+    assert np.allclose(f_all.data[:, 11, 11], factor * np.arange(nt))