Merge pull request #3 from ZoeLeibowitz/core_petsc

Core PETSc objects/functions.

devito/passes/iet/definitions.py

@@ -352,10 +352,10 @@ def place_definitions(self, iet, globs=None, **kwargs):
                 elif globs is not None:
                     # Track, to be handled by the EntryFunction being a global obj!
                     globs.add(i)
-            elif i.is_ObjectArray:
-                self._alloc_object_array_on_low_lat_mem(iet, i, storage)
             elif i.is_PointerArray:
                 self._alloc_pointed_array_on_high_bw_mem(iet, i, storage)
+            else:
+                self._alloc_object_array_on_low_lat_mem(iet, i, storage)
 
         # Handle postponed global objects
         includes = set()

devito/types/__init__.py

@@ -6,6 +6,7 @@
 from .object import *  # noqa
 from .lazy import *  # noqa
 from .misc import *  # noqa
+from .petsc import *  # noqa
 
 # Needed both within and outside Devito
 from .dimension import *  # noqa

devito/types/array.py

@@ -19,6 +19,12 @@ class ArrayBasic(AbstractFunction):
 
     is_ArrayBasic = True
 
+    __rkwargs__ = AbstractFunction.__rkwargs__ + ('is_const',)
+
+    def __init_finalize__(self, *args, **kwargs):
+        super().__init_finalize__(*args, **kwargs)
+        self._is_const = kwargs.get('is_const', False)
+
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):
         dimensions = kwargs['dimensions']
@@ -48,6 +54,10 @@ def shape(self):
     def shape_allocated(self):
         return self.symbolic_shape
 
+    @property
+    def is_const(self):
+        return self._is_const
+
 
 class Array(ArrayBasic):
 

devito/types/petsc.py

@@ -0,0 +1,92 @@
+from devito.tools import CustomDtype
+from devito.types import LocalObject
+from devito.types.array import ArrayBasic
+import numpy as np
+from cached_property import cached_property
+
+
+class DM(LocalObject):
+    """
+    PETSc Data Management object (DM).
+    """
+    dtype = CustomDtype('DM')
+
+
+class Mat(LocalObject):
+    """
+    PETSc Matrix object (Mat).
+    """
+    dtype = CustomDtype('Mat')
+
+
+class Vec(LocalObject):
+    """
+    PETSc Vector object (Vec).
+    """
+    dtype = CustomDtype('Vec')
+
+
+class PetscMPIInt(LocalObject):
+    """
+    PETSc datatype used to represent `int` parameters
+    to MPI functions.
+    """
+    dtype = CustomDtype('PetscMPIInt')
+
+
+class KSP(LocalObject):
+    """
+    PETSc KSP : Linear Systems Solvers.
+    Manages Krylov Methods.
+    """
+    dtype = CustomDtype('KSP')
+
+
+class PC(LocalObject):
+    """
+    PETSc object that manages all preconditioners (PC).
+    """
+    dtype = CustomDtype('PC')
+
+
+class KSPConvergedReason(LocalObject):
+    """
+    PETSc object - reason a Krylov method was determined
+    to have converged or diverged.
+    """
+    dtype = CustomDtype('KSPConvergedReason')
+
+
+class PETScArray(ArrayBasic):
+    """
+    PETScArrays are generated by the compiler only and represent
+    a customised variant of ArrayBasic. They are designed to
+    avoid generating a cast in the low-level code.
+    """
+
+    _data_alignment = False
+
+    @classmethod
+    def __dtype_setup__(cls, **kwargs):
+        return kwargs.get('dtype', np.float32)
+
+    @cached_property
+    def _C_ctype(self):
+        petsc_type = dtype_to_petsctype(self.dtype)
+        modifier = '*' * len(self.dimensions)
+        return CustomDtype(petsc_type, modifier=modifier)
+
+    @property
+    def _C_name(self):
+        return self.name
+
+
+def dtype_to_petsctype(dtype):
+    """Map numpy types to PETSc datatypes."""
+
+    return {
+        np.int32: 'PetscInt',
+        np.float32: 'PetscScalar',
+        np.int64: 'PetscInt',
+        np.float64: 'PetscScalar'
+    }[dtype]

tests/test_petsc.py

@@ -0,0 +1,65 @@
+from devito import Grid
+from devito.ir.iet import Call, ElementalFunction, Definition, DummyExpr
+from devito.passes.iet.languages.C import CDataManager
+from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP,
+                          PC, KSPConvergedReason, PETScArray)
+import numpy as np
+
+
+def test_petsc_local_object():
+    """
+    Test C++ support for PETSc LocalObjects.
+    """
+    lo0 = DM('da')
+    lo1 = Mat('A')
+    lo2 = Vec('x')
+    lo3 = PetscMPIInt('size')
+    lo4 = KSP('ksp')
+    lo5 = PC('pc')
+    lo6 = KSPConvergedReason('reason')
+
+    iet = Call('foo', [lo0, lo1, lo2, lo3, lo4, lo5, lo6])
+    iet = ElementalFunction('foo', iet, parameters=())
+
+    dm = CDataManager(sregistry=None)
+    iet = CDataManager.place_definitions.__wrapped__(dm, iet)[0]
+
+    assert 'DM da;' in str(iet)
+    assert 'Mat A;' in str(iet)
+    assert 'Vec x;' in str(iet)
+    assert 'PetscMPIInt size;' in str(iet)
+    assert 'KSP ksp;' in str(iet)
+    assert 'PC pc;' in str(iet)
+    assert 'KSPConvergedReason reason;' in str(iet)
+
+
+def test_petsc_functions():
+    """
+    Test C++ support for PETScArrays.
+    """
+    grid = Grid((2, 2))
+    x, y = grid.dimensions
+
+    ptr0 = PETScArray(name='ptr0', dimensions=grid.dimensions, dtype=np.float32)
+    ptr1 = PETScArray(name='ptr1', dimensions=grid.dimensions, dtype=np.float32,
+                      is_const=True)
+    ptr2 = PETScArray(name='ptr2', dimensions=grid.dimensions, dtype=np.float64,
+                      is_const=True)
+    ptr3 = PETScArray(name='ptr3', dimensions=grid.dimensions, dtype=np.int32)
+    ptr4 = PETScArray(name='ptr4', dimensions=grid.dimensions, dtype=np.int64,
+                      is_const=True)
+
+    defn0 = Definition(ptr0)
+    defn1 = Definition(ptr1)
+    defn2 = Definition(ptr2)
+    defn3 = Definition(ptr3)
+    defn4 = Definition(ptr4)
+
+    expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1)
+
+    assert str(defn0) == 'PetscScalar**restrict ptr0;'
+    assert str(defn1) == 'const PetscScalar**restrict ptr1;'
+    assert str(defn2) == 'const PetscScalar**restrict ptr2;'
+    assert str(defn3) == 'PetscInt**restrict ptr3;'
+    assert str(defn4) == 'const PetscInt**restrict ptr4;'
+    assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;'