API: Introducing complex (np.complex64/128) native support

devito/__init__.py

@@ -56,7 +56,8 @@ def reinit_compiler(val):
     """
     Re-initialize the Compiler.
     """
-    configuration['compiler'].__init__(suffix=configuration['compiler'].suffix,
+    configuration['compiler'].__init__(name=configuration['compiler'].name,
+                                       suffix=configuration['compiler'].suffix,
                                        mpi=configuration['mpi'])
     return val
 
@@ -65,7 +66,7 @@ def reinit_compiler(val):
 configuration.add('platform', 'cpu64', list(platform_registry),
                   callback=lambda i: platform_registry[i]())
 configuration.add('compiler', 'custom', list(compiler_registry),
-                  callback=lambda i: compiler_registry[i]())
+                  callback=lambda i: compiler_registry[i](name=i))
 
 # Setup language for shared-memory parallelism
 preprocessor = lambda i: {0: 'C', 1: 'openmp'}.get(i, i)  # Handles DEVITO_OPENMP deprec

devito/arch/compiler.py

@@ -180,6 +180,8 @@ def __init__(self):
     _cpp = False
 
     def __init__(self, **kwargs):
+        self._name = kwargs.pop('name', self.__class__.__name__)
+
         super().__init__(**kwargs)
 
         self.__lookup_cmds__()
@@ -223,13 +225,13 @@ def __new_with__(self, **kwargs):
         Create a new Compiler from an existing one, inherenting from it
         the flags that are not specified via ``kwargs``.
         """
-        return self.__class__(suffix=kwargs.pop('suffix', self.suffix),
+        return self.__class__(name=self.name, suffix=kwargs.pop('suffix', self.suffix),
                               mpi=kwargs.pop('mpi', configuration['mpi']),
                               **kwargs)
 
     @property
     def name(self):
-        return self.__class__.__name__
+        return self._name
 
     @property
     def version(self):
@@ -593,7 +595,7 @@ def __init_finalize__(self, **kwargs):
         self.cflags.remove('-O3')
         self.cflags.remove('-Wall')
 
-        self.cflags.append('-std=c++11')
+        self.cflags.append('-std=c++14')
 
         language = kwargs.pop('language', configuration['language'])
         platform = kwargs.pop('platform', configuration['platform'])
@@ -978,6 +980,7 @@ def __new_with__(self, **kwargs):
     'nvc++': NvidiaCompiler,
     'nvidia': NvidiaCompiler,
     'cuda': CudaCompiler,
+    'nvcc': CudaCompiler,
     'osx': ClangCompiler,
     'intel': OneapiCompiler,
     'icx': OneapiCompiler,

devito/data/allocators.py

@@ -92,8 +92,12 @@ def initialize(cls):
         return
 
     def alloc(self, shape, dtype, padding=0):
-        datasize = int(reduce(mul, shape))
-        ctype = dtype_to_ctype(dtype)
+        # For complex number, allocate double the size of its real/imaginary part
+        alloc_dtype = dtype(0).real.__class__
+        c_scale = 2 if np.issubdtype(dtype, np.complexfloating) else 1
+
+        datasize = int(reduce(mul, shape) * c_scale)
+        ctype = dtype_to_ctype(alloc_dtype)
 
         # Add padding, if any
         try:

devito/finite_differences/differentiable.py

@@ -14,8 +14,7 @@
 from devito.logger import warning
 from devito.tools import (as_tuple, filter_ordered, flatten, frozendict,
                           infer_dtype, is_integer, split)
-from devito.types import (Array, DimensionTuple, Evaluable, Indexed,
-                          StencilDimension)
+from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 
 __all__ = ['Differentiable', 'DiffDerivative', 'IndexDerivative', 'EvalDerivative',
            'Weights']
@@ -68,7 +67,7 @@ def grid(self):
 
     @cached_property
     def dtype(self):
-        dtypes = {f.dtype for f in self.find(Indexed)} - {None}
+        dtypes = {f.dtype for f in self._functions} - {None}
         return infer_dtype(dtypes)
 
     @cached_property

devito/ir/iet/visitors.py

@@ -13,6 +13,7 @@
 from sympy import IndexedBase
 from sympy.core.function import Application
 
+from devito.parameters import configuration, switchconfig
 from devito.exceptions import VisitorException
 from devito.ir.iet.nodes import (Node, Iteration, Expression, ExpressionBundle,
                                  Call, Lambda, BlankLine, Section, ListMajor)
@@ -176,7 +177,7 @@ class CGen(Visitor):
 
     def __init__(self, *args, compiler=None, **kwargs):
         super().__init__(*args, **kwargs)
-        self._compiler = compiler
+        self._compiler = compiler or configuration['compiler']
 
     # The following mappers may be customized by subclasses (that is,
     # backend-specific CGen-erators)
@@ -188,6 +189,16 @@ def __init__(self, *args, compiler=None, **kwargs):
     }
     _restrict_keyword = 'restrict'
 
+    @property
+    def compiler(self):
+        return self._compiler
+
+    def visit(self, o, *args, **kwargs):
+        # Make sure the visitor always is within the generating compiler
+        # in case the configuration is accessed
+        with switchconfig(compiler=self.compiler.name):
+            return super().visit(o, *args, **kwargs)
+
     def _gen_struct_decl(self, obj, masked=()):
         """
         Convert ctypes.Struct -> cgen.Structure.
@@ -376,10 +387,11 @@ def visit_tuple(self, o):
     def visit_PointerCast(self, o):
         f = o.function
         i = f.indexed
+        cstr = i._C_typedata
 
         if f.is_PointerArray:
             # lvalue
-            lvalue = c.Value(i._C_typedata, '**%s' % f.name)
+            lvalue = c.Value(cstr, '**%s' % f.name)
 
             # rvalue
             if isinstance(o.obj, ArrayObject):
@@ -388,7 +400,7 @@ def visit_PointerCast(self, o):
                 v = f._C_name
             else:
                 assert False
-            rvalue = '(%s**) %s' % (i._C_typedata, v)
+            rvalue = '(%s**) %s' % (cstr, v)
 
         else:
             # lvalue
@@ -399,10 +411,10 @@ def visit_PointerCast(self, o):
             if o.flat is None:
                 shape = ''.join("[%s]" % ccode(i) for i in o.castshape)
                 rshape = '(*)%s' % shape
-                lvalue = c.Value(i._C_typedata, '(*restrict %s)%s' % (v, shape))
+                lvalue = c.Value(cstr, '(*restrict %s)%s' % (v, shape))
             else:
                 rshape = '*'
-                lvalue = c.Value(i._C_typedata, '*%s' % v)
+                lvalue = c.Value(cstr, '*%s' % v)
             if o.alignment and f._data_alignment:
                 lvalue = c.AlignedAttribute(f._data_alignment, lvalue)
 
@@ -415,30 +427,30 @@ def visit_PointerCast(self, o):
                 else:
                     assert False
 
-                rvalue = '(%s %s) %s->%s' % (i._C_typedata, rshape, f._C_name, v)
+                rvalue = '(%s %s) %s->%s' % (cstr, rshape, f._C_name, v)
             else:
                 if isinstance(o.obj, Pointer):
                     v = o.obj.name
                 else:
                     v = f._C_name
 
-                rvalue = '(%s %s) %s' % (i._C_typedata, rshape, v)
+                rvalue = '(%s %s) %s' % (cstr, rshape, v)
 
         return c.Initializer(lvalue, rvalue)
 
     def visit_Dereference(self, o):
         a0, a1 = o.functions
         if a1.is_PointerArray or a1.is_TempFunction:
             i = a1.indexed
+            cstr = i._C_typedata
             if o.flat is None:
                 shape = ''.join("[%s]" % ccode(i) for i in a0.symbolic_shape[1:])
-                rvalue = '(%s (*)%s) %s[%s]' % (i._C_typedata, shape, a1.name,
+                rvalue = '(%s (*)%s) %s[%s]' % (cstr, shape, a1.name,
                                                 a1.dim.name)
-                lvalue = c.Value(i._C_typedata,
-                                 '(*restrict %s)%s' % (a0.name, shape))
+                lvalue = c.Value(cstr, '(*restrict %s)%s' % (a0.name, shape))
             else:
-                rvalue = '(%s *) %s[%s]' % (i._C_typedata, a1.name, a1.dim.name)
-                lvalue = c.Value(i._C_typedata, '*restrict %s' % a0.name)
+                rvalue = '(%s *) %s[%s]' % (cstr, a1.name, a1.dim.name)
+                lvalue = c.Value(cstr, '*restrict %s' % a0.name)
             if a0._data_alignment:
                 lvalue = c.AlignedAttribute(a0._data_alignment, lvalue)
         else:
@@ -590,7 +602,7 @@ def visit_MultiTraversable(self, o):
         return c.Collection(body)
 
     def visit_UsingNamespace(self, o):
-        return c.Statement('using namespace %s' % ccode(o.namespace))
+        return c.Statement('using namespace %s' % str(o.namespace))
 
     def visit_Lambda(self, o):
         body = []

devito/operator/operator.py

@@ -469,6 +469,8 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs):
 
         # Lower IET to a target-specific IET
         graph = Graph(iet, **kwargs)
+
+        # Specialize
         graph = cls._specialize_iet(graph, **kwargs)
 
         # Instrument the IET for C-level profiling
@@ -1347,7 +1349,8 @@ def parse_kwargs(**kwargs):
             raise InvalidOperator("Illegal `compiler=%s`" % str(compiler))
         kwargs['compiler'] = compiler_registry[compiler](platform=kwargs['platform'],
                                                          language=kwargs['language'],
-                                                         mpi=configuration['mpi'])
+                                                         mpi=configuration['mpi'],
+                                                         name=compiler)
     elif any([platform, language]):
         kwargs['compiler'] =\
             configuration['compiler'].__new_with__(platform=kwargs['platform'],

devito/passes/clusters/factorization.py

@@ -173,7 +173,6 @@ def _collect_nested(expr):
     Recursion helper for `collect_nested`.
     """
     # Return semantic (rebuilt expression, factorization candidates)
-
     if expr.is_Number:
         return expr, {'coeffs': expr}
     elif expr.is_Function:

devito/passes/iet/__init__.py

@@ -8,3 +8,4 @@
 from .instrument import *  # noqa
 from .languages import *  # noqa
 from .errors import *  # noqa
+from .dtypes import *  # noqa

devito/passes/iet/definitions.py

@@ -12,6 +12,7 @@
 from devito.ir import (Block, Call, Definition, DummyExpr, Return, EntryFunction,
                        FindSymbols, MapExprStmts, Transformer, make_callable)
 from devito.passes import is_gpu_create
+from devito.passes.iet.dtypes import lower_complex
 from devito.passes.iet.engine import iet_pass
 from devito.passes.iet.langbase import LangBB
 from devito.symbolics import (Byref, DefFunction, FieldFromPointer, IndexedPointer,
@@ -73,10 +74,12 @@ class DataManager:
     The language used to express data allocations, deletions, and host-device transfers.
     """
 
-    def __init__(self, rcompile=None, sregistry=None, platform=None, **kwargs):
+    def __init__(self, rcompile=None, sregistry=None, platform=None,
+                 compiler=None, **kwargs):
         self.rcompile = rcompile
         self.sregistry = sregistry
         self.platform = platform
+        self.compiler = compiler
 
     def _alloc_object_on_low_lat_mem(self, site, obj, storage):
         """
@@ -409,12 +412,18 @@ def place_casts(self, iet, **kwargs):
 
         return iet, {}
 
+    @iet_pass
+    def make_langtypes(self, iet):
+        iet, metadata = lower_complex(iet, self.lang, self.compiler)
+        return iet, metadata
+
     def process(self, graph):
         """
         Apply the `place_definitions` and `place_casts` passes.
         """
         self.place_definitions(graph, globs=set())
         self.place_casts(graph)
+        self.make_langtypes(graph)
 
 
 class DeviceAwareDataManager(DataManager):
@@ -564,6 +573,7 @@ def process(self, graph):
         self.place_devptr(graph)
         self.place_bundling(graph, writes_input=graph.writes_input)
         self.place_casts(graph)
+        self.make_langtypes(graph)
 
 
 def make_zero_init(obj):

devito/passes/iet/dtypes.py

@@ -0,0 +1,54 @@
+import numpy as np
+import ctypes
+
+from devito.ir import FindSymbols, Uxreplace
+
+__all__ = ['lower_complex']
+
+
+def lower_complex(iet, lang, compiler):
+    """
+    Add headers for complex arithmetic
+    """
+    # Check if there is complex numbers that always take dtype precedence
+    types = {f.dtype for f in FindSymbols().visit(iet)
+             if not issubclass(f.dtype, ctypes._Pointer)}
+
+    if not any(np.issubdtype(d, np.complexfloating) for d in types):
+        return iet, {}
+
+    lib = (lang['header-complex'],)
+
+    metadata = {}
+    if lang.get('complex-namespace') is not None:
+        metadata['namespaces'] = lang['complex-namespace']
+
+    # Some languges such as c++11 need some extra arithmetic definitions
+    if lang.get('def-complex'):
+        dest = compiler.get_jit_dir()
+        hfile = dest.joinpath('complex_arith.h')
+        with open(str(hfile), 'w') as ff:
+            ff.write(str(lang['def-complex']))
+        lib += (str(hfile),)
+
+    iet = _complex_dtypes(iet, lang)
+    metadata['includes'] = lib
+
+    return iet, metadata
+
+
+def _complex_dtypes(iet, lang):
+    """
+    Lower dtypes to language specific types
+    """
+    mapper = {}
+
+    for s in FindSymbols('indexeds|basics|symbolics').visit(iet):
+        if s.dtype in lang['types']:
+            mapper[s] = s._rebuild(dtype=lang['types'][s.dtype])
+
+    body = Uxreplace(mapper).visit(iet.body)
+    params = Uxreplace(mapper).visit(iet.parameters)
+    iet = iet._rebuild(body=body, parameters=params)
+
+    return iet

devito/passes/iet/langbase.py

@@ -31,6 +31,9 @@ def __getitem__(self, k):
             raise NotImplementedError("Missing required mapping for `%s`" % k)
         return self.mapper[k]
 
+    def get(self, k):
+        return self.mapper.get(k)
+
 
 class LangBB(metaclass=LangMeta):
 
@@ -200,6 +203,14 @@ def initialize(self, iet, options=None):
         """
         return iet, {}
 
+    @iet_pass
+    def make_langtypes(self, iet):
+        """
+        An `iet_pass` which transforms an IET such that the target language
+        types are introduced.
+        """
+        return iet, {}
+
     @property
     def Region(self):
         return self.lang.Region