diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 4ccd7cddfc..13eb88b502 100644 --- a/devito/data/allocators.py +++ b/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, diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 3c73bfd391..732858542e 100644 --- a/devito/petsc/iet/passes.py +++ b/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'])]) diff --git a/devito/petsc/memory.py b/devito/petsc/memory.py new file mode 100644 index 0000000000..7a035b2846 --- /dev/null +++ b/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' + ) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 334df4c802..29bad259cf 100644 --- a/devito/petsc/solve.py +++ b/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): diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index fcbf06ea69..02a5f33589 100644 --- a/devito/petsc/types/object.py +++ b/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 diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 782c5b7112..4dce232e1d 100644 --- a/devito/petsc/types/types.py +++ b/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 diff --git a/examples/petsc/memory_test.py b/examples/petsc/memory_test.py new file mode 100644 index 0000000000..6e6e3a88b7 --- /dev/null +++ b/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) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 47d6faabcd..71e11065ef 100644 --- a/tests/test_petsc.py +++ b/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)