Merge pull request #2564 from devitocodes/real-imag

dsl: Add Real, Imag, and Conj operators

Author: mloubout

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

@@ -30,8 +30,8 @@ jobs:
 
       matrix:
         name: [
-           pytest-ubuntu-py311-gcc11-noomp,
-           pytest-ubuntu-py312-gcc12-omp,
+           pytest-ubuntu-py311-gcc11-cxxnoomp,
+           pytest-ubuntu-py312-gcc12-cxxomp,
            pytest-ubuntu-py39-gcc14-omp,
            pytest-ubuntu-py310-gcc10-noomp,
            pytest-ubuntu-py312-gcc13-omp,
@@ -42,18 +42,18 @@ jobs:
         ]
         set: [base, adjoint]
         include:
-        - name: pytest-ubuntu-py311-gcc11-noomp
+        - name: pytest-ubuntu-py311-gcc11-cxxnoomp
           python-version: '3.11'
           os: ubuntu-22.04
           arch: "gcc-11"
-          language: "C"
+          language: "CXX"
           sympy: "1.11"
 
-        - name: pytest-ubuntu-py312-gcc12-omp
+        - name: pytest-ubuntu-py312-gcc12-cxxomp
           python-version: '3.12'
           os: ubuntu-24.04
           arch: "gcc-12"
-          language: "openmp"
+          language: "CXXopenmp"
           sympy: "1.13"
 
         - name: pytest-ubuntu-py39-gcc14-omp

devito/core/__init__.py

@@ -26,64 +26,82 @@
     DeviceNoopOmpOperator, DeviceNoopAccOperator,
     DeviceAdvOmpOperator, DeviceAdvAccOperator,
     DeviceFsgOmpOperator, DeviceFsgAccOperator,
-    DeviceCustomOmpOperator, DeviceCustomAccOperator
+    DeviceCustomOmpOperator, DeviceCustomAccOperator,
+    DeviceCustomCXXOmpOperator, DeviceNoopCXXOmpOperator,
+    DeviceAdvCXXOmpOperator, DeviceFsgCXXOmpOperator
 )
 from devito.operator.registry import operator_registry
 
 # Register CPU Operators
 operator_registry.add(Cpu64CustomOperator, Cpu64, 'custom', 'C')
 operator_registry.add(Cpu64CustomOperator, Cpu64, 'custom', 'openmp')
+operator_registry.add(Cpu64CustomOperator, Cpu64, 'custom', 'Copenmp')
 operator_registry.add(Cpu64CustomCXXOperator, Cpu64, 'custom', 'CXX')
 operator_registry.add(Cpu64CustomCXXOperator, Cpu64, 'custom', 'CXXopenmp')
 
 operator_registry.add(Cpu64NoopCOperator, Cpu64, 'noop', 'C')
 operator_registry.add(Cpu64NoopOmpOperator, Cpu64, 'noop', 'openmp')
+operator_registry.add(Cpu64NoopOmpOperator, Cpu64, 'noop', 'Copenmp')
 operator_registry.add(Cpu64CXXNoopCOperator, Cpu64, 'noop', 'CXX')
 operator_registry.add(Cpu64CXXNoopOmpOperator, Cpu64, 'noop', 'CXXopenmp')
 
 operator_registry.add(Cpu64AdvCOperator, Cpu64, 'advanced', 'C')
 operator_registry.add(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'openmp')
+operator_registry.add(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'Copenmp')
 operator_registry.add(Cpu64AdvCXXOperator, Cpu64, 'advanced', 'CXX')
 operator_registry.add(Cpu64AdvCXXOmpOperator, Cpu64, 'advanced', 'CXXopenmp')
 
 operator_registry.add(Cpu64FsgCOperator, Cpu64, 'advanced-fsg', 'C')
 operator_registry.add(Cpu64FsgOmpOperator, Cpu64, 'advanced-fsg', 'openmp')
+operator_registry.add(Cpu64FsgOmpOperator, Cpu64, 'advanced-fsg', 'Copenmp')
 operator_registry.add(Cpu64FsgCXXOperator, Cpu64, 'advanced-fsg', 'CXX')
 operator_registry.add(Cpu64FsgCXXOmpOperator, Cpu64, 'advanced-fsg', 'CXXopenmp')
 
 operator_registry.add(Intel64AdvCOperator, Intel64, 'advanced', 'C')
 operator_registry.add(Intel64AdvOmpOperator, Intel64, 'advanced', 'openmp')
+operator_registry.add(Intel64AdvOmpOperator, Intel64, 'advanced', 'Copenmp')
 operator_registry.add(Intel64CXXAdvCOperator, Intel64, 'advanced', 'CXX')
 operator_registry.add(Intel64AdvCXXOmpOperator, Intel64, 'advanced', 'CXXopenmp')
 
 operator_registry.add(Intel64FsgCOperator, Intel64, 'advanced-fsg', 'C')
 operator_registry.add(Intel64FsgOmpOperator, Intel64, 'advanced-fsg', 'openmp')
+operator_registry.add(Intel64FsgOmpOperator, Intel64, 'advanced-fsg', 'Copenmp')
 operator_registry.add(Intel64FsgCXXOperator, Intel64, 'advanced-fsg', 'CXX')
 operator_registry.add(Intel64FsgCXXOmpOperator, Intel64, 'advanced-fsg', 'CXXopenmp')
 
 operator_registry.add(ArmAdvCOperator, Arm, 'advanced', 'C')
 operator_registry.add(ArmAdvOmpOperator, Arm, 'advanced', 'openmp')
+operator_registry.add(ArmAdvOmpOperator, Arm, 'advanced', 'Copenmp')
 operator_registry.add(ArmAdvCXXOperator, Arm, 'advanced', 'CXX')
 operator_registry.add(ArmAdvCXXOmpOperator, Arm, 'advanced', 'CXXopenmp')
 
 operator_registry.add(PowerAdvCOperator, Power, 'advanced', 'C')
 operator_registry.add(PowerAdvOmpOperator, Power, 'advanced', 'openmp')
+operator_registry.add(PowerAdvOmpOperator, Power, 'advanced', 'Copenmp')
 operator_registry.add(PowerCXXAdvCOperator, Power, 'advanced', 'CXX')
 operator_registry.add(PowerAdvCXXOmpOperator, Power, 'advanced', 'CXXopenmp')
 
 # Register Device Operators
 operator_registry.add(DeviceCustomOmpOperator, Device, 'custom', 'C')
 operator_registry.add(DeviceCustomOmpOperator, Device, 'custom', 'openmp')
+operator_registry.add(DeviceCustomCXXOmpOperator, Device, 'custom', 'CXX')
+operator_registry.add(DeviceCustomCXXOmpOperator, Device, 'custom', 'CXXopenmp')
 operator_registry.add(DeviceCustomAccOperator, Device, 'custom', 'openacc')
 
 operator_registry.add(DeviceNoopOmpOperator, Device, 'noop', 'C')
 operator_registry.add(DeviceNoopOmpOperator, Device, 'noop', 'openmp')
+operator_registry.add(DeviceNoopCXXOmpOperator, Device, 'noop', 'CXX')
+operator_registry.add(DeviceNoopCXXOmpOperator, Device, 'noop', 'CXXopenmp')
 operator_registry.add(DeviceNoopAccOperator, Device, 'noop', 'openacc')
 
 operator_registry.add(DeviceAdvOmpOperator, Device, 'advanced', 'C')
 operator_registry.add(DeviceAdvOmpOperator, Device, 'advanced', 'openmp')
+operator_registry.add(DeviceAdvCXXOmpOperator, Device, 'advanced', 'CXX')
+operator_registry.add(DeviceAdvCXXOmpOperator, Device, 'advanced', 'CXXopenmp')
 operator_registry.add(DeviceAdvAccOperator, Device, 'advanced', 'openacc')
 
 operator_registry.add(DeviceFsgOmpOperator, Device, 'advanced-fsg', 'C')
 operator_registry.add(DeviceFsgOmpOperator, Device, 'advanced-fsg', 'openmp')
+operator_registry.add(DeviceFsgCXXOmpOperator, Device, 'advanced-fsg', 'CXX')
+operator_registry.add(DeviceFsgCXXOmpOperator, Device, 'advanced-fsg', 'CXXopenmp')
 operator_registry.add(DeviceFsgAccOperator, Device, 'advanced-fsg', 'openacc')

devito/core/cpu.py

@@ -322,7 +322,7 @@ def _make_iet_passes_mapper(cls, **kwargs):
 
 class Cpu64CustomCXXOperator(Cpu64CustomOperator):
 
-    _Target = CXXTarget
+    _Target = CXXOmpTarget
     LINEARIZE = True
 
 # Language level

devito/core/gpu.py

@@ -10,15 +10,17 @@
 from devito.passes.clusters import (Lift, tasking, memcpy_prefetch, blocking,
                                     buffering, cire, cse, factorize, fission, fuse,
                                     optimize_pows)
-from devito.passes.iet import (DeviceOmpTarget, DeviceAccTarget, mpiize,
-                               hoist_prodders, linearize, pthreadify,
+from devito.passes.iet import (DeviceOmpTarget, DeviceAccTarget, DeviceCXXOmpTarget,
+                               mpiize, hoist_prodders, linearize, pthreadify,
                                relax_incr_dimensions, check_stability)
 from devito.tools import as_tuple, timed_pass
 
 __all__ = ['DeviceNoopOperator', 'DeviceAdvOperator', 'DeviceCustomOperator',
            'DeviceNoopOmpOperator', 'DeviceAdvOmpOperator', 'DeviceFsgOmpOperator',
            'DeviceCustomOmpOperator', 'DeviceNoopAccOperator', 'DeviceAdvAccOperator',
-           'DeviceFsgAccOperator', 'DeviceCustomAccOperator']
+           'DeviceFsgAccOperator', 'DeviceCustomAccOperator', 'DeviceNoopCXXOmpOperator',
+           'DeviceAdvCXXOmpOperator', 'DeviceFsgCXXOmpOperator',
+           'DeviceCustomCXXOmpOperator']
 
 
 class DeviceOperatorMixin:
@@ -364,14 +366,29 @@ class DeviceNoopOmpOperator(DeviceOmpOperatorMixin, DeviceNoopOperator):
     pass
 
 
+class DeviceNoopCXXOmpOperator(DeviceNoopOmpOperator):
+    _Target = DeviceCXXOmpTarget
+    LINEARIZE = True
+
+
 class DeviceAdvOmpOperator(DeviceOmpOperatorMixin, DeviceAdvOperator):
     pass
 
 
+class DeviceAdvCXXOmpOperator(DeviceAdvOmpOperator):
+    _Target = DeviceCXXOmpTarget
+    LINEARIZE = True
+
+
 class DeviceFsgOmpOperator(DeviceOmpOperatorMixin, DeviceFsgOperator):
     pass
 
 
+class DeviceFsgCXXOmpOperator(DeviceFsgOmpOperator):
+    _Target = DeviceCXXOmpTarget
+    LINEARIZE = True
+
+
 class DeviceCustomOmpOperator(DeviceOmpOperatorMixin, DeviceCustomOperator):
 
     _known_passes = DeviceCustomOperator._known_passes + ('openmp',)
@@ -384,6 +401,11 @@ def _make_iet_passes_mapper(cls, **kwargs):
         return mapper
 
 
+class DeviceCustomCXXOmpOperator(DeviceCustomOmpOperator):
+    _Target = DeviceCXXOmpTarget
+    LINEARIZE = True
+
+
 # OpenACC
 
 class DeviceAccOperatorMixin:

devito/finite_differences/differentiable.py

@@ -17,12 +17,12 @@
 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, is_number)
+                          infer_dtype, extract_dtype, is_integer, split, is_number)
 from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 from devito.types.basic import AbstractFunction
 
 __all__ = ['Differentiable', 'DiffDerivative', 'IndexDerivative', 'EvalDerivative',
-           'Weights']
+           'Weights', 'Real', 'Imag', 'Conj']
 
 
 class Differentiable(sympy.Expr, Evaluable):
@@ -644,6 +644,46 @@ def __str__(self):
     __repr__ = __str__
 
 
+class ComplexPart(Differentiable, sympy.core.function.Application):
+    """Abstract class for `Real`, `Imag`, or `Conj` of an expression"""
+    _name = None
+
+    def __new__(cls, *args, **kwargs):
+        if len(args) != 1:
+            raise ValueError(f"{cls.__name__} expects exactly one arg;"
+                             f" {len(args)} were supplied instead.")
+
+        return super().__new__(cls, *args, **kwargs)
+
+    def __str__(self):
+        return f"{self.__class__.__name__}({self.args[0]})"
+
+    __repr__ = __str__
+
+
+class RealComplexPart(ComplexPart):
+
+    @cached_property
+    def dtype(self):
+        dtype = extract_dtype(self)
+        return dtype(0).real.__class__
+
+
+class Real(RealComplexPart):
+    """Get the real part of an expression"""
+    _name = 'real'
+
+
+class Imag(RealComplexPart):
+    """Get the imaginary part of an expression"""
+    _name = 'imag'
+
+
+class Conj(ComplexPart):
+    """Get the complex conjugate of an expression"""
+    _name = 'conj'
+
+
 class IndexSum(sympy.Expr, Evaluable):
 
     """

devito/finite_differences/tools.py

@@ -268,11 +268,10 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None, nweights
             raise ValueError(f"More weights ({nweights}) provided than the maximum "
                              f"stencil size ({order + 1}) for order {order} scheme")
         elif do > dw:
+            order = nweights - nweights % 2
             warning(f"Less weights ({nweights}) provided than the stencil size"
                     f"({order + 1}) for order {order} scheme."
-                    " Reducing order to {nweights//2}")
-            order = nweights - nweights % 2
-
+                    f" Reducing order to {order}")
     # Evaluation point
     x0 = sympify(((x0 or {}).get(dim) or expr.indices_ref[dim]))
 

devito/ir/iet/visitors.py

@@ -423,7 +423,10 @@ def visit_PointerCast(self, o):
             if o.flat is None:
                 shape = ''.join(f"[{self.ccode(i)}]" for i in o.castshape)
                 rshape = f'(*){shape}'
-                lvalue = c.Value(cstr, f'(*{self._restrict_keyword} {v}){shape}')
+                if shape:
+                    lvalue = c.Value(cstr, f'(*{self._restrict_keyword} {v}){shape}')
+                else:
+                    lvalue = c.Value(cstr, f'*{self._restrict_keyword} {v}')
             else:
                 rshape = '*'
                 lvalue = c.Value(cstr, f'*{v}')

devito/operator/operator.py

@@ -1337,6 +1337,10 @@ def parse_kwargs(**kwargs):
 
     if not opt or isinstance(opt, str):
         mode, options = opt, {}
+        # Legacy Operator(..., opt='openmp', ...) support
+        if mode == 'openmp':
+            mode = 'noop'
+            options = {'openmp': True}
     elif isinstance(opt, tuple):
         if len(opt) == 0:
             mode, options = 'noop', {}
@@ -1353,7 +1357,7 @@ def parse_kwargs(**kwargs):
     # `opt`, deprecated kwargs
     kwopenmp = kwargs.get('openmp', options.get('openmp'))
     if kwopenmp is None:
-        openmp = kwargs.get('language', configuration['language']) == 'openmp'
+        openmp = 'openmp' in kwargs.get('language', configuration['language'])
     else:
         openmp = kwopenmp
 
@@ -1399,7 +1403,9 @@ def parse_kwargs(**kwargs):
         kwargs['language'] = language
     elif kwopenmp is not None:
         # Handle deprecated `openmp` kwarg for backward compatibility
-        kwargs['language'] = 'openmp' if openmp else 'C'
+        omp = {'C': 'openmp', 'CXX': 'CXXopenmp'}.get(configuration['language'],
+                                                      'openmp')
+        kwargs['language'] = omp if openmp else 'C'
     else:
         kwargs['language'] = configuration['language']
 

devito/passes/clusters/aliases.py

@@ -1354,7 +1354,7 @@ def cost(self):
         # Not just the sum for the individual items' cost! There might be
         # redundancies, which we factor out here...
         counter = generator()
-        make = lambda: Symbol(name='dummy%d' % counter(), dtype=np.float32)
+        make = lambda _: Symbol(name='dummy%d' % counter(), dtype=np.float32)
 
         tot = 0
         for v in as_mapper(self, lambda i: i.ispace).values():

devito/passes/clusters/cse.py

@@ -14,7 +14,7 @@
 from devito.ir import Cluster, Scope, cluster_pass
 from devito.symbolics import estimate_cost, q_leaf, q_terminal
 from devito.symbolics.manipulation import _uxreplace
-from devito.tools import DAG, as_list, as_tuple, frozendict
+from devito.tools import DAG, as_list, as_tuple, frozendict, extract_dtype
 from devito.types import Eq, Symbol, Temp
 
 __all__ = ['cse']
@@ -78,7 +78,8 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     if cluster.is_fence:
         return cluster
 
-    make = lambda: CTemp(name=sregistry.make_name(), dtype=dtype)
+    make_dtype = lambda e: np.promote_types(e.dtype, dtype).type
+    make = lambda e: CTemp(name=sregistry.make_name(), dtype=make_dtype(e))
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)
 
@@ -118,7 +119,7 @@ def _cse(maybe_exprs, make, min_cost=1, mode='basic'):
             exprs = maybe_exprs
             scope = Scope(maybe_exprs)
         else:
-            exprs = [Eq(make(), e) for e in maybe_exprs]
+            exprs = [Eq(make(e), e) for e in maybe_exprs]
             scope = Scope([])
 
     # Some sub-expressions aren't really "common" -- that's the case of Dimension-
@@ -155,7 +156,7 @@ def _cse(maybe_exprs, make, min_cost=1, mode='basic'):
         candidates = [c for c in candidates if c.cost == cost]
 
         # Apply replacements
-        chosen = [(c, scheduled.get(c.key) or make()) for c in candidates]
+        chosen = [(c, scheduled.get(c.key) or make(c)) for c in candidates]
         exprs = _inject(exprs, chosen, scheduled)
 
     # Drop useless temporaries (e.g., r0=r1)
@@ -275,6 +276,10 @@ def __new__(cls, expr, conditionals=None, sources=()):
     def expr(self):
         return self[0]
 
+    @property
+    def dtype(self):
+        return extract_dtype(self.expr)
+
     @property
     def conditionals(self):
         return self[1]