compiler: fix for kernel

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -477,7 +477,15 @@ def callback(self, clusters, prefix, seen=None):
             key = lambda i: key0(i) and key1(i)  # noqa: B023
             ispace = c.ispace.project(key)
 
-            properties = c.properties.sequentialize()
+            # Only sequentialize the Dimensions that actually drive the
+            # halo (its ``loc_indices`` and the Dimensions they define).
+            # Sequentialising *every* outer Dimension would also bring
+            # along independent iterators like ``p_rec`` for an
+            # interpolation, vetoing downstream blocking of the real
+            # cluster the HaloTouch sits alongside.
+            relevant = set().union(*(i._defines for i in hs.loc_indices))
+            seq_dims = [d for d in c.properties if d in relevant]
+            properties = c.properties.sequentialize(seq_dims)
 
             halo_touch = c.rebuild(exprs=expr, ispace=ispace, properties=properties)
 

devito/ir/clusters/cluster.py

@@ -34,7 +34,7 @@ class EqBlock(CacheInstances):
     @classmethod
     def _preprocess_args(cls, exprs, ispace=null_ispace, guards=None,
                          properties=None, syncs=None, halo_scheme=None):
-        exprs = tuple(ClusterizedEq(e, ispace=ispace) for e in as_tuple(exprs))
+        exprs = tuple(clusterize_eq(e, ispace=ispace) for e in as_tuple(exprs))
         guards = Guards(guards or {})
         properties = Properties(properties or {})
         syncs = normalize_syncs(syncs or {})

devito/ir/equations/equation.py

@@ -13,8 +13,6 @@
 from devito.tools import (
     Pickable, Tag, as_hashable, filter_sorted, frozendict, reuse_if_unchanged
 )
-from devito.symbolics import IntDiv, limits_mapper, uxreplace
-from devito.tools import Pickable, Tag, frozendict
 from devito.types import (
     Eq, Inc, IncrInterpolation, Injection, InjectionMixin, Interpolation,
     InterpolationMixin, ReduceMax, ReduceMin, ReduceMinMax, SparseEq, SparseOpMixin,
@@ -208,6 +206,19 @@ def detect(cls, expr):
             (ReduceMin, OpMin),
             (ReduceMax, OpMax),
             (Inc, OpInc),
+            # An ``Interpolation`` looks like a plain ``Eq`` -- ``sf[p_*] =
+            # expr[rp_*]`` -- but the cluster scheduler iterates the rhs
+            # over the radius dims, so values are implicitly summed across
+            # ``rp_*``. Tagging it as ``OpInc`` makes the dependence
+            # analysis treat ``rp_*`` as reduction dims
+            # (``parallel_if_atomic``), which matches the lowered code
+            # (``sumrec += ... ; sf[p_*] = sumrec``) and stops the
+            # blocking heuristic from turning the tiny radius loops into
+            # thread blocks. The actual write-back flavour at ``sf[p_*]``
+            # is keyed off the Eq's class (``is_increment_writeback``) in
+            # ``lower_sparse_ops`` so this tag doesn't accidentally turn
+            # ``Interpolation`` assignments into increments.
+            (InterpolationMixin, OpInc),
         )
         for kls, op in reduction_mapper:
             if isinstance(expr, kls):
@@ -366,12 +377,17 @@ class LoweredSparseEq(SparseOpMixin, LoweredEq):
 
 class LoweredInterpolation(InterpolationMixin, LoweredSparseEq):
     """IR counterpart of ``Interpolation``."""
-    pass
+    # ``sf[p_*] = ...``: the write-back at the sparse position is an
+    # assignment. The Eq is still tagged as a reduction
+    # (``OpInc``/``is_Reduction``) because the rhs is summed over the
+    # radius dims; only the *outer* write-back to ``sf[p_*]`` is plain.
+    is_increment_writeback = False
 
 
 class LoweredIncrInterpolation(InterpolationMixin, LoweredSparseEq):
     """IR counterpart of ``IncrInterpolation``."""
-    pass
+    # ``sf[p_*] += ...``: the user asked for ``interpolate(..., increment=True)``.
+    is_increment_writeback = True
 
 
 class LoweredInjection(InjectionMixin, LoweredSparseEq):
@@ -458,12 +474,12 @@ class ClusterizedSparseEq(SparseOpMixin, ClusterizedEq):
 
 class ClusterizedInterpolation(InterpolationMixin, ClusterizedSparseEq):
     """Frozen counterpart of ``LoweredInterpolation``."""
-    pass
+    is_increment_writeback = False
 
 
 class ClusterizedIncrInterpolation(InterpolationMixin, ClusterizedSparseEq):
     """Frozen counterpart of ``LoweredIncrInterpolation``."""
-    pass
+    is_increment_writeback = True
 
 
 class ClusterizedInjection(InjectionMixin, ClusterizedSparseEq):

devito/ir/stree/algorithms.py

@@ -182,7 +182,12 @@ def preprocess(clusters, options=None, **kwargs):
             continue
 
         else:
-            dims = set(c.ispace.promote(lambda d: d.is_Block).itdims)
+            itdims = set(c.ispace.promote(lambda d: d.is_Block).itdims)
+            # Expand the iteration dims by ``_defines`` so that derived
+            # Dimensions (e.g. ``rp_*`` radius dims with the grid
+            # Dimension as parent) match the halo's distributed indices
+            # of the Dimensions they ultimately iterate over
+            dims = set().union(*[d._defines for d in itdims])
 
             found = []
             for c1 in list(queue):

devito/passes/iet/sparse.py

@@ -28,8 +28,8 @@
 from devito.ir.equations import DummyEq
 from devito.ir.equations.algorithms import lower_exprs
 from devito.ir.iet import (
-    Call, EntryFunction, Expression, FindNodes, HaloSpot, Increment, Iteration, List,
-    Transformer, make_callable
+    Call, Conditional, EntryFunction, Expression, ExpressionBundle, FindNodes, HaloSpot,
+    Increment, Iteration, List, Transformer, make_callable
 )
 from devito.passes.iet.engine import iet_pass
 from devito.types import Eq, InjectionMixin, InterpolationMixin, Symbol
@@ -91,9 +91,22 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
         efuncs.append(efunc)
 
         call = Call(efunc.name, list(efunc.parameters))
+
+        # Any HaloSpot living inside the nest (e.g. around the radius
+        # loops reading a grid Function) becomes invisible once the
+        # nest collapses into an opaque Call. Hoist its halo_scheme out
+        # so the halo exchange still happens at the right point in the
+        # parent IET, sitting next to the Call rather than buried in it.
+        # An injection's lhs is *written* (not read) so its inner halo
+        # entry for that Function is a reduction-only halo that the
+        # caller has nothing to read back -- drop it before hoisting.
+        reduced = {e.expr.lhs.function for e in exprs
+                   if isinstance(e.expr, InjectionMixin)}
+        prelude = _hoisted_halo_prelude(nest, reduced)
+
         parent = parents[nest]
         if parent is None:
-            mapper[nest] = call
+            mapper[nest] = List(body=prelude + (call,)) if prelude else call
             continue
 
         # Drop fields that the (now-opaque) Call only writes/increments,
@@ -120,7 +133,10 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
 def _find_outer_iteration(iet, expr):
     """
     Walk up the IET from ``expr`` and return the outermost Iteration
-    whose ``dim.root`` is the SparseFunction's sparse Dimension.
+    whose ``dim.root`` is the SparseFunction's sparse Dimension. This
+    is the entry point of the sparse-op nest in the parent IET; the
+    full nest (including any block Iterations the cluster pipeline
+    wrapped around the sparse loop) gets extracted into the efunc.
     """
     sparse_dim = expr.expr.interpolator.sfunction._sparse_dim
     for it in FindNodes(Iteration).visit(iet):
@@ -139,13 +155,60 @@ def _enclosing_halospot(iet, nest):
     return None
 
 
+def _hoisted_halo_prelude(nest, reduced=None):
+    """
+    Build a tuple of nodes that recreate, outside the sparse nest, any
+    HaloSpots that live inside it. Each hoisted HaloSpot is wrapped in
+    the Iterations from ``nest`` whose dim is referenced by the halo's
+    ``loc_indices`` (e.g. the ``ps`` loop wrapping a halo whose
+    ``loc_indices = {ps}``), so the lowering pipeline that turns the
+    HaloSpot into a halo-exchange Call still sees the indices it needs
+    in scope.
+
+    ``reduced``, when provided, lists Functions whose halo entries
+    should be dropped from the hoisted scheme -- these are the
+    Functions the sparse Call only writes/increments (e.g. an
+    injection's lhs), so the parent never reads them back and the
+    halo update would be redundant.
+    """
+    reduced = reduced or set()
+
+    inner = []
+    for hs in FindNodes(HaloSpot).visit(nest):
+        scheme = hs.halo_scheme.drop(reduced) if reduced else hs.halo_scheme
+        if not scheme.is_void:
+            inner.append(scheme)
+    if not inner:
+        return ()
+
+    iters = FindNodes(Iteration).visit(nest)
+    prelude = []
+    for scheme in inner:
+        loc_dims = set().union(*(d._defines for d in scheme.loc_indices))
+        wrappers = [it for it in iters if it.dim in loc_dims]
+        body = HaloSpot(List(body=[]), scheme)
+        for it in reversed(wrappers):
+            body = it._rebuild(nodes=body)
+        prelude.append(body)
+    return tuple(prelude)
+
+
 def _materialise_nest(nest, exprs):
     """
     Rewrite the sparse Dimension's Iteration body to compute the
     position/coefficient temps once per sparse point, then for any
     interpolation Expression wrap it with the scalar accumulator
     pattern. Multiple sparse-op Expressions sharing the same outer
     Iteration are materialised in one pass and reuse the same temps.
+
+    ``nest`` is the *outermost* sparse-Dimension Iteration, so that the
+    whole block-Iteration hierarchy (e.g. ``p_rec0_blk0`` -> ``p_rec``
+    on the GPU pipeline) is extracted into the efunc and downstream GPU
+    kernel synthesis can fold the block loops into a thread-grid
+    wrapping the kernel body. The temps and the accumulator pattern,
+    however, must live *inside* the innermost sparse Iteration -- one
+    set per sparse point, sitting beneath any thread-index/OOB-guard
+    prelude that the GPU kernel prep may have inserted.
     """
     # Position + coefficient temporaries as IET Expressions. These are
     # the same for every Expression in the group, so we emit them once.
@@ -155,24 +218,63 @@ def _materialise_nest(nest, exprs):
     temp_exprs = tuple(Expression(DummyEq(e.lhs, e.rhs))
                        for e in lower_exprs(sample.sparse_temps()))
 
-    # The radius nest is what runs once per sparse point. For each
-    # interpolation Expression in the group, build its
-    # accumulator-wrapped copy of the radius nest. Injection Exprs
-    # share a single copy of the radius nest (their ``Inc`` already
-    # carries the right ``weights * rhs`` form).
-    inner = nest.nodes[0] if len(nest.nodes) == 1 else List(body=nest.nodes)
+    # Find the innermost sparse-Dimension Iteration within ``nest`` --
+    # that's where the head Expressions actually live, beneath any block
+    # Iterations that the cluster pipeline wrapped around the sparse
+    # loop.
+    sparse_dim = sample.interpolator.sfunction._sparse_dim
+    inner_iter = nest
+    for it in FindNodes(Iteration).visit(nest):
+        if it.dim.root is sparse_dim and \
+                any(e in FindNodes(Expression).visit(it) for e in exprs):
+            inner_iter = it
+
+    # ``inner_iter`` may carry a GPU kernel prelude (thread-index
+    # ``ExpressionBundle`` and OOB ``Conditional``) that downstream
+    # kernel synthesis expects to find at the top of the block dim's
+    # body. The temps and the accumulator pattern go *after* that
+    # prelude.
+    head, body_nodes = _split_kernel_prelude(inner_iter.nodes)
+
+    radius_nest = body_nodes[0] if len(body_nodes) == 1 else List(body=body_nodes)
     interp_exprs = [e for e in exprs if isinstance(e.expr, InterpolationMixin)]
     inject_exprs = [e for e in exprs if isinstance(e.expr, InjectionMixin)]
 
-    body = []
+    new_body = []
     for expr in interp_exprs:
         siblings = [e for e in exprs if e is not expr]
-        body.append(_interp_inner_block(inner, expr, expr.expr.interpolator, siblings))
+        new_body.append(_interp_inner_block(
+            radius_nest, expr, expr.expr.interpolator, siblings))
     if inject_exprs:
         drop = {e: None for e in interp_exprs}
-        body.append(Transformer(drop, nested=True).visit(inner) if drop else inner)
+        new_body.append(Transformer(drop, nested=True).visit(radius_nest)
+                        if drop else radius_nest)
+
+    new_inner_iter = inner_iter._rebuild(
+        nodes=head + temp_exprs + tuple(new_body)
+    )
+    if new_inner_iter is inner_iter:
+        return nest
+    return Transformer({inner_iter: new_inner_iter}, nested=True).visit(nest)
 
-    return nest._rebuild(nodes=temp_exprs + tuple(body))
+
+def _split_kernel_prelude(nodes):
+    """
+    Split the contents of a sparse-Dimension Iteration into the GPU
+    kernel prelude (the thread-index ``ExpressionBundle`` and the
+    optional OOB ``Conditional``) and the remaining body. On non-cuda
+    pipelines the prelude is empty and the full ``nodes`` tuple is the
+    body.
+    """
+    head = ()
+    body = tuple(nodes)
+    if body and isinstance(body[0], ExpressionBundle):
+        head += (body[0],)
+        body = body[1:]
+        if body and isinstance(body[0], Conditional):
+            head += (body[0],)
+            body = body[1:]
+    return head, body
 
 
 def _interp_inner_block(inner, expr, interp, siblings):
@@ -219,10 +321,15 @@ def _interp_inner_block(inner, expr, interp, siblings):
 
     init = Expression(DummyEq(acc, 0))
     inc = Increment(DummyEq(acc, weighted_rhs))
-    # Honour the synthetic Eq's flavour: a SparseInc means the user
-    # asked for ``sf[p_*] += sum`` (interpolation with ``increment=True``);
-    # otherwise the standard write is ``sf[p_*] = sum``.
-    write_back_cls = Increment if eq.is_Reduction else Expression
+    # Honour the synthetic Eq's flavour: an ``IncrInterpolation`` means
+    # the user asked for ``sf[p_*] += sum`` (interpolation with
+    # ``increment=True``); a plain ``Interpolation`` is just ``sf[p_*] =
+    # sum``. We key off the leaf class' ``is_increment_writeback`` flag
+    # rather than ``is_Reduction`` because both flavours are tagged as
+    # reductions (``OpInc``) for dependence-analysis purposes -- the rhs
+    # is implicitly summed over the radius dims -- but only the
+    # ``IncrInterpolation`` flavour writes back with ``+=``.
+    write_back_cls = Increment if eq.is_increment_writeback else Expression
     write_back = write_back_cls(DummyEq(sf_lhs, acc))
 
     # Single substitution: drop siblings, replace ``expr`` with ``inc``.

devito/types/sparse.py

@@ -415,12 +415,17 @@ def dist_origin(self):
     @memoized_meth
     def _crdim(self, dim):
         """
-        The CustomDimension associated with the Dimension `dim` for
-        the radius of the interpolation/injection stencil
+        The CustomDimension associated with the grid Dimension ``dim``
+        for the radius of the interpolation/injection stencil. The
+        parent is ``dim`` itself so that ``_defines`` traces back to the
+        grid Dimension the radius slides over -- this is what dependence
+        analysis needs to recognise the implicit reduction over ``rp_*``
+        rather than treating ``rp_*`` as if they were derived from the
+        SparseFunction's sparse Dimension.
         """
         sname = self._sparse_dim.name
         return CustomDimension(f"r{sname}{dim.name}", -self.r+1,
-                               self.r, 2*self.r, self._sparse_dim)
+                               self.r, 2*self.r, dim)
 
     @memoized_meth
     def _cond_rdim(self, dim, cond):

tests/test_mpi.py

@@ -3474,15 +3474,19 @@ def test_issue_2448_v2(self, mode, setup):
 
         calls = [i for i in FindNodes(Call).visit(op2) if isinstance(i, HaloUpdateCall)]
 
-        assert len(calls) == 5
-        assert len(FindNodes(HaloUpdateCall).visit(body0(op2).body[1].body[0])) == 2
-        assert len(FindNodes(HaloUpdateCall).visit(body0(op2).body[1].body[1])) == 3
+        # The sparse-op efunc lowering hoists the interpolation halos so they
+        # share with the upstream stencil halo for the same Function. The
+        # initial-step halo merges ``v`` with ``v2``; the in-loop halo merges
+        # the same pair, leaving 3 distinct HaloUpdateCalls in total.
+        assert len(calls) == 3
+        assert len(FindNodes(HaloUpdateCall).visit(body0(op2).body[1].body[0])) == 1
+        assert len(FindNodes(HaloUpdateCall).visit(body0(op2).body[1].body[1])) == 2
         assert calls[0].arguments[0] is v
-        assert calls[1].arguments[0] is v2
-        assert calls[2].arguments[0] is tau
-        assert calls[2].arguments[1] is tau2
-        assert calls[3].arguments[0] is v
-        assert calls[4].arguments[0] is v2
+        assert calls[0].arguments[1] is v2
+        assert calls[1].arguments[0] is tau
+        assert calls[1].arguments[1] is tau2
+        assert calls[2].arguments[0] is v
+        assert calls[2].arguments[1] is v2
 
     @pytest.mark.parallel(mode=1)
     def test_issue_2448_v3(self, mode, setup):
@@ -3509,14 +3513,17 @@ def test_issue_2448_v3(self, mode, setup):
 
         calls = [i for i in FindNodes(Call).visit(op3) if isinstance(i, HaloUpdateCall)]
 
-        assert len(calls) == 4
+        # The interpolation halo for ``v2.forward`` merges with the in-loop
+        # halo for the same Function, so the per-loop tail collapses to a
+        # single merged ``v``/``v2`` call instead of two separate ones.
+        assert len(calls) == 3
         assert len(FindNodes(HaloUpdateCall).visit(body0(op3).body[1].body[0])) == 1
-        assert len(FindNodes(HaloUpdateCall).visit(body0(op3).body[1].body[1])) == 3
+        assert len(FindNodes(HaloUpdateCall).visit(body0(op3).body[1].body[1])) == 2
         assert calls[0].arguments[0] is v
         assert calls[1].arguments[0] is tau
         assert calls[1].arguments[1] is tau2
         assert calls[2].arguments[0] is v
-        assert calls[3].arguments[0] is v2
+        assert calls[2].arguments[1] is v2
 
     @pytest.mark.parallel(mode=1)
     def test_issue_2448_backward(self, mode):