compiler: add extra platforms and language to the custom compiler

devito/arch/compiler.py

@@ -172,19 +172,20 @@ def __init__(self):
     """
 
     fields = {'cc', 'ld'}
+    _cpp = False
 
     def __init__(self, **kwargs):
-        super(Compiler, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
         self.__lookup_cmds__()
 
         self.suffix = kwargs.get('suffix')
         if not kwargs.get('mpi'):
-            self.cc = self.CC if kwargs.get('cpp', False) is False else self.CXX
+            self.cc = self.CC if self._cpp is False else self.CXX
             self.cc = self.cc if self.suffix is None else ('%s-%s' %
                                                            (self.cc, self.suffix))
         else:
-            self.cc = self.MPICC if kwargs.get('cpp', False) is False else self.MPICXX
+            self.cc = self.MPICC if self._cpp is False else self.MPICXX
         self.ld = self.cc  # Wanted by the superclass
 
         self.cflags = ['-O3', '-g', '-fPIC', '-Wall', '-std=c99']
@@ -196,7 +197,7 @@ def __init__(self, **kwargs):
         self.defines = []
         self.undefines = []
 
-        self.src_ext = 'c' if kwargs.get('cpp', False) is False else 'cpp'
+        self.src_ext = 'c' if self._cpp is False else 'cpp'
 
         if platform.system() == "Linux":
             self.so_ext = '.so'
@@ -216,6 +217,11 @@ def __init__(self, **kwargs):
             # Knowing the version may still be useful to pick supported flags
             self.version = sniff_compiler_version(self.CC)
 
+        self.__init_finalize__(**kwargs)
+
+    def __init_finalize__(self, **kwargs):
+        pass
+
     def __new_with__(self, **kwargs):
         """
         Create a new Compiler from an existing one, inherenting from it
@@ -394,9 +400,7 @@ def add_ldflags(self, flags):
 
 class GNUCompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
+    def __init_finalize__(self, **kwargs):
         platform = kwargs.pop('platform', configuration['platform'])
 
         self.cflags += ['-Wno-unused-result',
@@ -443,9 +447,8 @@ def __lookup_cmds__(self):
 
 class ArmCompiler(GNUCompiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
+    def __init_finalize__(self, **kwargs):
+        GNUCompiler.__init_finalize__(self, **kwargs)
         platform = kwargs.pop('platform', configuration['platform'])
 
         # Graviton flag
@@ -455,8 +458,7 @@ def __init__(self, *args, **kwargs):
 
 class ClangCompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init_finalize__(self, **kwargs):
 
         self.cflags += ['-Wno-unused-result', '-Wno-unused-variable']
         if not configuration['safe-math']:
@@ -522,8 +524,7 @@ class AOMPCompiler(Compiler):
 
     """AMD's fork of Clang for OpenMP offloading on both AMD and NVidia cards."""
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init_finalize__(self, **kwargs):
 
         language = kwargs.pop('language', configuration['language'])
         platform = kwargs.pop('platform', configuration['platform'])
@@ -556,8 +557,7 @@ def __lookup_cmds__(self):
 
 class DPCPPCompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init_finalize__(self, **kwargs):
 
         self.cflags += ['-qopenmp', '-fopenmp-targets=spir64']
 
@@ -572,8 +572,9 @@ def __lookup_cmds__(self):
 
 class PGICompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, cpp=True, **kwargs)
+    _cpp = True
+
+    def __init_finalize__(self, **kwargs):
 
         self.cflags.remove('-std=c99')
         self.cflags.remove('-O3')
@@ -618,8 +619,9 @@ def __lookup_cmds__(self):
 
 class CudaCompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, cpp=True, **kwargs)
+    _cpp = True
+
+    def __init_finalize__(self, **kwargs):
 
         self.cflags.remove('-std=c99')
         self.cflags.remove('-Wall')
@@ -683,8 +685,9 @@ def __lookup_cmds__(self):
 
 class HipCompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, cpp=True, **kwargs)
+    _cpp = True
+
+    def __init_finalize__(self, **kwargs):
 
         self.cflags.remove('-std=c99')
         self.cflags.remove('-Wall')
@@ -712,8 +715,7 @@ def __lookup_cmds__(self):
 
 class IntelCompiler(Compiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init_finalize__(self, **kwargs):
 
         platform = kwargs.pop('platform', configuration['platform'])
         language = kwargs.pop('language', configuration['language'])
@@ -771,8 +773,8 @@ def __lookup_cmds__(self):
 
 class IntelKNLCompiler(IntelCompiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init_finalize__(self, **kwargs):
+        IntelCompiler.__init_finalize__(self, **kwargs)
 
         language = kwargs.pop('language', configuration['language'])
 
@@ -784,8 +786,8 @@ def __init__(self, *args, **kwargs):
 
 class OneapiCompiler(IntelCompiler):
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init_finalize__(self, **kwargs):
+        IntelCompiler.__init_finalize__(self, **kwargs)
 
         platform = kwargs.pop('platform', configuration['platform'])
         language = kwargs.pop('language', configuration['language'])
@@ -839,33 +841,47 @@ class CustomCompiler(Compiler):
 
     def __new__(cls, *args, **kwargs):
         platform = kwargs.pop('platform', configuration['platform'])
+        language = kwargs.pop('language', configuration['language'])
 
-        if any(i in environ for i in ['CC', 'CXX', 'CFLAGS', 'LDFLAGS']):
-            obj = super().__new__(cls)
-            obj.__init__(*args, **kwargs)
-            return obj
-        elif platform is M1:
-            return ClangCompiler(*args, **kwargs)
+        if platform is M1:
+            _base = ClangCompiler
+        elif platform is INTELGPUX:
+            _base = OneapiCompiler
+        elif platform is NVIDIAX:
+            if language == 'cuda':
+                _base = CudaCompiler
+            else:
+                _base = NvidiaCompiler
+        elif platform is AMDGPUX:
+            if language == 'hip':
+                _base = HipCompiler
+            else:
+                _base = AOMPCompiler
         else:
-            return GNUCompiler(*args, **kwargs)
-
-    def __init__(self, *args, **kwargs):
-        super(CustomCompiler, self).__init__(*args, **kwargs)
+            _base = GNUCompiler
 
-        default = '-O3 -g -march=native -fPIC -Wall -std=c99'
-        self.cflags = environ.get('CFLAGS', default).split(' ')
-        self.ldflags = environ.get('LDFLAGS', '-shared').split(' ')
+        obj = super().__new__(cls)
+        # Keep base to initialize accordingly
+        obj._base = _base
+        obj._cpp = _base._cpp
 
-        language = kwargs.pop('language', configuration['language'])
+        return obj
 
-        if language == 'openmp':
-            self.ldflags += environ.get('OMP_LDFLAGS', '-fopenmp').split(' ')
+    def __init_finalize__(self, **kwargs):
+        self._base.__init_finalize__(self, **kwargs)
+        # Update cflags
+        extrac = environ.get('CFLAGS', '').split(' ')
+        self.cflags = filter_ordered(self.cflags + extrac)
+        # Update ldflags
+        extrald = environ.get('LDFLAGS', '').split(' ')
+        self.ldflags = filter_ordered(self.ldflags + extrald)
 
     def __lookup_cmds__(self):
-        self.CC = environ.get('CC', 'gcc')
-        self.CXX = environ.get('CXX', 'g++')
-        self.MPICC = environ.get('MPICC', 'mpicc')
-        self.MPICXX = environ.get('MPICXX', 'mpicxx')
+        self._base.__lookup_cmds__(self)
+        self.CC = environ.get('CC', self.CC)
+        self.CXX = environ.get('CXX', self.CXX)
+        self.MPICC = environ.get('MPICC', self.MPICC)
+        self.MPICXX = environ.get('MPICXX', self.MPICXX)
 
 
 compiler_registry = {

devito/mpi/distributed.py

@@ -185,7 +185,7 @@ class Distributor(AbstractDistributor):
     """
 
     def __init__(self, shape, dimensions, input_comm=None, topology=None):
-        super(Distributor, self).__init__(shape, dimensions)
+        super().__init__(shape, dimensions)
 
         if configuration['mpi']:
             # First time we enter here, we make sure MPI is initialized
@@ -426,7 +426,7 @@ class SparseDistributor(AbstractDistributor):
     """
 
     def __init__(self, npoint, dimension, distributor):
-        super(SparseDistributor, self).__init__(npoint, dimension)
+        super().__init__(npoint, dimension)
         self._distributor = distributor
 
         # The dimension decomposition
@@ -523,7 +523,7 @@ def __init__(self, neighborhood):
         self._entries = [i for i in neighborhood if isinstance(i, tuple)]
 
         fields = [(''.join(j.name[0] for j in i), c_int) for i in self.entries]
-        super(MPINeighborhood, self).__init__('nb', 'neighborhood', fields)
+        super().__init__('nb', 'neighborhood', fields)
 
     @property
     def entries(self):
@@ -552,7 +552,7 @@ def _C_typedecl(self):
                                    for i, j in groups])
 
     def _arg_defaults(self):
-        values = super(MPINeighborhood, self)._arg_defaults()
+        values = super()._arg_defaults()
         for name, i in zip(self.fields, self.entries):
             setattr(values[self.name]._obj, name, self.neighborhood[i])
         return values

devito/types/dense.py

@@ -741,7 +741,8 @@ def _C_get_field(self, region, dim, side=None):
 
     def _halo_exchange(self):
         """Perform the halo exchange with the neighboring processes."""
-        if not MPI.Is_initialized() or MPI.COMM_WORLD.size == 1:
+        if not MPI.Is_initialized() or MPI.COMM_WORLD.size == 1 or \
+                not configuration['mpi']:
             # Nothing to do
             return
         if MPI.COMM_WORLD.size > 1 and self._distributor is None:

tests/test_mpi.py

@@ -419,6 +419,21 @@ def test_local_indices(self, shape, expected):
         assert all(i == slice(*j)
                    for i, j in zip(f.local_indices, expected[grid.distributor.myrank]))
 
+    @pytest.mark.parallel(mode=4)
+    @pytest.mark.parametrize('shape', [(1,), (2, 3), (4, 5, 6)])
+    def test_mpi4py_nodevmpi(self, shape):
+
+        with switchconfig(mpi=False):
+            # Mimic external mpi init
+            MPI.Init()
+            # Check that internal Function work correctly
+            grid = Grid(shape=shape)
+            f = Function(name="f", grid=grid, space_order=1)
+            assert f.data.shape == shape
+            assert f.data_with_halo.shape == tuple(s+2 for s in shape)
+            assert f.data._local.shape == shape
+            MPI.Finalize()
+
 
 class TestSparseFunction(object):