compiler: Shed reduction-only halos when lowering injections

After lifting lower_sparse_ops out of mpiize, an injection nest is
turned into a Call before optimize_halospots runs, so
_drop_reduction_halospots can no longer detect that the wrapping
HaloSpot's entry for the injected field is reduction-only. The stale
entry was left in place, and on save=True (no modulo buffering) the
hoist pass propagated the loop iteration variable out of the time
loop, producing an undeclared 'time' reference. Drop those entries at
lowering time so the resulting HaloSpot only carries entries with a
genuine read at the IET level.

Author: mloubout

devito/passes/iet/sparse.py

@@ -69,13 +69,16 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
             continue
         groups.setdefault(nest, []).append(expr)
 
-    # If a sparse-op nest sits inside a HaloSpot whose halo scheme is
-    # void (the reduction-only halo got dropped by
-    # ``_drop_reduction_halospots``), replace the HaloSpot rather than
-    # just the nest so we don't leave behind an empty HaloSpot — the
-    # MPI overlap machinery would otherwise try to wrap our Call with
-    # its own dynamic-args plumbing.
-    parents = {nest: _enclosing_void_halospot(iet, nest) for nest in groups}
+    # ``lower_sparse_ops`` runs before ``optimize_halospots``, so the
+    # halo-exchange optimiser hasn't yet had a chance to drop the
+    # reduction-only halo entries that the IR scheduler put around an
+    # injection nest (e.g. an entry for ``u`` at ``loc_indices={time:
+    # time+1}`` wrapping ``u[time+1] += ...``). Once the nest becomes a
+    # Call those expressions are no longer visible to
+    # ``_drop_reduction_halospots``, so we shed those entries here -- and
+    # if that empties the HaloSpot, replace it whole so the MPI overlap
+    # machinery doesn't wrap our Call with stale dynamic-args plumbing.
+    parents = {nest: _enclosing_halospot(iet, nest) for nest in groups}
 
     mapper = {}
     efuncs = []
@@ -87,7 +90,26 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
         efunc = make_callable(sregistry.make_name(prefix=prefix), new_nest)
         efuncs.append(efunc)
 
-        mapper[parents[nest] or nest] = Call(efunc.name, list(efunc.parameters))
+        call = Call(efunc.name, list(efunc.parameters))
+        parent = parents[nest]
+        if parent is None:
+            mapper[nest] = call
+            continue
+
+        # Drop fields that the (now-opaque) Call only writes/increments,
+        # since the wrapping HaloSpot's purpose was to ensure read-side
+        # coherency for them and the read no longer exists at the IET
+        # level. Interpolation reads its target field, so its entries
+        # stay.
+        reduced = {e.expr.lhs.function for e in exprs
+                   if isinstance(e.expr, InjectionMixin)}
+        hs = parent.halo_scheme.drop(reduced) if reduced else parent.halo_scheme
+        if hs.is_void:
+            mapper[parent] = call
+        elif hs is parent.halo_scheme:
+            mapper[nest] = call
+        else:
+            mapper[parent] = parent._rebuild(halo_scheme=hs, body=call)
 
     if not mapper:
         return iet, {}
@@ -107,14 +129,12 @@ def _find_outer_iteration(iet, expr):
     return None
 
 
-def _enclosing_void_halospot(iet, nest):
+def _enclosing_halospot(iet, nest):
     """
-    Return the HaloSpot directly wrapping ``nest`` if it carries an
-    empty (void) HaloScheme, otherwise None. Such HaloSpots are leftover
-    after ``_drop_reduction_halospots`` cleared all entries.
+    Return the HaloSpot directly wrapping ``nest``, if any.
     """
     for hs in FindNodes(HaloSpot).visit(iet):
-        if hs.is_void and nest in FindNodes(Iteration).visit(hs):
+        if nest in FindNodes(Iteration).visit(hs):
             return hs
     return None
 

tests/test_dse.py

@@ -2660,7 +2660,9 @@ def test_sparse_const(self):
 
         op = Operator(src.interpolate(u))
 
-        cond = FindNodes(Conditional).visit(op)
+        assert len(FindNodes(Conditional).visit(op)) == 0
+        print(op._func_table)
+        cond = FindNodes(Conditional).visit(op._func_table['interpolate_src0'])
         assert len(cond) == 1
         assert len(cond[0].args['then_body'][0].exprs) == 1
         assert all(e.is_scalar for e in cond[0].args['then_body'][0].exprs)
@@ -2914,12 +2916,12 @@ def test_fullopt(self):
         bns, _ = assert_blocking(op1, {'x0_blk0'})  # due to loop blocking
 
         assert summary0[('section0', None)].ops == 55
-        assert summary0[('section1', None)].ops == 44
+        assert summary0[('section1', None)].ops == 17
         assert np.isclose(summary0[('section0', None)].oi, 3.136, atol=0.001)
 
         assert summary1[('section0', None)].ops == 31
-        assert summary1[('section1', None)].ops == 88
-        assert summary1[('section2', None)].ops == 25
+        assert summary1[('section1', None)].ops == 17
+        assert summary1[('section2', None)].ops == 0
         assert np.isclose(summary1[('section0', None)].oi, 1.767, atol=0.001)
 
         assert np.allclose(u0.data, u1.data, atol=10e-5)
@@ -2966,7 +2968,8 @@ def tti_noopt(self):
 
         # Make sure no opts were applied
         op = wavesolver.op_fwd(False)
-        assert len(op._func_table) == 0
+        # Two funcs, one for src, one for rec
+        assert len(op._func_table) == 2
         assert summary[('section0', None)].ops == 753
 
         return v, rec
@@ -3024,7 +3027,7 @@ def test_fullopt_w_mpi(self, mode):
 
         # Run a quick check to be sure MPI-full-mode code was actually generated
         op = tti_agg.op_fwd(False)
-        assert len(op._func_table) == 7
+        assert len(op._func_table) == 9
         assert 'pokempi0' in op._func_table
 
     @switchconfig(profiling='advanced')