PETSc Memory Allocator

devito/data/allocators.py

@@ -4,6 +4,7 @@
 import mmap
 import os
 import sys
+from pathlib import Path
 
 import numpy as np
 
@@ -13,7 +14,7 @@
 
 __all__ = ['ALLOC_ALIGNED', 'ALLOC_NUMA_LOCAL', 'ALLOC_NUMA_ANY',
            'ALLOC_KNL_MCDRAM', 'ALLOC_KNL_DRAM', 'ALLOC_GUARD',
-           'default_allocator']
+           'ALLOC_PETSC', 'default_allocator']
 
 
 class AbstractMemoryAllocator:
@@ -334,6 +335,31 @@ def put_local(self):
         return self._node == 'local'
 
 
+class PetscMemoryAllocator(MemoryAllocator):
+    """
+    """
+    @classmethod
+    def initialize(cls):
+        from devito.petsc.utils import core_metadata
+        metadata = core_metadata()
+        lib_dir = Path(metadata['lib_dirs'][-1])
+
+        try:
+            cls.lib = ctypes.CDLL(lib_dir/'libpetsc.so')
+        except OSError:
+            cls.lib = None
+
+    def _alloc_C_libcall(self, size, ctype):
+        if not self.available():
+            raise RuntimeError("...")
+        c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype))
+        c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p)
+
+        from devito.petsc.memory import op_memory
+        op_memory(m=c_bytesize, result=ctypes.byref(c_pointer))
+        return c_pointer, (c_pointer, )
+
+
 class DataReference(MemoryAllocator):
 
     """
@@ -393,6 +419,7 @@ def alloc(self, shape, dtype, padding=0):
 ALLOC_KNL_MCDRAM = NumaAllocator(1)
 ALLOC_NUMA_ANY = NumaAllocator('any')
 ALLOC_NUMA_LOCAL = NumaAllocator('local')
+ALLOC_PETSC = PetscMemoryAllocator()
 
 custom_allocators = {
     'fallback': ALLOC_ALIGNED,

devito/petsc/iet/passes.py

@@ -14,7 +14,8 @@
                                 PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
                                 CallbackPointerIS, CallbackPointerDM, JacobianStruct,
-                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
+                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol,
+                                AllocateMemory, VoidPtrPtr, SizeT)
 from devito.petsc.types.macros import petsc_func_begin_user
 from devito.petsc.iet.nodes import PetscMetaData
 from devito.petsc.utils import core_metadata, petsc_languages
@@ -48,6 +49,9 @@ def lower_petsc(iet, **kwargs):
     if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)):
         return finalize(iet), core_metadata()
 
+    if any(filter(lambda i: isinstance(i.expr.rhs, AllocateMemory), data)):
+        return allocate_memory(iet), core_metadata()
+
     unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()}
     # Assumption is that all solves are on the same grid
     if len(unique_grids) > 1:
@@ -108,6 +112,24 @@ def finalize(iet):
     return iet._rebuild(body=finalize_body)
 
 
+def allocate_memory(iet):
+    """
+    Create function to allocate memory using PetscMalloc.
+    https://petsc.org/release/manualpages/Sys/PetscMalloc/
+    """
+    # Number of bytes to allocate
+    m = SizeT(name="m")
+    # Memory allocated
+    result = VoidPtrPtr(name='result')
+
+    allocate_body = petsc_call('PetscMalloc', [m, result])
+    allocate_body = CallableBody(
+        body=(petsc_func_begin_user, allocate_body),
+        retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+    )
+    return iet._rebuild(body=allocate_body, parameters=[m, result])
+
+
 def make_core_petsc_calls(objs, **kwargs):
     call_mpi = petsc_call_mpi('MPI_Comm_size', [objs['comm'], Byref(objs['size'])])
 

devito/petsc/memory.py

@@ -0,0 +1,13 @@
+from devito import Operator, switchconfig
+from devito.types import Symbol
+from devito.types.equation import PetscEq
+from devito.petsc.types import AllocateMemory
+
+
+dummy = Symbol(name='d')
+
+with switchconfig(language='petsc'):
+    op_memory = Operator(
+        [PetscEq(dummy, AllocateMemory(dummy))],
+        name='kernel_allocate_memory', opt='noop'
+    )

devito/petsc/solve.py

@@ -69,7 +69,6 @@ def build_function_eqns(self, eq, target, arrays):
         b, F_target, targets = separate_eqn(eq, target)
         formfunc = self.make_formfunc(eq, F_target, arrays, targets)
         formrhs = self.make_rhs(eq, b, arrays)
-
         return (formfunc, formrhs)
 
     def build_matvec_eqns(self, eq, target, arrays):

devito/petsc/types/object.py

@@ -1,4 +1,4 @@
-from ctypes import POINTER, c_char
+from ctypes import POINTER, c_char, c_size_t, c_void_p
 from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType
 from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension,
                           TimeDimension, ArrayObject, CustomDimension)
@@ -308,3 +308,15 @@ class ArgvSymbol(DataSymbol):
     @property
     def _C_ctype(self):
         return POINTER(POINTER(c_char))
+
+
+class VoidPtrPtr(DataSymbol):
+    @property
+    def _C_ctype(self):
+        return (POINTER(c_void_p))
+
+
+class SizeT(DataSymbol):
+    @property
+    def _C_ctype(self):
+        return c_size_t

devito/petsc/types/types.py

@@ -25,6 +25,10 @@ class Finalize(MetaData):
     pass
 
 
+class AllocateMemory(MetaData):
+    pass
+
+
 class LinearSolveExpr(MetaData):
     """
     A symbolic expression passed through the Operator, containing the metadata

examples/petsc/memory_test.py

@@ -0,0 +1,31 @@
+import os
+import numpy as np
+
+from devito import (Grid, Function, Eq, Operator, configuration,
+                    switchconfig)
+from devito.petsc import PETScSolve
+from devito.petsc.initialize import PetscInitialize
+configuration['compiler'] = 'custom'
+os.environ['CC'] = 'mpicc'
+
+PetscInitialize()
+
+nx = 81
+ny = 81
+
+grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64)
+
+u = Function(name='u', grid=grid, dtype=np.float64, space_order=2)
+v = Function(name='v', grid=grid, dtype=np.float64, space_order=2)
+
+v.data[:] = 5.0
+
+eq = Eq(v, u.laplace, subdomain=grid.interior)
+
+petsc = PETScSolve([eq], target=u)
+
+with switchconfig(language='petsc'):
+    op = Operator(petsc)
+    op.apply()
+
+print(op.ccode)

tests/test_petsc.py

@@ -835,8 +835,6 @@ def test_coupled_vs_non_coupled(self):
         enorm2 = norm(e)
         gnorm2 = norm(g)
 
-        print('enorm1:', enorm1)
-        print('enorm2:', enorm2)
         assert np.isclose(enorm1, enorm2, rtol=1e-16)
         assert np.isclose(gnorm1, gnorm2, rtol=1e-16)