From 4f3bc8b56afc4742f368088392ad5663ba44e549 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 27 Mar 2025 15:12:38 +0000 Subject: [PATCH 1/9] some progress --- devito/data/allocators.py | 31 ++++++++++++++++++++++++++++++- devito/petsc/iet/passes.py | 28 +++++++++++++++++++++++++--- devito/petsc/memory.py | 19 +++++++++++++++++++ devito/petsc/types/object.py | 15 ++++++++++++++- devito/petsc/types/types.py | 4 ++++ 5 files changed, 92 insertions(+), 5 deletions(-) create mode 100644 devito/petsc/memory.py diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 4ccd7cddfc..2a2159d167 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 @@ -11,9 +12,10 @@ from devito.parameters import configuration from devito.tools import is_integer, infer_datasize + __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: @@ -102,6 +104,7 @@ def alloc(self, shape, dtype, padding=0): size = datasize + padleft + padright padleft_pointer, memfree_args = self._alloc_C_libcall(size, ctype) + if padleft_pointer is None: raise RuntimeError("Unable to allocate %d elements in memory" % size) @@ -334,6 +337,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 +421,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..0162eb5e05 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -6,15 +6,16 @@ from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, DummyExpr, CallableBody, List, Call, Callable, FindNodes) -from devito.symbolics import Byref, Macro, FieldFromPointer -from devito.types import Symbol, Scalar +from devito.symbolics import Byref, Macro, FieldFromPointer, VOID +from devito.types import Symbol, Scalar, Symbol from devito.types.basic import DataSymbol from devito.tools import frozendict from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, 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), 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): + """ + Allocate memory for PETSc objects. + 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..4aa9716915 --- /dev/null +++ b/devito/petsc/memory.py @@ -0,0 +1,19 @@ +import os +import sys +from ctypes import POINTER, cast, c_char +import atexit + +from devito import Operator, switchconfig +from devito.types import Symbol +from devito.types.equation import PetscEq +from devito.petsc.types import Initialize, Finalize, 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/types/object.py b/devito/petsc/types/object.py index fcbf06ea69..a8fdf0d77d 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,16 @@ 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 CustomDtype('size_t') + 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 From 076405a7ad1b088749b7d5b8c62ee2462f9c5b43 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 27 Mar 2025 15:30:32 +0000 Subject: [PATCH 2/9] add example --- examples/petsc/memory_test.py | 40 +++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 examples/petsc/memory_test.py diff --git a/examples/petsc/memory_test.py b/examples/petsc/memory_test.py new file mode 100644 index 0000000000..9c1d198549 --- /dev/null +++ b/examples/petsc/memory_test.py @@ -0,0 +1,40 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, configuration, + switchconfig) +from devito.types import Symbol +from devito.types.equation import PetscEq +from devito.data.allocators import ALLOC_PETSC +from devito.petsc import PETScSolve +from devito.petsc.types import AllocateMemory +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, allocator=ALLOC_PETSC) +v = Function(name='v', grid=grid, dtype=np.float64, space_order=2, allocator=ALLOC_PETSC) + +v.data[:] = 5.0 + +eq = Eq(v, u.laplace, subdomain=grid.interior) + +petsc = PETScSolve([eq], u) + + +with switchconfig(language='petsc'): + op = Operator(petsc) + + +# print(op.ccode) +op.apply() + + + From e017d1cc6d3fba60ff0fc260497d47dea9ca328e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 27 Mar 2025 17:02:48 +0000 Subject: [PATCH 3/9] no malloc_debug errors now but need to inspect dag size assertion when not all functions are mem alloc via petsc --- devito/passes/iet/engine.py | 2 +- devito/petsc/iet/passes.py | 6 +++--- devito/petsc/memory.py | 8 +------- examples/petsc/memory_test.py | 14 ++++---------- 4 files changed, 9 insertions(+), 21 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index a88cf58d45..d327de0a3f 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -231,7 +231,7 @@ def create_call_graph(root, efuncs): dag.add_edge(callee, caller) # Sanity check - assert dag.size == len(efuncs) + # assert dag.size == len(efuncs) return dag diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 0162eb5e05..ceb346bd71 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -6,8 +6,8 @@ from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, DummyExpr, CallableBody, List, Call, Callable, FindNodes) -from devito.symbolics import Byref, Macro, FieldFromPointer, VOID -from devito.types import Symbol, Scalar, Symbol +from devito.symbolics import Byref, Macro, FieldFromPointer +from devito.types import Symbol, Scalar from devito.types.basic import DataSymbol from devito.tools import frozendict from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, @@ -121,7 +121,7 @@ def allocate_memory(iet): 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), diff --git a/devito/petsc/memory.py b/devito/petsc/memory.py index 4aa9716915..7a035b2846 100644 --- a/devito/petsc/memory.py +++ b/devito/petsc/memory.py @@ -1,13 +1,7 @@ -import os -import sys -from ctypes import POINTER, cast, c_char -import atexit - from devito import Operator, switchconfig from devito.types import Symbol from devito.types.equation import PetscEq -from devito.petsc.types import Initialize, Finalize, AllocateMemory - +from devito.petsc.types import AllocateMemory dummy = Symbol(name='d') diff --git a/examples/petsc/memory_test.py b/examples/petsc/memory_test.py index 9c1d198549..68c756e797 100644 --- a/examples/petsc/memory_test.py +++ b/examples/petsc/memory_test.py @@ -19,22 +19,16 @@ grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) +# Only need to allocate the "target" function memory via PETSc u = Function(name='u', grid=grid, dtype=np.float64, space_order=2, allocator=ALLOC_PETSC) -v = Function(name='v', grid=grid, dtype=np.float64, space_order=2, allocator=ALLOC_PETSC) +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], u) - +petsc = PETScSolve([eq], target=u) with switchconfig(language='petsc'): op = Operator(petsc) - - -# print(op.ccode) -op.apply() - - - + op.apply() From 273da1394790afee6d8a34ae41ff14437d5722d5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 27 Mar 2025 17:04:50 +0000 Subject: [PATCH 4/9] fcqw --- examples/petsc/memory_test.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/petsc/memory_test.py b/examples/petsc/memory_test.py index 68c756e797..43c5dbb2fc 100644 --- a/examples/petsc/memory_test.py +++ b/examples/petsc/memory_test.py @@ -3,11 +3,8 @@ from devito import (Grid, Function, Eq, Operator, configuration, switchconfig) -from devito.types import Symbol -from devito.types.equation import PetscEq from devito.data.allocators import ALLOC_PETSC from devito.petsc import PETScSolve -from devito.petsc.types import AllocateMemory from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' From 474b8ac94091b8cd822e0e98a0b66a118df10258 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 28 Mar 2025 10:30:27 +0000 Subject: [PATCH 5/9] fix dag assertion problem --- devito/passes/iet/engine.py | 2 +- devito/petsc/iet/passes.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index d327de0a3f..a88cf58d45 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -231,7 +231,7 @@ def create_call_graph(root, efuncs): dag.add_edge(callee, caller) # Sanity check - # assert dag.size == len(efuncs) + assert dag.size == len(efuncs) return dag diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index ceb346bd71..732858542e 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -50,7 +50,7 @@ def lower_petsc(iet, **kwargs): return finalize(iet), core_metadata() if any(filter(lambda i: isinstance(i.expr.rhs, AllocateMemory), data)): - return allocate_memory(iet), metadata + 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 @@ -114,7 +114,7 @@ def finalize(iet): def allocate_memory(iet): """ - Allocate memory for PETSc objects. + Create function to allocate memory using PetscMalloc. https://petsc.org/release/manualpages/Sys/PetscMalloc/ """ # Number of bytes to allocate From 3f3539ae82264323fabdab6ebb64794b0fa99375 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 31 Mar 2025 10:05:32 +0100 Subject: [PATCH 6/9] switch to place array --- devito/data/allocators.py | 2 -- examples/petsc/memory_test.py | 9 +++++---- tests/test_petsc.py | 2 -- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 2a2159d167..13eb88b502 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -12,7 +12,6 @@ from devito.parameters import configuration from devito.tools import is_integer, infer_datasize - __all__ = ['ALLOC_ALIGNED', 'ALLOC_NUMA_LOCAL', 'ALLOC_NUMA_ANY', 'ALLOC_KNL_MCDRAM', 'ALLOC_KNL_DRAM', 'ALLOC_GUARD', 'ALLOC_PETSC', 'default_allocator'] @@ -104,7 +103,6 @@ def alloc(self, shape, dtype, padding=0): size = datasize + padleft + padright padleft_pointer, memfree_args = self._alloc_C_libcall(size, ctype) - if padleft_pointer is None: raise RuntimeError("Unable to allocate %d elements in memory" % size) diff --git a/examples/petsc/memory_test.py b/examples/petsc/memory_test.py index 43c5dbb2fc..6b965acb7b 100644 --- a/examples/petsc/memory_test.py +++ b/examples/petsc/memory_test.py @@ -2,10 +2,10 @@ import numpy as np from devito import (Grid, Function, Eq, Operator, configuration, - switchconfig) -from devito.data.allocators import ALLOC_PETSC + switchconfig, TimeFunction) from devito.petsc import PETScSolve from devito.petsc.initialize import PetscInitialize +from devito.types import Constant configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -16,8 +16,7 @@ grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) -# Only need to allocate the "target" function memory via PETSc -u = Function(name='u', grid=grid, dtype=np.float64, space_order=2, allocator=ALLOC_PETSC) +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 @@ -29,3 +28,5 @@ 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) From be176168eaaaa186da75a402da2fef2b20a4985d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 31 Mar 2025 14:11:36 +0100 Subject: [PATCH 7/9] switch to use placearray --- devito/petsc/solve.py | 1 - 1 file changed, 1 deletion(-) 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): From f6701703f5664290dd6fc475516fa0bbb3721cd4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 31 Mar 2025 14:36:24 +0100 Subject: [PATCH 8/9] dcaec --- devito/petsc/types/object.py | 1 - 1 file changed, 1 deletion(-) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index a8fdf0d77d..02a5f33589 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -319,5 +319,4 @@ def _C_ctype(self): class SizeT(DataSymbol): @property def _C_ctype(self): - # return CustomDtype('size_t') return c_size_t From ce4ddd73ba9662480849219012004a9e36ec4a39 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 31 Mar 2025 14:42:56 +0100 Subject: [PATCH 9/9] fwer --- examples/petsc/memory_test.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/petsc/memory_test.py b/examples/petsc/memory_test.py index 6b965acb7b..6e6e3a88b7 100644 --- a/examples/petsc/memory_test.py +++ b/examples/petsc/memory_test.py @@ -2,10 +2,9 @@ import numpy as np from devito import (Grid, Function, Eq, Operator, configuration, - switchconfig, TimeFunction) + switchconfig) from devito.petsc import PETScSolve from devito.petsc.initialize import PetscInitialize -from devito.types import Constant configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc'