mpi: fix allreduce iteration space

.github/workflows/pytest-core-nompi.yml

@@ -38,7 +38,8 @@ jobs:
            pytest-ubuntu-py310-gcc9-omp,
            pytest-osx-py312-clang-omp,
            pytest-docker-py310-gcc-omp,
-           pytest-docker-py310-icx-omp
+           pytest-docker-py310-icx-omp,
+           pytest-ubuntu-py313-gcc14-omp
         ]
         set: [base, adjoint]
         include:
@@ -105,6 +106,13 @@ jobs:
           language: "openmp"
           sympy: "1.13"
 
+        - name: pytest-ubuntu-py313-gcc14-omp
+          python-version: '3.13'
+          os: ubuntu-24.04
+          arch: "gcc-14"
+          language: "openmp"
+          sympy: "1.14"
+
         - set: base
           test-set: 'not adjoint'
 
@@ -155,8 +163,11 @@ jobs:
 
     - name: Set pip flags for latest python (3.12)
       run: |
-          if [ "${{ matrix.python-version }}" == '3.12' ]; then
-              echo "PIPFLAGS='--break-system-packages'" >> $GITHUB_ENV
+          ver="${{ matrix.python-version }}"
+          major=${ver%%.*}
+          minor=${ver#*.}; minor=${minor%%.*}
+          if [ "$major" -eq 3 ] && [ "$minor" -ge 12 ]; then
+            echo "PIPFLAGS='--break-system-packages'" >> $GITHUB_ENV
           fi
 
     - name: Install dependencies

devito/ir/clusters/algorithms.py

@@ -8,7 +8,7 @@
 from devito.exceptions import CompilationError
 from devito.finite_differences.elementary import Max, Min
 from devito.ir.support import (Any, Backward, Forward, IterationSpace, erange,
-                               pull_dims, null_ispace)
+                               pull_dims)
 from devito.ir.equations import OpMin, OpMax, identity_mapper
 from devito.ir.clusters.analysis import analyze
 from devito.ir.clusters.cluster import Cluster, ClusterGroup
@@ -493,9 +493,9 @@ def reduction_comms(clusters):
         processed.append(c)
 
     # Leftover reductions are placed at the very end
-    if fifo:
-        exprs = [Eq(dr.var, dr) for dr in fifo]
-        processed.append(Cluster(exprs=exprs, ispace=null_ispace))
+    for ispace, reds in groupby(fifo, key=lambda r: r.ispace):
+        exprs = [Eq(dr.var, dr) for dr in reds]
+        processed.append(Cluster(exprs=exprs, ispace=ispace))
 
     return processed
 

devito/ir/iet/nodes.py

@@ -407,7 +407,10 @@ def reads(self):
     @cached_property
     def write(self):
         """The Function written by the Expression."""
-        return self.expr.lhs.base.function
+        try:
+            return self.expr.rhs.writes[0]
+        except (AttributeError, IndexError):
+            return self.expr.lhs.base.function
 
     @cached_property
     def dimensions(self):

devito/ir/iet/visitors.py

@@ -1295,6 +1295,28 @@ def __init__(self, mapper, nested=False):
         self.mapper = mapper
         self.nested = nested
 
+    def transform(self, o, handle, **kwargs):
+        if handle is None:
+            # None -> drop `o`
+            return None
+        elif isinstance(handle, Iterable):
+            # Iterable -> inject `handle` into `o`'s children
+            if not o.children:
+                raise CompilationError("Cannot inject nodes in a leaf node")
+            if self.nested:
+                children = [self._visit(i, **kwargs) for i in o.children]
+            else:
+                children = o.children
+            children = (tuple(handle) + children[0],) + tuple(children[1:])
+            return o._rebuild(*children, **o.args_frozen)
+        else:
+            # Replace `o` with `handle`
+            if self.nested:
+                children = [self._visit(i, **kwargs) for i in handle.children]
+                return handle._rebuild(*children, **handle.args_frozen)
+            else:
+                return handle
+
     def visit_object(self, o, **kwargs):
         return o
 
@@ -1307,29 +1329,11 @@ def visit_tuple(self, o, **kwargs):
     def visit_Node(self, o, **kwargs):
         if o in self.mapper:
             handle = self.mapper[o]
-            if handle is None:
-                # None -> drop `o`
-                return None
-            elif isinstance(handle, Iterable):
-                # Iterable -> inject `handle` into `o`'s children
-                if not o.children:
-                    raise CompilationError("Cannot inject nodes in a leaf node")
-                if self.nested:
-                    children = [self._visit(i, **kwargs) for i in o.children]
-                else:
-                    children = o.children
-                children = (tuple(handle) + children[0],) + tuple(children[1:])
-                return o._rebuild(*children, **o.args_frozen)
-            else:
-                # Replace `o` with `handle`
-                if self.nested:
-                    children = [self._visit(i, **kwargs) for i in handle.children]
-                    return handle._rebuild(*children, **handle.args_frozen)
-                else:
-                    return handle
-        else:
-            children = [self._visit(i, **kwargs) for i in o.children]
-            return o._rebuild(*children, **o.args_frozen)
+            return self.transform(o, handle, **kwargs)
+        children = [self._visit(i, **kwargs) for i in o.children]
+        if o._traversable and not any(children):
+            return None
+        return o._rebuild(*children, **o.args_frozen)
 
     def visit_Operator(self, o, **kwargs):
         raise ValueError("Cannot apply a Transformer visitor to an Operator directly")

pyproject.toml

@@ -12,7 +12,7 @@ description = "Finite Difference DSL for symbolic computation."
 license = { file = "LICENSE.md" }
 readme = "README.md"
 keywords = ["finite-difference", "DSL", "symbolic", "jit", "devito"]
-requires-python = ">=3.10,<3.13"
+requires-python = ">=3.10,<3.14"
 authors = [
   { name = "Imperial College London", email = "g.gorman@imperial.ac.uk" },
   { name = "Fabio Luporini",  email = "fabio@devitocodes.com" },

tests/test_mpi.py

@@ -2014,6 +2014,24 @@ def test_merge_smart_if_within_conditional(self, mode):
         for n in FindNodes(Conditional).visit(op1):
             assert len(FindNodes(HaloUpdateCall).visit(n)) == 0
 
+    @pytest.mark.parallel(mode=2)
+    def test_allreduce_time(self, mode):
+        space_order = 8
+        nx, ny = 11, 11
+
+        grid = Grid(shape=(nx, ny))
+        tt = grid.time_dim
+        nt = 10
+
+        ux = TimeFunction(name="ux", grid=grid, time_order=1, space_order=space_order)
+        g = TimeFunction(name="g", grid=grid, dimensions=(tt, ), shape=(nt,))
+
+        op = Operator([Eq(ux.forward, ux + tt), Inc(g, ux)], name="Op")
+        assert_structure(op, ['t,x,y', 't'], 'txy')
+
+        op.apply(time_m=0, time_M=nt-1)
+        assert np.isclose(np.max(g.data), 4356.0)
+
 
 class TestOperatorAdvanced: