compiler: Lift lower_sparse_ops out of mpiize into _lower_iet

Sparse-op lowering belongs at the IET-pass level alongside other
target-independent structural lowerings, not buried inside the MPI
pass. Lifting it also fixes a missed host-device transfer on GPU:
the position/coefficient temps it materialises read the coordinate
SubFunctions, but those reads were invisible to place_transfers
because Graph.data_movs was snapshotted before mpiize ran.

Graph.data_movs is now a property recomputed from the live efunc
set so downstream passes see ExprStmts introduced after Graph
construction.

Author: mloubout

devito/operator/operator.py

@@ -33,7 +33,7 @@
 from devito.parameters import configuration
 from devito.passes import (
     Graph, error_mapper, generate_implicit, generate_macros, is_on_device, lower_dtypes,
-    lower_index_derivatives, minimize_symbols, optimize_pows, unevaluate
+    lower_index_derivatives, lower_sparse_ops, minimize_symbols, optimize_pows, unevaluate
 )
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (
@@ -490,13 +490,12 @@ def _lower_iet(cls, uiet, **kwargs):
         parameters = derive_parameters(uiet, True)
         iet = EntryFunction(name, uiet, 'int', parameters, ())
 
-        # Lower IET to a target-specific IET. Sparse-op lowering
-        # (``lower_sparse_ops``) is now run from inside ``mpiize``,
-        # between the halo-optimisation phase and the halo-exchange
-        # injection, so the reduction heuristic gets a chance to drop
-        # redundant halo exchanges around the sparse-op nest before
-        # the nest is sealed into an ElementalFunction.
+        # Materialise the sparse-op iteration nests into ElementalFunctions
+        # before target specialisation, so the position/coefficient temps
+        # the IET pass emits are visible to downstream passes (e.g. the
+        # data-transfer placement on device).
         graph = Graph(iet, **kwargs)
+        lower_sparse_ops(graph, **kwargs)
         graph = cls._specialize_iet(graph, **kwargs)
 
         # Instrument the IET for C-level profiling

devito/passes/iet/engine.py

@@ -81,16 +81,29 @@ def __init__(self, iet, options=None, sregistry=None, **kwargs):
         writes = FindSymbols('writes').visit(iet)
         self.writes_input = frozenset(f for f in writes if f.is_Input)
 
-        # All symbols requiring host-device data transfers when running
-        # on device
-        self.data_movs = rmovs, wmovs = set(), set()
-        gpu_fit = (options or {}).get('gpu-fit', ())
-        for i in FindNodes(ExprStmt).visit(iet):
-            wmovs.update({w for w in i.writes
-                          if needs_transfer(w, gpu_fit)})
-        for i in FindNodes(ExprStmt).visit(iet):
-            rmovs.update({f for f in i.functions
-                          if needs_transfer(f, gpu_fit) and f not in wmovs})
+        self._gpu_fit = (options or {}).get('gpu-fit', ())
+
+    @property
+    def data_movs(self):
+        """
+        ``(reads, writes)`` of symbols requiring host-device data transfers.
+
+        Recomputed on access from the current state of the Graph (root
+        plus every reachable efunc) so passes that introduce ExprStmts
+        after construction — e.g. ``lower_sparse_ops`` materialising the
+        sparse-op position temps inside an efunc — are reflected.
+        """
+        rmovs, wmovs = set(), set()
+        for efunc in self.efuncs.values():
+            for i in FindNodes(ExprStmt).visit(efunc):
+                wmovs.update(w for w in i.writes
+                             if needs_transfer(w, self._gpu_fit))
+        for efunc in self.efuncs.values():
+            for i in FindNodes(ExprStmt).visit(efunc):
+                rmovs.update(f for f in i.functions
+                             if needs_transfer(f, self._gpu_fit)
+                             and f not in wmovs)
+        return rmovs, wmovs
 
     @property
     def root(self):

devito/passes/iet/mpi.py

@@ -11,7 +11,6 @@
 from devito.mpi.reduction_scheme import DistReduce
 from devito.mpi.routines import HaloExchangeBuilder, ReductionBuilder
 from devito.passes.iet.engine import iet_pass
-from devito.passes.iet.sparse import lower_sparse_ops
 from devito.symbolics import VectorAccess, search
 from devito.tools import generator
 from devito.types import TensorMove
@@ -390,11 +389,6 @@ def mpiize(graph, **kwargs):
     Perform three IET passes:
 
         * Optimization of halo exchanges
-        * Lower sparse operations (Interpolation/Injection) into Calls
-          to ElementalFunctions. This runs after halo optimisation so
-          the reduction heuristic gets a chance to drop redundant halo
-          exchanges around the sparse-op nest before it is sealed into
-          an efunc.
         * Injection of code for halo exchanges
         * Injection of code for reductions
     """
@@ -403,8 +397,6 @@ def mpiize(graph, **kwargs):
     if options['opt-comms']:
         optimize_halospots(graph, **kwargs)
 
-    lower_sparse_ops(graph, **kwargs)
-
     mpimode = options['mpi']
     if mpimode:
         make_halo_exchanges(graph, mpimode=mpimode, **kwargs)