compiler: fix complex reductions for gnu

Author: mloubout

devito/mpi/distributed.py

@@ -40,7 +40,7 @@
     def cleanup():
         devito_mpi_finalize()
     atexit.register(cleanup)
-except ImportError as e:
+except (RuntimeError, ImportError) as e:
     # Dummy fallback in case mpi4py/MPI aren't available
     class NoneMetaclass(type):
         def __getattr__(self, name):

devito/passes/iet/languages/openmp.py

@@ -2,10 +2,12 @@
 from packaging.version import Version
 
 import cgen as c
+import numpy as np
 from sympy import And, Ne, Not
 
 from devito.arch import AMDGPUX, NVIDIAX, INTELGPUX, PVC
 from devito.arch.compiler import GNUCompiler, NvidiaCompiler
+from devito.exceptions import InvalidOperator
 from devito.ir import (Call, Conditional, DeviceCall, List, Pragma, Prodder,
                        ParallelBlock, PointerCast, While, FindSymbols)
 from devito.passes.iet.definitions import DataManager, DeviceAwareDataManager
@@ -18,6 +20,7 @@
 from devito.passes.iet.languages.C import CBB
 from devito.passes.iet.languages.CXX import CXXBB
 from devito.symbolics import CondEq, DefFunction
+from devito.symbolics.extended_sympy import UnaryOp
 from devito.tools import filter_ordered
 
 __all__ = ['SimdOmpizer', 'Ompizer', 'OmpIteration', 'OmpRegion',
@@ -113,6 +116,44 @@ def _generate(self):
         return self.pragma % (joins(*items), n)
 
 
+class RealExt(UnaryOp):
+
+    _op = '__real__ '
+
+
+class ImagExt(UnaryOp):
+
+    _op = '__imag__ '
+
+
+def atomic_add(i, pragmas):
+    lhs, rhs = i.expr.lhs, i.expr.rhs
+    if (np.issubdtype(lhs.dtype, np.complexfloating)
+       and np.issubdtype(rhs.dtype, np.complexfloating)):
+        # Complex i, complex j
+        # Atomic add real and imaginary parts separately
+        lhsr, rhsr = RealExt(lhs), RealExt(rhs)
+        lhsi, rhsi = ImagExt(lhs), ImagExt(rhs)
+        real = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
+                          pragmas=pragmas)
+        imag = i._rebuild(expr=i.expr._rebuild(lhs=lhsi, rhs=rhsi),
+                          pragmas=pragmas)
+        return List(body=[real, imag])
+
+    elif (np.issubdtype(lhs.dtype, np.complexfloating)
+          and not np.issubdtype(rhs.dtype, np.complexfloating)):
+        # Complex i, real j
+        # Atomic add j to real part of i
+        lhsr, rhsr = RealExt(lhs), rhs
+        real = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
+                          pragmas=pragmas)
+        return real
+    else:
+        # Real i, complex j
+        raise InvalidOperator("Atomic add not implemented for real "
+                              "Functions with complex increments")
+
+
 class AbstractOmpBB(LangBB):
 
     mapper = {
@@ -134,7 +175,8 @@ class AbstractOmpBB(LangBB):
         'simd-for-aligned': lambda n, *a:
             SimdForAligned('omp simd aligned(%s:%d)', arguments=(n, *a)),
         'atomic':
-            Pragma('omp atomic update')
+            Pragma('omp atomic update'),
+        'split-atomic': lambda i: atomic_add(i, Pragma('omp atomic update'))
     }
 
     Region = OmpRegion
@@ -241,6 +283,13 @@ def _support_array_reduction(cls, compiler):
         else:
             return True
 
+    @classmethod
+    def _support_complex_reduction(cls, compiler):
+        if isinstance(compiler, GNUCompiler):
+            # Gcc doesn't supports complex reduction
+            return False
+        return True
+
 
 class Ompizer(AbstractOmpizer):
     langbb = OmpBB

devito/passes/iet/parpragma.py

@@ -42,6 +42,10 @@ class PragmaSimdTransformer(PragmaTransformer):
     def _support_array_reduction(cls, compiler):
         return True
 
+    @classmethod
+    def _support_complex_reduction(cls, compiler):
+        return False
+
     @property
     def simd_reg_nbytes(self):
         return self.platform.simd_reg_nbytes
@@ -238,8 +242,13 @@ def _make_reductions(self, partree):
             # Implement reduction
             mapper = {partree.root: partree.root._rebuild(reduction=reductions)}
         elif all(i is OpInc for _, _, i in reductions):
-            # Use atomic increments
-            mapper = {i: i._rebuild(pragmas=self.langbb['atomic']) for i in exprs}
+            test2 = not self._support_complex_reduction(self.compiler) and \
+                any(np.iscomplexobj(i.dtype(0)) for i, _, _ in reductions)
+            if test2:
+                mapper = {i: self.langbb['split-atomic'](i) for i in exprs}
+            else:
+                # Use atomic increments
+                mapper = {i: i._rebuild(pragmas=self.langbb['atomic']) for i in exprs}
         else:
             raise NotImplementedError
 

tests/test_dtypes.py

@@ -2,9 +2,13 @@
 import pytest
 import sympy
 
+from conftest import skipif
+
 from devito import (
     Constant, Eq, Function, Grid, Operator, exp, log, sin, configuration
 )
+from devito.arch.compiler import GNUCompiler
+from devito.exceptions import InvalidOperator
 from devito.ir.cgen.printer import BasePrinter
 from devito.passes.iet.langbase import LangBB
 from devito.passes.iet.languages.C import CBB, CPrinter
@@ -13,6 +17,7 @@
 from devito.symbolics.extended_dtypes import ctypes_vector_mapper
 from devito.types.basic import Basic, Scalar, Symbol
 from devito.types.dense import TimeFunction
+from devito.types.sparse import SparseTimeFunction
 
 # Mappers for language-specific types and headers
 _languages: dict[str, type[LangBB]] = {
@@ -274,3 +279,45 @@ def test_complex_space_deriv(dtype: np.dtype[np.complexfloating]) -> None:
     dfdy = h.data.T[1:-1, 1:-1]
     assert np.allclose(dfdx, np.ones((5, 5), dtype=dtype))
     assert np.allclose(dfdy, np.ones((5, 5), dtype=dtype))
+
+
+@skipif('device')
+@pytest.mark.parametrize('dtypeu', [np.float32, np.complex64, np.complex128])
+def test_complex_reduction(dtypeu: np.dtype[np.complexfloating]) -> None:
+    """
+    Tests reductions over complex-valued functions.
+    """
+    grid = Grid((11, 11))
+
+    u = TimeFunction(name="u", grid=grid, space_order=2, time_order=1, dtype=dtypeu)
+    for dtypes in [dtypeu, dtypeu(0).real.__class__]:
+        u.data.fill(0)
+        s = SparseTimeFunction(name="s", grid=grid, npoint=1, nt=10, dtype=dtypes)
+        if np.issubdtype(dtypes, np.complexfloating):
+            s.data[:] = 1 + 2j
+            expected = 8. + 16.j
+        else:
+            s.data[:] = 1
+            expected = 8.
+        s.coordinates.data[:] = [.5, .5]
+
+        # s complex and u real should error
+        if np.issubdtype(dtypeu, np.floating) and \
+           np.issubdtype(dtypes, np.complexfloating):
+            with pytest.raises(InvalidOperator):
+                op = Operator([Eq(u.forward, u)] + s.inject(u.forward, expr=s))
+            continue
+        else:
+            op = Operator([Eq(u.forward, u)] + s.inject(u.forward, expr=s))
+        op()
+
+        if isinstance(configuration['compiler'], GNUCompiler) and \
+           np.issubdtype(dtypeu, np.complexfloating):
+            ru = '__real__ u[t1][rsx + posx + 2][rsy + posy + 2]'
+            iu = '__imag__ u[t1][rsx + posx + 2][rsy + posy + 2]'
+            assert f'{ru} += __real__ r0' in str(op)
+            assert f'{iu} += __imag__ r0' in str(op)
+        else:
+            assert 'u[t1][rsx + posx + 2][rsy + posy + 2] += r0' in str(op)
+
+        assert np.isclose(u.data[0, 5, 5], expected)