compiler: simplify complex reductions

Author: mloubout

devito/ir/cgen/printer.py

@@ -349,7 +349,8 @@ def _print_ListInitializer(self, expr):
         return f"{{{', '.join(self._print(i) for i in expr.params)}}}"
 
     def _print_IndexedPointer(self, expr):
-        return f"{expr.base}{''.join(f'[{self._print(i)}]' for i in expr.index)}"
+        base = self._print(expr.base)
+        return f"{base}{''.join(f'[{self._print(i)}]' for i in expr.index)}"
 
     def _print_IntDiv(self, expr):
         lhs = self._print(expr.lhs)

devito/passes/iet/languages/C.py

@@ -1,11 +1,11 @@
 import numpy as np
 from sympy.printing.c import C99CodePrinter
 
-from devito.exceptions import InvalidOperator
-from devito.ir import Call, BasePrinter, List
+from devito.ir import Call, BasePrinter
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.orchestration import Orchestrator
 from devito.passes.iet.langbase import LangBB
+from devito.passes.iet.languages.utils import _atomic_add_split
 from devito.symbolics import c_complex, c_double_complex
 from devito.symbolics.extended_sympy import UnaryOp
 from devito.tools import dtype_to_cstr
@@ -28,38 +28,14 @@ def atomic_add(i, pragmas, split=False):
     if not split:
         return i._rebuild(pragmas=pragmas)
     # Complex reduction, split using a temp pointer
-    # Transforns lhs += rhs into
+    # Transforms lhs += rhs into
     # {
     #   pragmas
     #   __real__ lhs += __real__ rhs;
     #   pragmas
     #   __imag__ lhs += __imag__ rhs;
     # }
-    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")
+    return _atomic_add_split(i, pragmas, RealExt, ImagExt)
 
 
 class CBB(LangBB):

devito/passes/iet/languages/CXX.py

@@ -1,17 +1,13 @@
-from ctypes import POINTER
-
 import numpy as np
 from sympy.printing.cxx import CXX11CodePrinter
 
-from devito import Real, Imag
-from devito.exceptions import InvalidOperator
-from devito.ir import Call, UsingNamespace, BasePrinter, DummyExpr, List
+from devito.ir import Call, UsingNamespace, BasePrinter
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.orchestration import Orchestrator
 from devito.passes.iet.langbase import LangBB
+from devito.passes.iet.languages.utils import _atomic_add_split
 from devito.symbolics import c_complex, c_double_complex, IndexedPointer, cast, Byref
-from devito.tools import dtype_to_cstr, dtype_to_ctype
-from devito.types import Pointer
+from devito.tools import dtype_to_cstr
 
 __all__ = ['CXXBB', 'CXXDataManager', 'CXXOrchestrator']
 
@@ -70,49 +66,30 @@ def std_arith(prefix=None):
 """
 
 
+def split_pointer(i, idx):
+    dtype = i.dtype(0).real.__class__
+    ptr = cast(dtype, stars='*')(Byref(i), reinterpret=True)
+    return IndexedPointer(ptr, idx)
+
+
+cxx_imag = lambda i: split_pointer(i, 1)
+cxx_real = lambda i: split_pointer(i, 0)
+
+
 def atomic_add(i, pragmas, split=False):
     # Base case, real reduction
     if not split:
         return i._rebuild(pragmas=pragmas)
     # Complex reduction, split using a temp pointer
     # Transforns lhs += rhs into
     # {
-    #   float * lhs = reinterpret_cast<float*>(&lhs);
     #   pragmas
-    #   lhs[0] += std::real(rhs);
+    #   reinterpret_cast<float*>(&lhs)[0] += std::real(rhs);
     #   pragmas
-    #   lhs[1] += std::imag(rhs);
+    #   reinterpret_cast<float*>(&lhs)[1] += std::imag(rhs);
     # }
     # Make a temp pointer
-    lhs, rhs = i.expr.lhs, i.expr.rhs
-    rdtype = lhs.dtype(0).real.__class__
-    plhs = Pointer(name=f'p{lhs.name}', dtype=POINTER(dtype_to_ctype(rdtype)))
-    peq = DummyExpr(plhs, cast(rdtype, stars='*')(Byref(lhs), reinterpret=True))
-
-    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 = IndexedPointer(plhs, 0), Real(rhs)
-        lhsi, rhsi = IndexedPointer(plhs, 1), Imag(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=[peq, 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 = IndexedPointer(plhs, 0), rhs
-        real = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
-                          pragmas=pragmas)
-        return List(body=[peq, real])
-    else:
-        # Real i, complex j
-        raise InvalidOperator("Atomic add not implemented for real "
-                              "Functions with complex increments")
+    return _atomic_add_split(i, pragmas, cxx_real, cxx_imag)
 
 
 class CXXBB(LangBB):

devito/passes/iet/languages/utils.py

@@ -1,5 +1,42 @@
-__all__ = ['joins']
+import numpy as np
+
+from devito import Real, Imag
+from devito.exceptions import InvalidOperator
+from devito.ir import List
+
+__all__ = ['joins', '_atomic_add_split']
 
 
 def joins(*symbols):
     return ",".join(sorted([i.name for i in symbols]))
+
+
+def _atomic_add_split(i, pragmas, real, imag):
+    # Complex reduction, split between real and imaginary parts.
+    # real is a function i -> real(i)
+    # imag is a function i -> imag(i)
+    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 = real(lhs), Real(rhs)
+        lhsi, rhsi = imag(lhs), Imag(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 = real(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")

devito/passes/iet/linearization.py

@@ -72,7 +72,7 @@ def key1(f, d):
     if f.is_regular:
         # For paddable objects the following holds:
         # `same dim + same halo + same padding_dtype => same (auto-)padding`
-        return (d, f._size_halo[d], f.__padding_dtype__)
+        return (d, f._size_halo[d], f._size_padding[d])
     else:
         return False
 

tests/test_dtypes.py

@@ -326,12 +326,13 @@ def test_complex_reduction(dtypeu: np.dtype[np.complexfloating]) -> None:
         if gnu and np.issubdtype(dtypeu, np.complexfloating):
             if 'CXX' in op._language:
                 rd = dtype_to_cstr(dtypeu(0).real.__class__)
-                assert f'{rd} * p{u.name} = reinterpret_cast<{rd}*>(&uL0' in str(op)
-                assert f'p{u.name}[0] += std::real(r0)' in str(op)
-                assert f'p{u.name}[1] += std::imag(r0)' in str(op)
+                fu = f'reinterpret_cast<{rd}*>(&{ustr})'
+                assert f'{fu}[0] += std::real(r0)' in str(op)
+                assert f'{fu}[1] += std::imag(r0)' in str(op)
             else:
-                assert f'__real__ {ustr} += __real__ r0' in str(op)
-                assert f'__imag__ {ustr} += __imag__ r0' in str(op)
+                ext = '' if dtypeu == np.complex128 else 'f'
+                assert f'__real__ {ustr} += creal{ext}(r0)' in str(op)
+                assert f'__imag__ {ustr} += cimag{ext}(r0)' in str(op)
         else:
             assert f'{ustr} += r0' in str(op)