misc: Fix conflict with main

Author: ZoeLeibowitz

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

@@ -19,7 +19,7 @@ on:
 jobs:
   test-mpi-basic:
     name: pytest-mpi
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         python-version: ['3.9', '3.11']
@@ -30,6 +30,7 @@ jobs:
       OMP_NUM_THREADS: "1"
       CC: "gcc-9"
       CXX: "g++-9"
+      RDMAV_FORK_SAFE: 1
 
     steps:
     - uses: actions/setup-python@v5
@@ -78,20 +79,20 @@ jobs:
               os: ubuntu-latest
               # Need safe math for icx due to inaccuracy with mpi+sinc interpolation
               mpiflag: "-e DEVITO_SAFE_MATH=1"
-      
+
       steps:
       - name: Checkout devito
         uses: actions/checkout@v4
-      
+
       - name: Build docker image
         run: |
           docker build . --file docker/Dockerfile.devito --tag devito_img --build-arg base=devitocodes/bases:cpu-${{ matrix.arch }}
-      
+
       - name: Test with pytest
         run: |
           docker run --init -t --rm -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} -e OMP_NUM_THREADS=1 --name testrun devito_img pytest tests/test_mpi.py
 
       - name: Test examples with MPI
         run: |
           docker run --init -t --rm ${{ matrix.mpiflag }} -e DEVITO_MPI=1 -e OMP_NUM_THREADS=1 --name examplerun devito_img mpiexec -n 2 pytest examples/seismic/acoustic
-          docker run --init -t --rm -e DEVITO_MPI=1 -e OMP_NUM_THREADS=1 --name examplerun devito_img mpiexec -n 2 pytest examples/seismic/tti
+          docker run --init -t --rm -e DEVITO_MPI=1 -e OMP_NUM_THREADS=1 --name examplerun devito_img mpiexec -n 2 pytest examples/seismic/tti

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

@@ -34,7 +34,7 @@ jobs:
         name: [
            pytest-ubuntu-py311-gcc11-noomp,
            pytest-ubuntu-py312-gcc12-omp,
-           pytest-ubuntu-py39-gcc7-omp,
+           pytest-ubuntu-py39-gcc14-omp,
            pytest-ubuntu-py310-gcc10-noomp,
            pytest-ubuntu-py312-gcc13-omp,
            pytest-ubuntu-py39-gcc9-omp,
@@ -58,16 +58,16 @@ jobs:
           language: "openmp"
           sympy: "1.13"
 
-        - name: pytest-ubuntu-py39-gcc7-omp
+        - name: pytest-ubuntu-py39-gcc14-omp
           python-version: '3.9'
-          os: ubuntu-20.04
-          arch: "gcc-7"
+          os: ubuntu-24.04
+          arch: "gcc-14"
           language: "openmp"
           sympy: "1.9"
 
         - name: pytest-ubuntu-py310-gcc10-noomp
           python-version: '3.10'
-          os: ubuntu-20.04
+          os: ubuntu-22.04
           arch: "gcc-10"
           language: "C"
           sympy: "1.11"
@@ -81,7 +81,7 @@ jobs:
 
         - name: pytest-ubuntu-py39-gcc9-omp
           python-version: '3.9'
-          os: ubuntu-20.04
+          os: ubuntu-22.04
           arch: "custom"
           language: "openmp"
           sympy: "1.10"
@@ -154,7 +154,7 @@ jobs:
             echo "/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
           fi
       id: set-tests
-    
+
     - name: Set pip flags for latest python (3.12)
       run: |
           if [ "${{ matrix.python-version }}" == '3.12' ]; then
@@ -174,7 +174,7 @@ jobs:
 
     - name: Test with pytest
       run: |
-        ${{ env.RUN_CMD }} pytest -k "${{ matrix.test-set }}" -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+        ${{ env.RUN_CMD }} pytest -k "${{ matrix.test-set }}" -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml tests/
 
     - name: Upload coverage to Codecov
       if: "!contains(matrix.name, 'docker')"

devito/finite_differences/derivative.py

@@ -405,7 +405,7 @@ def _eval_at(self, func):
             return self
         # For basic equation of the form f = Derivative(g, ...) we can just
         # compare staggering
-        if self.expr.staggered == func.staggered:
+        if self.expr.staggered == func.staggered and self.expr.is_Function:
             return self
 
         x0 = func.indices_ref.getters

devito/finite_differences/differentiable.py

@@ -17,7 +17,7 @@
 from devito.finite_differences.tools import make_shift_x0, coeff_priority
 from devito.logger import warning
 from devito.tools import (as_tuple, filter_ordered, flatten, frozendict,
-                          infer_dtype, is_integer, split)
+                          infer_dtype, is_integer, split, is_number)
 from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 from devito.types.basic import AbstractFunction
 
@@ -548,20 +548,19 @@ def __new__(cls, *args, **kwargs):
         nested, others = split(args, lambda e: isinstance(e, Mul))
         args = flatten(e.args for e in nested) + list(others)
 
+        # Gather all numbers and simplify
+        nums, others = split(args, lambda e: is_number(e))
+        scalar = sympy.Mul(*nums)
+
         # a*0 -> 0
-        if any(i == 0 for i in args):
+        if scalar == 0:
             return sympy.S.Zero
 
         # a*1 -> a
-        args = [i for i in args if i != 1]
-
-        # a*-1 -> a*-1
-        # a*-1*-1 -> a
-        # a*-1*-1*-1 -> a*-1
-        nminus = len([i for i in args if i == sympy.S.NegativeOne])
-        args = [i for i in args if i != sympy.S.NegativeOne]
-        if nminus % 2 == 1:
-            args.append(sympy.S.NegativeOne)
+        if scalar - 1 == 0:
+            args = others
+        else:
+            args = [scalar] + others
 
         # Reorder for homogeneity with pure SymPy types
         _mulsort(args)

devito/finite_differences/operators.py

@@ -165,13 +165,19 @@ def diag(func, size=None):
         size of the diagonal matrix (size x size).
         Defaults to the number of spatial dimensions when unspecified
     """
+    from devito.types.tensor import TensorFunction, TensorTimeFunction
+    if isinstance(func, TensorFunction):
+        if func.is_TensorValued:
+            return func._new(*func.shape, lambda i, j: func[i, i] if i == j else 0)
+        else:
+            n = func.shape[0]
+            return func._new(n, n, lambda i, j: func[i] if i == j else 0)
+
     dim = size or len(func.dimensions)
     dim = dim-1 if func.is_TimeDependent else dim
     to = getattr(func, 'time_order', 0)
 
-    from devito.types.tensor import TensorFunction, TensorTimeFunction
     tens_func = TensorTimeFunction if func.is_TimeDependent else TensorFunction
-
     comps = [[func if i == j else 0 for i in range(dim)] for j in range(dim)]
     return tens_func(name='diag', grid=func.grid, space_order=func.space_order,
                      components=comps, time_order=to, diagonal=True)

devito/finite_differences/rsfd.py

@@ -1,6 +1,5 @@
 from functools import wraps
 
-from devito.types import NODE
 from devito.types.dimension import StencilDimension
 from .differentiable import Weights, DiffDerivative
 from .tools import generate_indices, fd_weights_registry
@@ -101,12 +100,7 @@ def check_staggering(func):
     def wrapper(expr, dim, x0=None, expand=True):
         grid = expr.grid
         x0 = {k: v for k, v in x0.items() if k.is_Space}
-        if expr.staggered is NODE or expr.staggered is None:
-            cond = x0 == {} or x0 == all_staggered(grid) or x0 == grid_node(grid)
-        elif expr.staggered == grid.dimensions:
-            cond = x0 == {} or x0 == all_staggered(grid) or x0 == grid_node(grid)
-        else:
-            cond = False
+        cond = x0 == {} or x0 == all_staggered(grid) or x0 == grid_node(grid)
         if cond:
             return func(expr, dim, x0=x0, expand=expand)
         else:
@@ -117,7 +111,8 @@ def wrapper(expr, dim, x0=None, expand=True):
 @check_staggering
 def d45(expr, dim, x0=None, expand=True):
     """
-    RSFD approximation of the derivative of `expr` along `dim` at point `x0`.
+    Rotated staggered grid finite-differences (RSFD) discretization
+    of the derivative of `expr` along `dim` at point `x0`.
 
     Parameters
     ----------
@@ -132,7 +127,8 @@ def d45(expr, dim, x0=None, expand=True):
     """
     # Make sure the grid supports RSFD
     if expr.grid.dim not in [2, 3]:
-        raise ValueError('RSFD only supported in 2D and 3D')
+        raise ValueError('Rotated staggered grid finite-differences (RSFD)'
+                         ' only supported in 2D and 3D')
 
     # Diagonals weights
     w = dir_weights[(dim.name, expr.grid.dim)]

devito/ir/support/syncs.py

@@ -3,6 +3,7 @@
 """
 
 from collections import defaultdict
+from functools import cached_property
 
 from devito.data import FULL
 from devito.tools import Pickable, as_tuple, filter_ordered, frozendict
@@ -163,6 +164,21 @@ def update(self, ops):
             m[d] = set(self.get(d, [])) | set(v)
         return Ops(m)
 
+    def _get_sync(self, cls, dims=None):
+        if dims is None:
+            dims = list(self)
+        for d in dims:
+            for s in self.get(d, []):
+                if isinstance(s, cls):
+                    # NOTE: Remember there can only be one SyncOp of a given
+                    # type per `Ops` object
+                    return s
+        return None
+
+    @cached_property
+    def initarray(self):
+        return self._get_sync(InitArray)
+
 
 def normalize_syncs(*args, strict=True):
     if not args:

devito/mpi/halo_scheme.py

@@ -31,17 +31,32 @@ class HaloSchemeEntry(EnrichedTuple):
 
     __rargs__ = ('loc_indices', 'loc_dirs', 'halos', 'dims')
 
-    def __init__(self, loc_indices, loc_dirs, halos, dims, getters=None):
-        self.loc_indices = frozendict(loc_indices)
-        self.loc_dirs = frozendict(loc_dirs)
-        self.halos = frozenset(halos)
-        self.dims = frozenset(dims)
+    def __new__(cls, loc_indices, loc_dirs, halos, dims, getters=None):
+        items = [frozendict(loc_indices), frozendict(loc_dirs),
+                 frozenset(halos), frozenset(dims)]
+        kwargs = dict(zip(cls.__rargs__, items))
+        return super().__new__(cls, *items, getters=cls.__rargs__, **kwargs)
 
     def __hash__(self):
-        return hash((self.loc_indices,
-                     self.loc_dirs,
-                     self.halos,
-                     self.dims))
+        return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims))
+
+    def union(self, other):
+        """
+        Return a new HaloSchemeEntry that is the union of this and `other`.
+        The `loc_indices` and `loc_dirs` must be the same, otherwise an
+        exception is raised.
+        """
+        if self.loc_indices != other.loc_indices or \
+           self.loc_dirs != other.loc_dirs:
+            raise HaloSchemeException(
+                "Inconsistency found while building a HaloScheme"
+            )
+
+        halos = self.halos | other.halos
+        dims = self.dims | other.dims
+
+        return HaloSchemeEntry(self.loc_indices, self.loc_dirs, halos, dims,
+                               getters=self.getters)
 
 
 Halo = namedtuple('Halo', 'dim side')
@@ -410,9 +425,11 @@ def add(self, f, hse):
         """
         Create a new HaloScheme that contains all entries in `self` plus the one
         passed in input. If `f` already exists in `self`, the old value is
-        overridden.
+        augmented with `hse` (i.e., the halos are unioned).
         """
-        fmapper = dict(self.fmapper.items())
+        fmapper = dict(self.fmapper)
+        if f in fmapper:
+            hse = fmapper[f].union(hse)
         fmapper[f] = hse
         return HaloScheme.build(fmapper, self.honored)
 

devito/mpi/routines.py

@@ -1217,7 +1217,7 @@ def _arg_defaults(self, allocator, alias, args=None):
 
     def _arg_values(self, args=None, **kwargs):
         # Any will do
-        for f in self.target.handles:
+        for f in self.target.components:
             try:
                 alias = kwargs[f.name]
                 break

devito/symbolics/extended_sympy.py

@@ -654,6 +654,10 @@ def __new__(cls, name, arguments=None, template=None, **kwargs):
 
         return obj
 
+    def _eval_is_commutative(self):
+        # DefFunction defaults to commutative
+        return True
+
     @property
     def name(self):
         return self._name