Merge pull request #2662 from devitocodes/red-space

mpi: fix allreduce iteration space

Author: mloubout

.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

.github/workflows/pytest-gpu.yml

@@ -74,7 +74,7 @@ jobs:
 
     - name: Build docker image
       run: |
-          docker build . --rm --pull --file docker/Dockerfile.devito --tag ${{ matrix.name }} --build-arg base=${{ matrix.base }}
+          docker build . --rm --pull --file docker/Dockerfile.devito --tag ${{ matrix.name }}-${{ github.ref }} --build-arg base=${{ matrix.base }}
 
     - name: Remove dangling layers
       run: |
@@ -87,13 +87,13 @@ jobs:
         docker ps -aq | xargs -r docker rm -f
         # Run tests
         ci_env=`bash <(curl -s https://codecov.io/env)`
-        docker run ${{ matrix.flags }} ${{ matrix.name }} ${{ matrix.test_drive_cmd }}
-        docker run ${{ matrix.flags }} $ci_env -e CI=true ${{ matrix.name }} pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ matrix.test_files }}
+        docker run ${{ matrix.flags }} ${{ matrix.name }}-${{ github.ref }} ${{ matrix.test_drive_cmd }}
+        docker run ${{ matrix.flags }} $ci_env -e CI=true ${{ matrix.name }}-${{ github.ref }} pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ matrix.test_files }}
 
     - name: Test examples
       run: |
-        docker run ${{ matrix.flags }} ${{ matrix.name }} pytest ${{ matrix.test_examples }}
+        docker run ${{ matrix.flags }} ${{ matrix.name }}-${{ github.ref }} pytest ${{ matrix.test_examples }}
 
     - name: Test examples with MPI
       run: |
-        docker run ${{ matrix.flags }} --env DEVITO_MPI=1 ${{ matrix.name }} mpiexec -n 2 pytest ${{ matrix.test_examples }}
+        docker run ${{ matrix.flags }} --env DEVITO_MPI=1 ${{ matrix.name }}-${{ github.ref }} mpiexec -n 2 pytest ${{ matrix.test_examples }}

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

@@ -30,7 +30,7 @@
            'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration',
            'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma',
            'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace',
-           'Using', 'CallableBody', 'Transfer']
+           'Using', 'CallableBody', 'Transfer', 'EmptyList']
 
 # First-class IET nodes
 
@@ -218,6 +218,11 @@ def __repr__(self):
                                       len(self.body), len(self.footer))
 
 
+class EmptyList(List):
+    """A plain List node without a body, such as a header/footer only"""
+    _traversable = []
+
+
 class Block(List):
 
     """A sequence of Nodes, wrapped in a block {...}."""

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) and any(o.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")

devito/passes/iet/misc.py

@@ -7,7 +7,7 @@
 from devito.finite_differences import Max, Min
 from devito.finite_differences.differentiable import SafeInv
 from devito.logger import warning
-from devito.ir import (Any, Forward, DummyExpr, Iteration, List, Prodder,
+from devito.ir import (Any, Forward, DummyExpr, Iteration, EmptyList, Prodder,
                        FindApplications, FindNodes, FindSymbols, Transformer,
                        Uxreplace, filter_iterations, retrieve_iteration_tree,
                        pull_dims)
@@ -47,7 +47,7 @@ def avoid_denormals(iet, platform=None, **kwargs):
               cgen.Statement('_MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON)'),
               cgen.Line())
 
-    body = iet.body._rebuild(body=(List(header=header),) + iet.body.body)
+    body = iet.body._rebuild(body=(EmptyList(header=header),) + iet.body.body)
     iet = iet._rebuild(body=body)
 
     return iet, {'includes': ('xmmintrin.h', 'pmmintrin.h')}

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: