compiler: Block strided subsampling and align SubGrids to global indices

devito/ir/equations/algorithms.py

@@ -156,6 +156,8 @@ def _lower_exprs(expressions, subs):
         # Add dimensions map to the mapper in case dimensions are used
         # as an expression, i.e. Eq(u, x, subdomain=xleft)
         mapper.update(dimension_map)
+        # Dimensions used as an expression are global indices, unlike loop indices
+        mapper.update(_make_global_index_map(expr, dimension_map))
         # Add the user-supplied substitutions
         mapper.update(subs)
         # Apply mapper to expression
@@ -184,14 +186,39 @@ def _make_dimension_map(expr):
         # Some Relationals may be pure SymPy objects, thus lacking the SubDomain
         dimension_map = {}
     else:
-        functions = [f for f in retrieve_functions(expr) if f._is_on_subdomain]
+        functions = retrieve_functions(expr)
         for f in functions:
-            dimension_map.update({d: expr.subdomain.dimension_map[d.root]
-                                  for d in f.space_dimensions if d.is_Sub})
+            if f._is_on_subdomain:
+                dimension_map.update({d: expr.subdomain.dimension_map[d.root]
+                                      for d in f.space_dimensions if d.is_Sub})
+
+        # A subsampled Function, e.g. on a SubGrid, is subsampled from the
+        # SubDomain's Dimensions rather than from the whole Grid's
+        for f in functions:
+            dimension_map.update({
+                d: d._rebuild(parent=expr.subdomain.dimension_map[d.parent])
+                for d in f.dimensions
+                if d.is_Conditional and d.parent in expr.subdomain.dimension_map
+            })
 
     return frozendict(dimension_map)
 
 
+def _make_global_index_map(expr, dimension_map):
+    """
+    Map the Dimensions of a distributed Grid to the global indices they stand for
+    when used as an expression, i.e. Eq(u, x + y) -> Eq(u[x, y], x + io_x + y + io_y).
+    Dimensions indexing an Indexed are unaffected as the Indexed is replaced as a whole.
+    """
+    if not isinstance(expr, Eq):
+        return {}
+
+    grids = {f.grid.root for f in retrieve_functions(expr) if f.grid is not None}
+    return {d: dimension_map.get(d, d) + o
+            for grid in grids if grid.distributor.is_parallel
+            for d, o in zip(grid.dimensions, grid.origin_isymbols, strict=True)}
+
+
 def concretize_subdims(exprs, **kwargs):
     """
     Given a list of expressions, return a new list where all user-defined
@@ -390,7 +417,7 @@ def generate_conditionals(expr, input_expr, ordering):
 
     # Replace the ConditionalDimensions in `expr`
     for d, cond in conditionals.items():
-        index = d.index
+        index = d.index - d.offset
         if d.condition is not None and expr.has(d):
             index = index - relational_min(cond, d.parent)
         shift = relational_shift(cond, d.parent, d.symbolic_factor)

devito/ir/support/guards.py

@@ -61,7 +61,7 @@ class GuardFactor(Guard, CondEq, Pickable):
     A guard for factor-based ConditionalDimensions.
 
     Given the ConditionalDimension `d` with factor `k`, create the
-    symbolic relational `d.parent % k == 0`.
+    symbolic relational `(d.parent - d.offset) % k == 0`.
     """
 
     __rargs__ = ('d',)
@@ -76,7 +76,7 @@ def __new__(cls, *args, **kwargs):
         d, = args
         assert d.is_Conditional
 
-        obj = super().__new__(cls, d.parent % d.symbolic_factor, 0)
+        obj = super().__new__(cls, (d.parent - d.offset) % d.symbolic_factor, 0)
         obj.d = d
 
         return obj

devito/mpi/distributed.py

@@ -68,6 +68,7 @@ def __getattr__(self, name):
     'Distributor',
     'SparseDistributor',
     'SubDistributor',
+    'SubGridDistributor',
     'devito_mpi_finalize',
     'devito_mpi_init',
 ]
@@ -375,10 +376,15 @@ def __init__(self, shape, dimensions, input_comm=None, topology=None):
             self._comm = MPI.COMM_NULL
             self._topology = tuple(1 for _ in range(len(shape)))
 
-        # The domain decomposition
+        self.__decomposition_setup__()
+
+    def __decomposition_setup__(self):
+        """
+        Set up the domain decomposition.
+        """
         self._decomposition = [
             Decomposition(np.array_split(range(i), j), c)
-            for i, j, c in zip(shape, self.topology, self.mycoords, strict=True)
+            for i, j, c in zip(self.glb_shape, self.topology, self.mycoords, strict=True)
         ]
 
     @cached_property
@@ -485,6 +491,41 @@ def _rebuild(self, shape=None, dimensions=None, comm=None):
                            comm or self.comm)
 
 
+class SubGridDistributor(Distributor):
+
+    """
+    Decompose a SubGrid over a set of MPI processes in a manner consistent with
+    the parent distributor: a rank owns the SubGrid points on its parent points.
+
+    Parameters
+    ----------
+    subgrid : SubGrid
+        The SubGrid to be decomposed.
+    """
+
+    def __init__(self, subgrid):
+        self._parent = subgrid.parent.distributor
+        self._factors = subgrid.factor
+
+        super().__init__(subgrid.shape, subgrid.dimensions, self.parent.comm,
+                         tuple(self.parent.topology))
+
+    def __decomposition_setup__(self):
+        """
+        Set up the decomposition, aligned with that of the parent Distributor.
+        """
+        self._decomposition = [
+            Decomposition([d[d % f == 0] // f for d in dec], c)
+            for dec, f, c in zip(self.parent.decomposition, self._factors,
+                                 self.mycoords, strict=True)
+        ]
+
+    @property
+    def parent(self):
+        """The parent distributor of this SubGridDistributor."""
+        return self._parent
+
+
 class SubDistributor(DenseDistributor):
 
     """

devito/operator/operator.py

@@ -43,7 +43,7 @@
     split, timed_pass, timed_region
 )
 from devito.types import Buffer, Evaluable, device_layer, disk_layer, host_layer
-from devito.types.dimension import SubDimension, Thickness
+from devito.types.dimension import SubDimension, SubsamplingOffset, Thickness
 from devito.warnings import warn
 
 __all__ = ['Operator']
@@ -415,6 +415,11 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
         # Operation count before specialization
         init_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense)
 
+        # Strengthen factor-based ConditionalDimension guards into strided
+        # iteration. Before specialization, so that blocking tiles the swapped
+        # Dimensions
+        clusters = subsampling(clusters)
+
         clusters = cls._specialize_clusters(clusters, **kwargs)
 
         # Operation count after specialization
@@ -425,11 +430,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
         # Generate implicit Clusters from higher level abstractions
         clusters = generate_implicit(clusters)
 
-        # Strengthen factor-based ConditionalDimension guards into strided
-        # iteration. Before `lower_index_derivatives` so that the Dimension
-        # swap is seen by everything downstream
-        clusters = subsampling(clusters)
-
         # Lower all remaining high order symbolic objects
         clusters = lower_local_sums(clusters, **kwargs)
         clusters = lower_index_derivatives(clusters, **kwargs)
@@ -695,15 +695,19 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
             # No need to do anything more if only checking the memory
             return args
 
+        # Process the symbols localized by MPI. A SubDimension's thicknesses are
+        # needed even when they only bound another Dimension, e.g. a subsampled one
+        localized = OrderedSet(*[p for p in self.parameters
+                                 if isinstance(p, (SubsamplingOffset, Thickness))])
+        localized.update(t for d in toposort if isinstance(d, SubDimension)
+                         for t in d.thickness)
+        for p in localized:
+            args.update(p._arg_values(grid=grid, **kwargs))
+
         # Process Dimensions
         for d in reversed(toposort):
             args.update(d._arg_values(self._dspace[d], grid, **kwargs))
 
-        # Process Thicknesses
-        for p in self.parameters:
-            if isinstance(p, Thickness):
-                args.update(p._arg_values(grid=grid, **kwargs))
-
         # Process Objects
         for o in self.objects:
             args.update(o._arg_values(grid=grid, **kwargs))

devito/passes/clusters/subsampling.py

@@ -20,8 +20,8 @@ def subsampling(clusters):
           if (x%xf == 0)                 ->     g[x/xf] = ...
             g[x/xf] = ...
 
-The loop is taken over the subsampled Dimension `cd` itself, with a *unit*
-    stride, and every access to the parent `d` is rescaled to `cd*factor`:
+    The loop is taken over the subsampled Dimension `cd` itself, with a *unit*
+    stride, and every access to the parent `d` is mapped to `cd*factor + offset`:
 
         for (x = x_m; x <= x_M; x += 1)      for (xs = xs_m; xs <= xs_M; xs++)
           if (x%xf == 0)                 ->     g[xs] = u[xs*xf]
@@ -30,8 +30,8 @@ def subsampling(clusters):
     A non-unitary increment is deliberately avoided: it would escape the
     Lamport-style dependence analysis the compiler relies on. The loop bounds
     are `cd`'s own runtime values, which `ConditionalDimension._arg_values`
-    derives as `ceil(x_m/factor) .. floor(x_M/factor)`, so no anchoring of the
-    loop minimum is needed either.
+    derives by inverting that map, so no anchoring of the loop minimum is needed
+    either.
 
     Parameters
     ----------
@@ -73,17 +73,18 @@ def callback(self, clusters, prefix):
 
         cd = guard.d
         factor = cd.symbolic_factor
+        offset = cd.offset
 
         processed = []
         for c in clusters:
             # Iterate the subsampled Dimension itself, with unit stride, and
-            # rescale every parent access `d` to `cd*factor`
+            # map every parent access `d` to `cd*factor + offset`
             ispace = c.ispace.switch(d, cd)
 
-            # `IntDiv(d, factor)` is how the subsampled access was lowered;
-            # with `d -> cd*factor` it collapses to `cd`, so rewrite it
-            # directly rather than leaving `(cd*factor)/factor` behind
-            subs = {IntDiv(d, factor): cd, d: cd*factor}
+            # `IntDiv(d - offset, factor)` is how the subsampled access was
+            # lowered; with `d -> cd*factor + offset` it collapses to `cd`, so
+            # rewrite it directly rather than leaving the division behind
+            subs = {IntDiv(d - offset, factor): cd, d: cd*factor + offset}
             exprs = [uxreplace(e, subs) for e in c.exprs]
 
             # `cd` inherits `d`'s computational properties: it covers the very
@@ -112,6 +113,10 @@ def _is_strideable(cluster, d, guard):
     if d not in cluster.ispace.itdims:
         return False
 
+    # Time keeps its guard, which later passes (e.g. buffering) rely upon
+    if d.is_Time:
+        return False
+
     # Skipping points is only sound if no dependence is carried along `d`;
     # otherwise the discarded iterations may be doing work that `d`'s
     # remaining iterations rely upon