From a9801e7403d60dc97b516b6fecea9ea24891df05 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 10:46:24 +0000 Subject: [PATCH 001/151] compiler: Add FuncPtrCall --- devito/ir/iet/nodes.py | 15 ++++++++++++++- devito/ir/iet/visitors.py | 17 +++++++++++++++++ tests/test_iet.py | 20 ++++++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index d45c9af939..3430247244 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -28,7 +28,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer'] + 'CallableBody', 'Transfer', 'FuncPtrCall'] # First-class IET nodes @@ -358,6 +358,19 @@ def writes(self): return self._writes +class FuncPtrCall(Call): + + """ + Function Pointer Argument to a Call. + """ + def __init__(self, name, return_type, parameter_type, **kwargs): + + super().__init__(name=name) + + self.return_type = return_type + self.parameter_type = parameter_type + + class Expression(ExprStmt, Node): """ diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 9749825407..38d64cd374 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -520,6 +520,12 @@ def visit_Call(self, o, nested_call=False): else: return c.Initializer(c.Value(rettype, retobj._C_name), call) + def visit_FuncPtrCall(self, o, nested_call=False): + name = o.name + return_type = o.return_type + parameter_type = o.parameter_type + return FuncPtrArg(name, return_type, parameter_type) + def visit_Conditional(self, o): try: then_body, else_body = self._blankline_logic(o.children) @@ -1427,3 +1433,14 @@ def sorted_efuncs(efuncs): CommCallable: 1 } return sorted_priority(efuncs, priority) + + +class FuncPtrArg(c.Generable): + + def __init__(self, name, return_type, parameter_type): + self.name = name + self.return_type = return_type + self.parameter_type = parameter_type + + def generate(self): + yield "(%s (*)(%s))%s" % (self.return_type, self.parameter_type, self.name) diff --git a/tests/test_iet.py b/tests/test_iet.py index ff963d5518..cc3322c73d 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -128,6 +128,26 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(expected) +def test_funcptrcall_cgen(): + + a = Symbol('a') + b = Symbol('b') + foo0 = Callable('foo0', Definition(a), 'void', parameters=[b]) + foo0_arg = FuncPtrCall(foo0.name, foo0.retval, 'int') + code0 = CGen().visit(foo0_arg) + assert str(code0) == '(void (*)(int))foo0' + + # test nested calls with a FuncPtrCall as an argument + call = Call('foo1', [ + Call('foo2', [foo0_arg]) + ]) + code1 = CGen().visit(call) + assert str(code1) == 'foo1(foo2((void (*)(int))foo0));' + + callees = FindNodes(Call).visit(call) + assert len(callees) == 3 + + def test_list_denesting(): l0 = List(header=cgen.Line('a'), body=List(header=cgen.Line('b'))) l1 = l0._rebuild(body=List(header=cgen.Line('c'))) From 3f73add6492d965aa911efd4b108b003a320f75e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 11:46:01 +0000 Subject: [PATCH 002/151] misc: Update docstring --- devito/ir/iet/nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 3430247244..ce4ca967c4 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -361,7 +361,7 @@ def writes(self): class FuncPtrCall(Call): """ - Function Pointer Argument to a Call. + Function Pointer Call. """ def __init__(self, name, return_type, parameter_type, **kwargs): From ccd10b9e31d12f0fc5be5461ab977906b17690e5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 11:04:27 +0000 Subject: [PATCH 003/151] compiler: Edit names, move position of Callback and other small edits --- devito/ir/iet/nodes.py | 35 ++++++++++++++++++++++------------- devito/ir/iet/visitors.py | 20 +++++++++----------- tests/test_iet.py | 9 +++++---- 3 files changed, 36 insertions(+), 28 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index ce4ca967c4..853d4532d1 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -358,19 +358,6 @@ def writes(self): return self._writes -class FuncPtrCall(Call): - - """ - Function Pointer Call. - """ - def __init__(self, name, return_type, parameter_type, **kwargs): - - super().__init__(name=name) - - self.return_type = return_type - self.parameter_type = parameter_type - - class Expression(ExprStmt, Node): """ @@ -1146,6 +1133,28 @@ def defines(self): return tuple(self.parameters) +class Callback(Call): + + """ + Callback as a function pointer. + + Parameters + ---------- + name : str + The name of the callback. + retval : str + The return type of the callback. + param_types : str or list of str + The return type for each argument of the callback. + """ + def __init__(self, name, retval, param_types): + + super().__init__(name=name) + + self.retval = retval + self.param_types = as_tuple(param_types) + + class Section(List): """ diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 38d64cd374..d2f5d9db3b 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -520,12 +520,6 @@ def visit_Call(self, o, nested_call=False): else: return c.Initializer(c.Value(rettype, retobj._C_name), call) - def visit_FuncPtrCall(self, o, nested_call=False): - name = o.name - return_type = o.return_type - parameter_type = o.parameter_type - return FuncPtrArg(name, return_type, parameter_type) - def visit_Conditional(self, o): try: then_body, else_body = self._blankline_logic(o.children) @@ -630,6 +624,9 @@ def visit_Lambda(self, o): (', '.join(captures), ', '.join(decls), ''.join(extra))) return LambdaCollection([top, c.Block(body)]) + def visit_Callback(self, o, nested_call=False): + return CallbackArg(o.name, o.retval, o.param_types) + def visit_HaloSpot(self, o): body = flatten(self._visit(i) for i in o.children) return c.Collection(body) @@ -1435,12 +1432,13 @@ def sorted_efuncs(efuncs): return sorted_priority(efuncs, priority) -class FuncPtrArg(c.Generable): +class CallbackArg(c.Generable): - def __init__(self, name, return_type, parameter_type): + def __init__(self, name, retval, param_types): self.name = name - self.return_type = return_type - self.parameter_type = parameter_type + self.retval = retval + self.param_types = param_types def generate(self): - yield "(%s (*)(%s))%s" % (self.return_type, self.parameter_type, self.name) + param_types_str = ', '.join([str(t) for t in self.param_types]) + yield "(%s (*)(%s))%s" % (self.retval, param_types_str, self.name) diff --git a/tests/test_iet.py b/tests/test_iet.py index cc3322c73d..b5cff00be0 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -10,7 +10,8 @@ from devito.ir.iet import (Call, Callable, Conditional, DeviceCall, DummyExpr, Iteration, List, KernelLaunch, Lambda, ElementalFunction, CGen, FindSymbols, filter_iterations, make_efunc, - retrieve_iteration_tree, Transformer) + retrieve_iteration_tree, Transformer, Callback, FindNodes, + Definition) from devito.ir import SymbolRegistry from devito.passes.iet.engine import Graph from devito.passes.iet.languages.C import CDataManager @@ -128,16 +129,16 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(expected) -def test_funcptrcall_cgen(): +def test_callback_cgen(): a = Symbol('a') b = Symbol('b') foo0 = Callable('foo0', Definition(a), 'void', parameters=[b]) - foo0_arg = FuncPtrCall(foo0.name, foo0.retval, 'int') + foo0_arg = Callback(foo0.name, foo0.retval, 'int') code0 = CGen().visit(foo0_arg) assert str(code0) == '(void (*)(int))foo0' - # test nested calls with a FuncPtrCall as an argument + # test nested calls with a Callback as an argument. call = Call('foo1', [ Call('foo2', [foo0_arg]) ]) From f34ae6d83a171df00f6f2910f530b299467f5dd8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 5 Jan 2024 11:20:48 +0000 Subject: [PATCH 004/151] misc: Update docstring for Callback --- devito/ir/iet/nodes.py | 8 ++++++++ tests/test_iet.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 853d4532d1..c2ecebcd69 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1146,7 +1146,15 @@ class Callback(Call): The return type of the callback. param_types : str or list of str The return type for each argument of the callback. + + Notes + ----- + The reason Callback is an IET type rather than a SymPy type is + due to the fact that, when represented at the SymPy level, the IET + engine fails to bind the callback to a specific Call. Consequently, + errors occur during the creation of the call graph. """ + def __init__(self, name, retval, param_types): super().__init__(name=name) diff --git a/tests/test_iet.py b/tests/test_iet.py index b5cff00be0..1c36830084 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -138,7 +138,7 @@ def test_callback_cgen(): code0 = CGen().visit(foo0_arg) assert str(code0) == '(void (*)(int))foo0' - # test nested calls with a Callback as an argument. + # Test nested calls with a Callback as an argument. call = Call('foo1', [ Call('foo2', [foo0_arg]) ]) From ae28dac6f17adbc5c4922c5f8c59f5784d951697 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sat, 6 Jan 2024 10:37:33 +0000 Subject: [PATCH 005/151] misc: Clean up blank lines --- devito/ir/iet/nodes.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index c2ecebcd69..c0f2a32c4f 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1134,7 +1134,6 @@ def defines(self): class Callback(Call): - """ Callback as a function pointer. @@ -1156,9 +1155,7 @@ class Callback(Call): """ def __init__(self, name, retval, param_types): - super().__init__(name=name) - self.retval = retval self.param_types = as_tuple(param_types) From 20fd2d12403d1c7011d0cbdf537898d12fed0ece Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 006/151] misc: Split Dockerfile.cpu into .cpu and .intel --- docker/Dockerfile.intel | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index e013fccc03..1fab7147d6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,6 +35,7 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu +<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -50,6 +51,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +======= +# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal +RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg +RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list + +RUN apt-get update -y && apt-get dist-upgrade -y && \ + apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ + intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ + libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -76,7 +87,11 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD # ICX OpenMP image +======= +# ICX image +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -93,6 +108,35 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD +======= +# ICX hpc image +############################################################## +FROM oneapi as icx-hpc + +# Install both icc and icx to avoid missing dependencies +RUN apt-get update -y && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic + +# Missig components +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager +RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ + chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ + rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# Devito config +ENV DEVITO_ARCH="icx" +ENV DEVITO_LANGUAGE="openmp" +# MPICC compiler for mpi4py +ENV MPICC=mpiicc +ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' + +############################################################## +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -107,6 +151,12 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl +<<<<<<< HEAD +======= +# NOTE: the name of this file ends with ".cpu" but this is a GPU image. +# It then feels a bit akward, so some restructuring might be needed + +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 2a6b923d48d87169155090f3e96f6dad6ff9fe4f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 27 Nov 2023 19:16:56 +0000 Subject: [PATCH 007/151] try petscfunc from abstractfunc --- devito/passes/iet/petsc.py | 136 ++++++++++++++++++++++++++++++++ devito/tools/dtypes_lowering.py | 20 ++++- 2 files changed, 154 insertions(+), 2 deletions(-) create mode 100644 devito/passes/iet/petsc.py diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py new file mode 100644 index 0000000000..a78159212d --- /dev/null +++ b/devito/passes/iet/petsc.py @@ -0,0 +1,136 @@ +from devito.types.basic import AbstractSymbol, AbstractFunction +from devito.tools import petsc_type_to_ctype, dtype_to_ctype, dtype_to_cstr +import numpy as np +from sympy import Expr +from ctypes import POINTER + + + +class PETScObject(AbstractSymbol): + """ + PETScObjects. + """ + @property + def _C_ctype(self): + ctype = petsc_type_to_ctype(self._dtype) + return type(self._dtype, (ctype,), {}) + + + +# may need to also inherit from Expr +class PETScFunction(AbstractFunction): + """ + PETScFunctions. + """ + + @classmethod + def __dtype_setup__(cls, **kwargs): + grid = kwargs.get('grid') + dtype = kwargs.get('dtype') + if dtype is not None: + return dtype + elif grid is not None: + return grid.dtype + else: + return np.float32 + + @property + def dimensions(self): + """Tuple of Dimensions representing the object indices.""" + return self._dimensions + + @classmethod + def __indices_setup__(cls, **kwargs): + grid = kwargs.get('grid') + shape = kwargs.get('shape') + dimensions = kwargs.get('dimensions') + + if dimensions is None and shape is None and grid is None: + return (), () + + elif grid is None: + if dimensions is None: + raise TypeError("Need either `grid` or `dimensions`") + elif dimensions is None: + dimensions = grid.dimensions + + return dimensions, dimensions + + # @classmethod + # def __shape_setup__(cls, **kwargs): + # grid = kwargs.get('grid') + # dimensions = kwargs.get('dimensions') + # shape = kwargs.get('shape') + + # if dimensions is None and shape is None and grid is None: + # return None + + # elif grid is None: + # if shape is None: + # raise TypeError("Need either `grid` or `shape`") + # elif shape is None: + # if dimensions is not None and dimensions != grid.dimensions: + # raise TypeError("Need `shape` as not all `dimensions` are in `grid`") + # shape = grid.shape + # elif dimensions is None: + # raise TypeError("`dimensions` required if both `grid` and " + # "`shape` are provided") + + # return shape + + # @classmethod + # def __indices_setup__(cls, *args, **kwargs): + # grid = kwargs.get('grid') + # dimensions = kwargs.get('dimensions') + # if grid is None: + # if dimensions is None: + # raise TypeError("Need either `grid` or `dimensions`") + # elif dimensions is None: + # dimensions = grid.dimensions + + # return tuple(dimensions), tuple(dimensions) + + @property + def _C_ctype(self): + # from IPython import embed; embed() + ctypename = 'Petsc%s' % dtype_to_cstr(self._dtype).capitalize() + ctype = dtype_to_ctype(self.dtype) + r = POINTER(type(ctypename, (ctype,), {})) + for n in range(len(self.dimensions)-1): + r = POINTER(r) + return r + + # @property + # def _C_ctype(self): + # # from IPython import embed; embed() + # ctypename = 'Petsc%s' % dtype_to_cstr(self._dtype).capitalize() + # ctype = dtype_to_ctype(self.dtype) + # r = type(ctypename, (ctype,), {}) + # # for n in range(len(self.dimensions)-1): + # # r = POINTER(r) + # from IPython import embed; embed() + # return r + + @property + def _C_name(self): + return self.name + + + + +from devito.ir import Definition +da = PETScObject('da', dtype='DM') +tmp = Definition(da) +# print(tmp) + +from devito import * +grid = Grid((2, 2)) +x, y = grid.dimensions +# pointer and const functionality +ptr1 = PETScFunction(name='ptr1', dtype=np.int32, dimensions=grid.dimensions, shape=grid.shape, is_const=True) +defn1 = Definition(ptr1) +from IPython import embed; embed() +print(str(defn1)) + + + diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index b5b564a4d7..4df760c525 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -13,8 +13,7 @@ 'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype', 'dtype_to_cstr', 'dtype_to_ctype', 'dtype_to_mpitype', 'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', - 'is_external_ctype', 'infer_dtype', 'CustomDtype'] - + 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'petsc_type_to_ctype'] # *** Custom np.dtypes @@ -313,3 +312,20 @@ def infer_dtype(dtypes): else: # E.g., mixed integer arithmetic return max(dtypes, key=lambda i: np.dtype(i).itemsize, default=None) + + +def petsc_type_to_ctype(dtype): + """ + Map PETSc types to ctypes type. + """ + return { + # 'PetscInt': ctypes.c_int, + # 'PetscScalar': ctypes.c_float, + 'Mat': ctypes.c_void_p, + 'Vec': ctypes.c_void_p, + 'PetscErrorCode': ctypes.c_int, + 'KSP': ctypes.c_void_p, + 'PC': ctypes.c_void_p, + 'DM': ctypes.c_void_p, + 'PetscMPIInt': ctypes.c_int + }[dtype] From 2c0d33e0e5bbe891ec8db09d313a70e2bf237e4b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 30 Nov 2023 23:58:44 +0000 Subject: [PATCH 008/151] latest. --- devito/passes/iet/petsc.py | 100 +++++++++++++------------------- devito/tools/dtypes_lowering.py | 24 +++----- 2 files changed, 48 insertions(+), 76 deletions(-) diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index a78159212d..1a589d9d50 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -1,21 +1,30 @@ from devito.types.basic import AbstractSymbol, AbstractFunction -from devito.tools import petsc_type_to_ctype, dtype_to_ctype, dtype_to_cstr +from devito.tools import petsc_type_to_ctype, dtype_to_ctype, dtype_to_cstr, dtype_to_petsctype import numpy as np from sympy import Expr -from ctypes import POINTER +from devito.types import LocalObject +from ctypes import POINTER, c_void_p +from devito.ir import Definition -class PETScObject(AbstractSymbol): - """ - PETScObjects. - """ - @property - def _C_ctype(self): - ctype = petsc_type_to_ctype(self._dtype) - return type(self._dtype, (ctype,), {}) +# class PETScObject(AbstractSymbol): +# @property +# def _C_ctype(self): +# ctype = petsc_type_to_ctype(self._dtype) +# return type(self._dtype, (ctype,), {}) +class PETScDM(LocalObject): + dtype = type('DM', (c_void_p,), {}) + + +da = PETScDM('da') +defn1 = Definition(da) +print(defn1) + + + # may need to also inherit from Expr class PETScFunction(AbstractFunction): @@ -36,25 +45,7 @@ def __dtype_setup__(cls, **kwargs): @property def dimensions(self): - """Tuple of Dimensions representing the object indices.""" return self._dimensions - - @classmethod - def __indices_setup__(cls, **kwargs): - grid = kwargs.get('grid') - shape = kwargs.get('shape') - dimensions = kwargs.get('dimensions') - - if dimensions is None and shape is None and grid is None: - return (), () - - elif grid is None: - if dimensions is None: - raise TypeError("Need either `grid` or `dimensions`") - elif dimensions is None: - dimensions = grid.dimensions - - return dimensions, dimensions # @classmethod # def __shape_setup__(cls, **kwargs): @@ -78,38 +69,26 @@ def __indices_setup__(cls, **kwargs): # return shape - # @classmethod - # def __indices_setup__(cls, *args, **kwargs): - # grid = kwargs.get('grid') - # dimensions = kwargs.get('dimensions') - # if grid is None: - # if dimensions is None: - # raise TypeError("Need either `grid` or `dimensions`") - # elif dimensions is None: - # dimensions = grid.dimensions + @classmethod + def __indices_setup__(cls, *args, **kwargs): + grid = kwargs.get('grid') + dimensions = kwargs.get('dimensions') + if grid is None: + if dimensions is None: + raise TypeError("Need either `grid` or `dimensions`") + elif dimensions is None: + dimensions = grid.dimensions - # return tuple(dimensions), tuple(dimensions) + return tuple(dimensions), tuple(dimensions) @property def _C_ctype(self): - # from IPython import embed; embed() - ctypename = 'Petsc%s' % dtype_to_cstr(self._dtype).capitalize() + petsc_type = dtype_to_petsctype(self.dtype) ctype = dtype_to_ctype(self.dtype) - r = POINTER(type(ctypename, (ctype,), {})) - for n in range(len(self.dimensions)-1): + r = type(petsc_type, (ctype,), {}) + for n in range(len(self.dimensions)): r = POINTER(r) return r - - # @property - # def _C_ctype(self): - # # from IPython import embed; embed() - # ctypename = 'Petsc%s' % dtype_to_cstr(self._dtype).capitalize() - # ctype = dtype_to_ctype(self.dtype) - # r = type(ctypename, (ctype,), {}) - # # for n in range(len(self.dimensions)-1): - # # r = POINTER(r) - # from IPython import embed; embed() - # return r @property def _C_name(self): @@ -118,19 +97,18 @@ def _C_name(self): -from devito.ir import Definition -da = PETScObject('da', dtype='DM') -tmp = Definition(da) +# from devito.ir import Definition +# da = PETScObject('da', dtype='DM') +# tmp = Definition(da) # print(tmp) from devito import * grid = Grid((2, 2)) x, y = grid.dimensions -# pointer and const functionality -ptr1 = PETScFunction(name='ptr1', dtype=np.int32, dimensions=grid.dimensions, shape=grid.shape, is_const=True) -defn1 = Definition(ptr1) -from IPython import embed; embed() -print(str(defn1)) +ptr1 = PETScFunction(name='ptr1', dtype=np.float32, dimensions=grid.dimensions, shape=grid.shape) +defn2 = Definition(ptr1) +# from IPython import embed; embed() +print(str(defn2)) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 4df760c525..c1d917e3c8 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -13,7 +13,7 @@ 'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype', 'dtype_to_cstr', 'dtype_to_ctype', 'dtype_to_mpitype', 'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', - 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'petsc_type_to_ctype'] + 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'dtype_to_petsctype'] # *** Custom np.dtypes @@ -312,20 +312,14 @@ def infer_dtype(dtypes): else: # E.g., mixed integer arithmetic return max(dtypes, key=lambda i: np.dtype(i).itemsize, default=None) - -def petsc_type_to_ctype(dtype): - """ - Map PETSc types to ctypes type. - """ + +def dtype_to_petsctype(dtype): + """Map numpy types to PETSc datatypes.""" + return { - # 'PetscInt': ctypes.c_int, - # 'PetscScalar': ctypes.c_float, - 'Mat': ctypes.c_void_p, - 'Vec': ctypes.c_void_p, - 'PetscErrorCode': ctypes.c_int, - 'KSP': ctypes.c_void_p, - 'PC': ctypes.c_void_p, - 'DM': ctypes.c_void_p, - 'PetscMPIInt': ctypes.c_int + np.int32: 'PetscInt', + np.float32: 'PetscScalar', + # np.int64: + # np.float64: }[dtype] From 12ab5bd6f7bf36bb109cfd72847aca4898d32cba Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 4 Dec 2023 14:44:33 +0000 Subject: [PATCH 009/151] remove data allignment --- devito/passes/iet/petsc.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index 1a589d9d50..b2913e3ce7 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -31,6 +31,7 @@ class PETScFunction(AbstractFunction): """ PETScFunctions. """ + _data_alignment = False @classmethod def __dtype_setup__(cls, **kwargs): @@ -106,9 +107,11 @@ def _C_name(self): grid = Grid((2, 2)) x, y = grid.dimensions ptr1 = PETScFunction(name='ptr1', dtype=np.float32, dimensions=grid.dimensions, shape=grid.shape) -defn2 = Definition(ptr1) + +from IPython import embed; embed() +# defn2 = Definition(ptr1) # from IPython import embed; embed() -print(str(defn2)) +# print(/defn2) From 2e940ad36d151b9ff333d9412e28d509df1b689b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 4 Dec 2023 17:47:08 +0000 Subject: [PATCH 010/151] new petsc types --- devito/passes/iet/petsc.py | 117 -------------------------------- devito/tools/dtypes_lowering.py | 2 - devito/types/__init__.py | 1 + devito/types/petsc.py | 77 +++++++++++++++++++++ tests/test_petsc.py | 54 +++++++++++++++ 5 files changed, 132 insertions(+), 119 deletions(-) create mode 100644 devito/types/petsc.py create mode 100644 tests/test_petsc.py diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index b2913e3ce7..e69de29bb2 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -1,117 +0,0 @@ -from devito.types.basic import AbstractSymbol, AbstractFunction -from devito.tools import petsc_type_to_ctype, dtype_to_ctype, dtype_to_cstr, dtype_to_petsctype -import numpy as np -from sympy import Expr -from devito.types import LocalObject -from ctypes import POINTER, c_void_p -from devito.ir import Definition - - - -# class PETScObject(AbstractSymbol): -# @property -# def _C_ctype(self): -# ctype = petsc_type_to_ctype(self._dtype) -# return type(self._dtype, (ctype,), {}) - - -class PETScDM(LocalObject): - dtype = type('DM', (c_void_p,), {}) - - -da = PETScDM('da') -defn1 = Definition(da) -print(defn1) - - - - -# may need to also inherit from Expr -class PETScFunction(AbstractFunction): - """ - PETScFunctions. - """ - _data_alignment = False - - @classmethod - def __dtype_setup__(cls, **kwargs): - grid = kwargs.get('grid') - dtype = kwargs.get('dtype') - if dtype is not None: - return dtype - elif grid is not None: - return grid.dtype - else: - return np.float32 - - @property - def dimensions(self): - return self._dimensions - - # @classmethod - # def __shape_setup__(cls, **kwargs): - # grid = kwargs.get('grid') - # dimensions = kwargs.get('dimensions') - # shape = kwargs.get('shape') - - # if dimensions is None and shape is None and grid is None: - # return None - - # elif grid is None: - # if shape is None: - # raise TypeError("Need either `grid` or `shape`") - # elif shape is None: - # if dimensions is not None and dimensions != grid.dimensions: - # raise TypeError("Need `shape` as not all `dimensions` are in `grid`") - # shape = grid.shape - # elif dimensions is None: - # raise TypeError("`dimensions` required if both `grid` and " - # "`shape` are provided") - - # return shape - - @classmethod - def __indices_setup__(cls, *args, **kwargs): - grid = kwargs.get('grid') - dimensions = kwargs.get('dimensions') - if grid is None: - if dimensions is None: - raise TypeError("Need either `grid` or `dimensions`") - elif dimensions is None: - dimensions = grid.dimensions - - return tuple(dimensions), tuple(dimensions) - - @property - def _C_ctype(self): - petsc_type = dtype_to_petsctype(self.dtype) - ctype = dtype_to_ctype(self.dtype) - r = type(petsc_type, (ctype,), {}) - for n in range(len(self.dimensions)): - r = POINTER(r) - return r - - @property - def _C_name(self): - return self.name - - - - -# from devito.ir import Definition -# da = PETScObject('da', dtype='DM') -# tmp = Definition(da) -# print(tmp) - -from devito import * -grid = Grid((2, 2)) -x, y = grid.dimensions -ptr1 = PETScFunction(name='ptr1', dtype=np.float32, dimensions=grid.dimensions, shape=grid.shape) - -from IPython import embed; embed() -# defn2 = Definition(ptr1) -# from IPython import embed; embed() -# print(/defn2) - - - diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index c1d917e3c8..272141d821 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -320,6 +320,4 @@ def dtype_to_petsctype(dtype): return { np.int32: 'PetscInt', np.float32: 'PetscScalar', - # np.int64: - # np.float64: }[dtype] diff --git a/devito/types/__init__.py b/devito/types/__init__.py index 6ec8bdfd16..ff024f24be 100644 --- a/devito/types/__init__.py +++ b/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 diff --git a/devito/types/petsc.py b/devito/types/petsc.py new file mode 100644 index 0000000000..f4e999e334 --- /dev/null +++ b/devito/types/petsc.py @@ -0,0 +1,77 @@ +from devito.types.basic import AbstractFunction +from devito.tools import dtype_to_petsctype, CustomDtype +import numpy as np +from devito.types import LocalObject + + +class DM(LocalObject): + dtype = CustomDtype('DM') + + +class Mat(LocalObject): + dtype = CustomDtype('Mat') + + +class Vec(LocalObject): + dtype = CustomDtype('Vec') + + +class PetscMPIInt(LocalObject): + dtype = CustomDtype('PetscMPIInt') + + +class KSP(LocalObject): + dtype = CustomDtype('KSP') + + +class PC(LocalObject): + dtype = CustomDtype('PC') + + +class KSPConvergedReason(LocalObject): + dtype = CustomDtype('KSPConvergedReason') + + +class PETScFunction(AbstractFunction): + """ + PETScFunctions. + """ + _data_alignment = False + + @classmethod + def __dtype_setup__(cls, **kwargs): + grid = kwargs.get('grid') + dtype = kwargs.get('dtype') + if dtype is not None: + return dtype + elif grid is not None: + return grid.dtype + else: + return np.float32 + + @classmethod + def __indices_setup__(cls, *args, **kwargs): + grid = kwargs.get('grid') + dimensions = kwargs.get('dimensions') + if grid is None: + if dimensions is None: + raise TypeError("Need either `grid` or `dimensions`") + elif dimensions is None: + dimensions = grid.dimensions + + return tuple(dimensions), tuple(dimensions) + + @property + def dimensions(self): + return self._dimensions + + @property + def _C_ctype(self): + petsc_type = dtype_to_petsctype(self.dtype) + modifier = '*' * len(self.dimensions) + customtype = CustomDtype(petsc_type, modifier=modifier) + return customtype + + @property + def _C_name(self): + return self.name diff --git a/tests/test_petsc.py b/tests/test_petsc.py new file mode 100644 index 0000000000..2ea5cbf9ba --- /dev/null +++ b/tests/test_petsc.py @@ -0,0 +1,54 @@ +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, PETScFunction) +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 PETScFunctions. + """ + grid = Grid((2, 2)) + x, y = grid.dimensions + + ptr0 = PETScFunction(name='ptr0', dtype=np.float32, dimensions=grid.dimensions, + shape=grid.shape) + ptr1 = PETScFunction(name='ptr1', dtype=np.float32, grid=grid) + + defn0 = Definition(ptr0) + defn1 = Definition(ptr1) + + expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1) + + assert str(defn0) == 'PetscScalar**restrict ptr0;' + assert str(defn1) == 'PetscScalar**restrict ptr1;' + assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' From 0815c703e86beb6eb6e47ee20458e405198e6b64 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 7 Dec 2023 10:58:32 +0000 Subject: [PATCH 011/151] add const func in petscfunction --- devito/tools/dtypes_lowering.py | 2 ++ devito/types/petsc.py | 34 +++++++++++++++++++++++++++++++++ tests/test_petsc.py | 3 +++ 3 files changed, 39 insertions(+) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 272141d821..6f54b2d0a0 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -320,4 +320,6 @@ def dtype_to_petsctype(dtype): return { np.int32: 'PetscInt', np.float32: 'PetscScalar', + np.int64: 'PetscInt', + np.float64: 'PetscScalar' }[dtype] diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f4e999e334..f1e1ea6660 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -5,30 +5,54 @@ 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') @@ -38,6 +62,12 @@ class PETScFunction(AbstractFunction): """ _data_alignment = False + def __init_finalize__(self, *args, **kwargs): + + super(PETScFunction, self).__init_finalize__(*args, **kwargs) + + self._is_const = kwargs.get('is_const', False) + @classmethod def __dtype_setup__(cls, **kwargs): grid = kwargs.get('grid') @@ -75,3 +105,7 @@ def _C_ctype(self): @property def _C_name(self): return self.name + + @property + def is_const(self): + return self._is_const diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 2ea5cbf9ba..1021662783 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -43,12 +43,15 @@ def test_petsc_functions(): ptr0 = PETScFunction(name='ptr0', dtype=np.float32, dimensions=grid.dimensions, shape=grid.shape) ptr1 = PETScFunction(name='ptr1', dtype=np.float32, grid=grid) + ptr2 = PETScFunction(name='ptr2', dtype=np.float32, grid=grid, is_const=True) defn0 = Definition(ptr0) defn1 = Definition(ptr1) + defn2 = Definition(ptr2) expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1) assert str(defn0) == 'PetscScalar**restrict ptr0;' assert str(defn1) == 'PetscScalar**restrict ptr1;' + assert str(defn2) == 'const PetscScalar**restrict ptr2;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' From cc38a90e15e0c97279a3bf9e7c8b691445bee279 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 13:37:34 +0000 Subject: [PATCH 012/151] dsl: Simplify PETScFunction class and add diff dtypes inside test_petsc_functions --- devito/types/petsc.py | 11 +++-------- tests/test_petsc.py | 8 +++++++- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f1e1ea6660..f07cb298ad 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -27,7 +27,7 @@ class Vec(LocalObject): class PetscMPIInt(LocalObject): """ - PETSc datatype used to represent ‘int’ parameters + PETSc datatype used to represent `int` parameters to MPI functions. """ dtype = CustomDtype('PetscMPIInt') @@ -64,7 +64,7 @@ class PETScFunction(AbstractFunction): def __init_finalize__(self, *args, **kwargs): - super(PETScFunction, self).__init_finalize__(*args, **kwargs) + super().__init_finalize__(*args, **kwargs) self._is_const = kwargs.get('is_const', False) @@ -91,16 +91,11 @@ def __indices_setup__(cls, *args, **kwargs): return tuple(dimensions), tuple(dimensions) - @property - def dimensions(self): - return self._dimensions - @property def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) modifier = '*' * len(self.dimensions) - customtype = CustomDtype(petsc_type, modifier=modifier) - return customtype + return CustomDtype(petsc_type, modifier=modifier) @property def _C_name(self): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 1021662783..94dfe54fbd 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -43,15 +43,21 @@ def test_petsc_functions(): ptr0 = PETScFunction(name='ptr0', dtype=np.float32, dimensions=grid.dimensions, shape=grid.shape) ptr1 = PETScFunction(name='ptr1', dtype=np.float32, grid=grid) - ptr2 = PETScFunction(name='ptr2', dtype=np.float32, grid=grid, is_const=True) + ptr2 = PETScFunction(name='ptr2', dtype=np.float64, grid=grid, is_const=True) + ptr3 = PETScFunction(name='ptr3', dtype=np.int32, grid=grid, is_const=True) + ptr4 = PETScFunction(name='ptr4', dtype=np.int64, grid=grid, 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) == 'PetscScalar**restrict ptr1;' assert str(defn2) == 'const PetscScalar**restrict ptr2;' + assert str(defn3) == 'const PetscInt**restrict ptr3;' + assert str(defn4) == 'const PetscInt**restrict ptr4;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' From 1dc5695714f48ce9a52ffa62f940e6ccc24118d0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 15:02:49 +0000 Subject: [PATCH 013/151] dsl: Inherit from DiscreteFunction for PETScFunction instead?' --- devito/types/petsc.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f07cb298ad..8d7b61935b 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -2,6 +2,7 @@ from devito.tools import dtype_to_petsctype, CustomDtype import numpy as np from devito.types import LocalObject +# from devito.types.dense import DiscreteFunction class DM(LocalObject): @@ -104,3 +105,42 @@ def _C_name(self): @property def is_const(self): return self._is_const + + +# class PETScFunction(DiscreteFunction): +# """ +# PETScFunctions. +# """ +# _data_alignment = False + +# 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): +# grid = kwargs.get('grid') +# dimensions = kwargs.get('dimensions') +# if grid is None: +# if dimensions is None: +# raise TypeError("Need either `grid` or `dimensions`") +# elif dimensions is None: +# dimensions = grid.dimensions + +# return tuple(dimensions), tuple(dimensions) + +# @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 + +# @property +# def is_const(self): +# return self._is_const From 61524457d30b7d35d1259e1589d9bad39cd1ef7c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 10:18:11 +0000 Subject: [PATCH 014/151] compiler: Switch PETScFunction to inherit from Array --- devito/types/petsc.py | 69 ++----------------------------------------- tests/test_petsc.py | 18 ++++++----- 2 files changed, 12 insertions(+), 75 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 8d7b61935b..b5f77c68e0 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,8 +1,5 @@ -from devito.types.basic import AbstractFunction from devito.tools import dtype_to_petsctype, CustomDtype -import numpy as np -from devito.types import LocalObject -# from devito.types.dense import DiscreteFunction +from devito.types import LocalObject, Array class DM(LocalObject): @@ -57,7 +54,7 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScFunction(AbstractFunction): +class PETScFunction(Array): """ PETScFunctions. """ @@ -69,29 +66,6 @@ def __init_finalize__(self, *args, **kwargs): self._is_const = kwargs.get('is_const', False) - @classmethod - def __dtype_setup__(cls, **kwargs): - grid = kwargs.get('grid') - dtype = kwargs.get('dtype') - if dtype is not None: - return dtype - elif grid is not None: - return grid.dtype - else: - return np.float32 - - @classmethod - def __indices_setup__(cls, *args, **kwargs): - grid = kwargs.get('grid') - dimensions = kwargs.get('dimensions') - if grid is None: - if dimensions is None: - raise TypeError("Need either `grid` or `dimensions`") - elif dimensions is None: - dimensions = grid.dimensions - - return tuple(dimensions), tuple(dimensions) - @property def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) @@ -105,42 +79,3 @@ def _C_name(self): @property def is_const(self): return self._is_const - - -# class PETScFunction(DiscreteFunction): -# """ -# PETScFunctions. -# """ -# _data_alignment = False - -# 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): -# grid = kwargs.get('grid') -# dimensions = kwargs.get('dimensions') -# if grid is None: -# if dimensions is None: -# raise TypeError("Need either `grid` or `dimensions`") -# elif dimensions is None: -# dimensions = grid.dimensions - -# return tuple(dimensions), tuple(dimensions) - -# @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 - -# @property -# def is_const(self): -# return self._is_const diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 94dfe54fbd..bfb729619f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -40,12 +40,14 @@ def test_petsc_functions(): grid = Grid((2, 2)) x, y = grid.dimensions - ptr0 = PETScFunction(name='ptr0', dtype=np.float32, dimensions=grid.dimensions, - shape=grid.shape) - ptr1 = PETScFunction(name='ptr1', dtype=np.float32, grid=grid) - ptr2 = PETScFunction(name='ptr2', dtype=np.float64, grid=grid, is_const=True) - ptr3 = PETScFunction(name='ptr3', dtype=np.int32, grid=grid, is_const=True) - ptr4 = PETScFunction(name='ptr4', dtype=np.int64, grid=grid, is_const=True) + ptr0 = PETScFunction(name='ptr0', dimensions=grid.dimensions, dtype=np.float32) + ptr1 = PETScFunction(name='ptr1', dimensions=grid.dimensions, dtype=np.float32, + is_const=True) + ptr2 = PETScFunction(name='ptr2', dimensions=grid.dimensions, dtype=np.float64, + is_const=True) + ptr3 = PETScFunction(name='ptr3', dimensions=grid.dimensions, dtype=np.int32) + ptr4 = PETScFunction(name='ptr4', dimensions=grid.dimensions, dtype=np.int64, + is_const=True) defn0 = Definition(ptr0) defn1 = Definition(ptr1) @@ -56,8 +58,8 @@ def test_petsc_functions(): expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1) assert str(defn0) == 'PetscScalar**restrict ptr0;' - assert str(defn1) == 'PetscScalar**restrict ptr1;' + assert str(defn1) == 'const PetscScalar**restrict ptr1;' assert str(defn2) == 'const PetscScalar**restrict ptr2;' - assert str(defn3) == 'const PetscInt**restrict ptr3;' + assert str(defn3) == 'PetscInt**restrict ptr3;' assert str(defn4) == 'const PetscInt**restrict ptr4;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' From cbedc60eeb8da0e66373d107c3c229192bd1b776 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:03:02 +0000 Subject: [PATCH 015/151] compiler: Edit PETScFunction to inherit from ArrayBasic --- devito/passes/iet/definitions.py | 10 ++++++++++ devito/types/basic.py | 1 + devito/types/petsc.py | 27 +++++++++++++++++++-------- tests/test_petsc.py | 8 ++++---- 4 files changed, 34 insertions(+), 12 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 3532169754..7bb3527c57 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -260,6 +260,14 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage): storage.update(obj, site, allocs=decl) + def _alloc_petsc_array_on_low_lat_mem(self, site, obj, storage): + """ + Allocate a PETScArray in the low latency memory. + """ + decl = Definition(obj) + definition = (decl) + storage.update(obj, site, standalones=definition) + def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage): """ Allocate the following objects in the high bandwidth memory: @@ -412,6 +420,8 @@ def place_definitions(self, iet, globs=None, **kwargs): elif i.is_PointerArray: self._alloc_pointed_array_on_high_bw_mem(iet, i, storage) + elif i.is_PETScArray: + self._alloc_petsc_array_on_low_lat_mem(iet, i, storage) # Handle postponed global objects includes = set() diff --git a/devito/types/basic.py b/devito/types/basic.py index bd04be8564..d127d03282 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -265,6 +265,7 @@ class Basic(CodeSymbol): is_Array = False is_PointerArray = False is_ObjectArray = False + is_PETScArray = False is_Bundle = False is_Object = False is_LocalObject = False diff --git a/devito/types/petsc.py b/devito/types/petsc.py index b5f77c68e0..30ca184e0f 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,5 +1,9 @@ -from devito.tools import dtype_to_petsctype, CustomDtype -from devito.types import LocalObject, Array +from devito.tools import (dtype_to_petsctype, CustomDtype, + dtype_to_ctype) +from devito.types import LocalObject +from devito.types.array import ArrayBasic +from ctypes import POINTER +import numpy as np class DM(LocalObject): @@ -54,23 +58,30 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScFunction(Array): - """ - PETScFunctions. - """ +class PETScFunction(ArrayBasic): + _data_alignment = False + is_PETScArray = True + def __init_finalize__(self, *args, **kwargs): super().__init_finalize__(*args, **kwargs) self._is_const = kwargs.get('is_const', False) + @classmethod + def __dtype_setup__(cls, **kwargs): + return kwargs.get('dtype', np.float32) + @property def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) - modifier = '*' * len(self.dimensions) - return CustomDtype(petsc_type, modifier=modifier) + ctype = dtype_to_ctype(self.dtype) + r = type(petsc_type, (ctype,), {}) + for n in range(len(self.dimensions)): + r = POINTER(r) + return r @property def _C_name(self): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index bfb729619f..8a7a645851 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -57,9 +57,9 @@ def test_petsc_functions(): 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(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(defn4) == 'const PetscInt **restrict ptr4;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' From 469e3080e340990f3e838df6ec5bf272cafc30c5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:08:40 +0000 Subject: [PATCH 016/151] misc: Fix pytest for test_petsc.py --- tests/test_petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8a7a645851..5dc60e2aff 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -60,6 +60,6 @@ def test_petsc_functions(): 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(defn3) == 'PetscInt **restrict ptr3;' assert str(defn4) == 'const PetscInt **restrict ptr4;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' From 912922333a5ef88e1c3e8bc891e8f542b00ccf09 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:20:16 +0000 Subject: [PATCH 017/151] compiler: Edit _C_ctype property of PETScFunction to utilise CustomDtype --- devito/types/petsc.py | 12 ++++-------- tests/test_petsc.py | 10 +++++----- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 30ca184e0f..45daf4243f 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,8 +1,6 @@ -from devito.tools import (dtype_to_petsctype, CustomDtype, - dtype_to_ctype) +from devito.tools import dtype_to_petsctype, CustomDtype from devito.types import LocalObject from devito.types.array import ArrayBasic -from ctypes import POINTER import numpy as np @@ -77,11 +75,9 @@ def __dtype_setup__(cls, **kwargs): @property def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) - ctype = dtype_to_ctype(self.dtype) - r = type(petsc_type, (ctype,), {}) - for n in range(len(self.dimensions)): - r = POINTER(r) - return r + modifier = '*' * len(self.dimensions) + customtype = CustomDtype(petsc_type, modifier=modifier) + return customtype @property def _C_name(self): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 5dc60e2aff..bfb729619f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -57,9 +57,9 @@ def test_petsc_functions(): 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(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;' From acff6877ca26ec58b32c5104c09586550e4d5341 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:22:33 +0000 Subject: [PATCH 018/151] compiler: Edit _C_ctype property of PETScFunction --- devito/types/petsc.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 45daf4243f..b8ffae94dd 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -76,8 +76,7 @@ def __dtype_setup__(cls, **kwargs): def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) modifier = '*' * len(self.dimensions) - customtype = CustomDtype(petsc_type, modifier=modifier) - return customtype + return CustomDtype(petsc_type, modifier=modifier) @property def _C_name(self): From 3362ad8ef1959986eae4b166fdc338a992e91eec Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 019/151] misc: Update file due to incorrect rebase --- docker/Dockerfile.intel | 50 ----------------------------------------- 1 file changed, 50 deletions(-) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index 1fab7147d6..e013fccc03 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,7 +35,6 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu -<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -51,16 +50,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev -======= -# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal -RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg -RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list - -RUN apt-get update -y && apt-get dist-upgrade -y && \ - apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ - intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ - libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -87,11 +76,7 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD # ICX OpenMP image -======= -# ICX image ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -108,35 +93,6 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD -======= -# ICX hpc image -############################################################## -FROM oneapi as icx-hpc - -# Install both icc and icx to avoid missing dependencies -RUN apt-get update -y && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic - -# Missig components -# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager -RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ - chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ - rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -# Devito config -ENV DEVITO_ARCH="icx" -ENV DEVITO_LANGUAGE="openmp" -# MPICC compiler for mpi4py -ENV MPICC=mpiicc -ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' - -############################################################## ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -151,12 +107,6 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl -<<<<<<< HEAD -======= -# NOTE: the name of this file ends with ".cpu" but this is a GPU image. -# It then feels a bit akward, so some restructuring might be needed - ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 5b15b4e15c5f6702ecf7ee4cef77889d7178f2d3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Jan 2024 11:46:05 +0000 Subject: [PATCH 020/151] compiler: Change name to PETScArray from PETScFunction --- devito/passes/iet/definitions.py | 10 ---------- devito/types/petsc.py | 2 +- tests/test_petsc.py | 20 ++++++++++---------- 3 files changed, 11 insertions(+), 21 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 7bb3527c57..3532169754 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -260,14 +260,6 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage): storage.update(obj, site, allocs=decl) - def _alloc_petsc_array_on_low_lat_mem(self, site, obj, storage): - """ - Allocate a PETScArray in the low latency memory. - """ - decl = Definition(obj) - definition = (decl) - storage.update(obj, site, standalones=definition) - def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage): """ Allocate the following objects in the high bandwidth memory: @@ -420,8 +412,6 @@ def place_definitions(self, iet, globs=None, **kwargs): elif i.is_PointerArray: self._alloc_pointed_array_on_high_bw_mem(iet, i, storage) - elif i.is_PETScArray: - self._alloc_petsc_array_on_low_lat_mem(iet, i, storage) # Handle postponed global objects includes = set() diff --git a/devito/types/petsc.py b/devito/types/petsc.py index b8ffae94dd..117f0d8bd7 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -56,7 +56,7 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScFunction(ArrayBasic): +class PETScArray(ArrayBasic): _data_alignment = False diff --git a/tests/test_petsc.py b/tests/test_petsc.py index bfb729619f..fa05a32a3d 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -2,7 +2,7 @@ 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, PETScFunction) + PC, KSPConvergedReason, PETScArray) import numpy as np @@ -35,19 +35,19 @@ def test_petsc_local_object(): def test_petsc_functions(): """ - Test C++ support for PETScFunctions. + Test C++ support for PETScArrays. """ grid = Grid((2, 2)) x, y = grid.dimensions - ptr0 = PETScFunction(name='ptr0', dimensions=grid.dimensions, dtype=np.float32) - ptr1 = PETScFunction(name='ptr1', dimensions=grid.dimensions, dtype=np.float32, - is_const=True) - ptr2 = PETScFunction(name='ptr2', dimensions=grid.dimensions, dtype=np.float64, - is_const=True) - ptr3 = PETScFunction(name='ptr3', dimensions=grid.dimensions, dtype=np.int32) - ptr4 = PETScFunction(name='ptr4', dimensions=grid.dimensions, dtype=np.int64, - is_const=True) + 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) From 7f3920a03a20b6e23a1428c0894973fc3014bf9d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 11:12:00 +0000 Subject: [PATCH 021/151] compiler: Add intermediate class between PETScArray and ArrayBasic --- devito/types/basic.py | 1 - devito/types/petsc.py | 28 ++++++++++++++++++++-------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/devito/types/basic.py b/devito/types/basic.py index d127d03282..bd04be8564 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -265,7 +265,6 @@ class Basic(CodeSymbol): is_Array = False is_PointerArray = False is_ObjectArray = False - is_PETScArray = False is_Bundle = False is_Object = False is_LocalObject = False diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 117f0d8bd7..66c0b12abc 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -56,11 +56,20 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScArray(ArrayBasic): +class AbstractArray(ArrayBasic): + + """ + A customised version of ArrayBasic that allows objects to + be generated without a cast. This is particularly useful for + PETSc objects. Also, AbstractArray objects can + be explicitly specifed as constant. + """ _data_alignment = False - is_PETScArray = True + is_AbstractArray = True + + __rkwargs__ = ArrayBasic.__rkwargs__ + ('is_const',) def __init_finalize__(self, *args, **kwargs): @@ -72,12 +81,6 @@ def __init_finalize__(self, *args, **kwargs): def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) - @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 @@ -85,3 +88,12 @@ def _C_name(self): @property def is_const(self): return self._is_const + + +class PETScArray(AbstractArray): + + @property + def _C_ctype(self): + petsc_type = dtype_to_petsctype(self.dtype) + modifier = '*' * len(self.dimensions) + return CustomDtype(petsc_type, modifier=modifier) From 15e74001315bd0916a05a58538e0d2f4ad87c6f6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 13:33:58 +0000 Subject: [PATCH 022/151] compiler: Remove PETSc specific checks in place_definitions --- devito/passes/iet/definitions.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 3532169754..490b3a4ea8 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -407,11 +407,10 @@ def place_definitions(self, iet, globs=None, **kwargs): # 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() From ea9be50f2ab51228d70de1c9e691e3efc9195d73 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 14:08:14 +0000 Subject: [PATCH 023/151] compiler: Remove PETSc specific check from place_definitions and edit PETScArray --- devito/passes/iet/definitions.py | 13 +++++++++++- devito/types/petsc.py | 35 ++++++++------------------------ 2 files changed, 21 insertions(+), 27 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 490b3a4ea8..791d16c83b 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -260,6 +260,17 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage): storage.update(obj, site, allocs=decl) + def _alloc_special_type_on_low_lat_mem(self, site, obj, storage): + """ + Allocate special types in the low latency memory e.g PETScArrays. + These types are typically customised versions of ArrayBasic + that deviate from the standard types + generated by the compiler. + """ + decl = Definition(obj) + + storage.update(obj, site, allocs=decl) + def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage): """ Allocate the following objects in the high bandwidth memory: @@ -410,7 +421,7 @@ def place_definitions(self, iet, globs=None, **kwargs): 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) + self._alloc_special_type_on_low_lat_mem(iet, i, storage) # Handle postponed global objects includes = set() diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 66c0b12abc..bef6b486cb 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -56,44 +56,27 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class AbstractArray(ArrayBasic): - +class PETScArray(ArrayBasic): """ - A customised version of ArrayBasic that allows objects to - be generated without a cast. This is particularly useful for - PETSc objects. Also, AbstractArray objects can - be explicitly specifed as constant. + 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 - is_AbstractArray = True - - __rkwargs__ = ArrayBasic.__rkwargs__ + ('is_const',) - - def __init_finalize__(self, *args, **kwargs): - - super().__init_finalize__(*args, **kwargs) - - self._is_const = kwargs.get('is_const', False) + is_PETScArray = True @classmethod def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) - @property - def _C_name(self): - return self.name - - @property - def is_const(self): - return self._is_const - - -class PETScArray(AbstractArray): - @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 From 905a7e2d0345bba57ecf6f4867517ea74ecca562 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 16:14:31 +0000 Subject: [PATCH 024/151] compiler: Remove is_PETScArray --- devito/types/petsc.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index bef6b486cb..e2c6cfcd08 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -65,8 +65,6 @@ class PETScArray(ArrayBasic): _data_alignment = False - is_PETScArray = True - @classmethod def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) From 415ad20853d646a4bcc75a09633135c67f200f69 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 18:53:04 +0000 Subject: [PATCH 025/151] compiler: Remove elif is_ObjectArray and use alloc_object_.. inside the else statement --- devito/passes/iet/definitions.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 791d16c83b..490b3a4ea8 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -260,17 +260,6 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage): storage.update(obj, site, allocs=decl) - def _alloc_special_type_on_low_lat_mem(self, site, obj, storage): - """ - Allocate special types in the low latency memory e.g PETScArrays. - These types are typically customised versions of ArrayBasic - that deviate from the standard types - generated by the compiler. - """ - decl = Definition(obj) - - storage.update(obj, site, allocs=decl) - def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage): """ Allocate the following objects in the high bandwidth memory: @@ -421,7 +410,7 @@ def place_definitions(self, iet, globs=None, **kwargs): elif i.is_PointerArray: self._alloc_pointed_array_on_high_bw_mem(iet, i, storage) else: - self._alloc_special_type_on_low_lat_mem(iet, i, storage) + self._alloc_object_array_on_low_lat_mem(iet, i, storage) # Handle postponed global objects includes = set() From adc54282524e8de69bda86773241bdf3ec9edd35 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 8 Jan 2024 10:04:22 +0000 Subject: [PATCH 026/151] misc: Delete empty file --- devito/passes/iet/petsc.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 devito/passes/iet/petsc.py diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py deleted file mode 100644 index e69de29bb2..0000000000 From 0a571c7b53675ec4f14c0aa0fc459f6665d2599d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 8 Jan 2024 10:12:19 +0000 Subject: [PATCH 027/151] misc: Move dtype_to_petsctype to petsc.py and clean up blank lines --- devito/tools/dtypes_lowering.py | 13 +------------ devito/types/petsc.py | 16 ++++++++++++++-- 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 6f54b2d0a0..e0167289bb 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -13,7 +13,7 @@ 'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype', 'dtype_to_cstr', 'dtype_to_ctype', 'dtype_to_mpitype', 'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', - 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'dtype_to_petsctype'] + 'is_external_ctype', 'infer_dtype', 'CustomDtype'] # *** Custom np.dtypes @@ -312,14 +312,3 @@ def infer_dtype(dtypes): else: # E.g., mixed integer arithmetic return max(dtypes, key=lambda i: np.dtype(i).itemsize, default=None) - - -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] diff --git a/devito/types/petsc.py b/devito/types/petsc.py index e2c6cfcd08..d6d07b6e2c 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,7 +1,8 @@ -from devito.tools import dtype_to_petsctype, CustomDtype +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): @@ -69,7 +70,7 @@ class PETScArray(ArrayBasic): def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) - @property + @cached_property def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) modifier = '*' * len(self.dimensions) @@ -78,3 +79,14 @@ def _C_ctype(self): @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] From ce3cfb3cd1426ea677e13d83ec2b8bc85b02451f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 8 Jan 2024 10:14:28 +0000 Subject: [PATCH 028/151] misc: Clean up blank line --- devito/tools/dtypes_lowering.py | 1 + 1 file changed, 1 insertion(+) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index e0167289bb..b5b564a4d7 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -15,6 +15,7 @@ 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', 'is_external_ctype', 'infer_dtype', 'CustomDtype'] + # *** Custom np.dtypes # NOTE: the following is inspired by pyopencl.cltypes From 245242a4c0baca60a1d351b435a08f1e8dad52b2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Fri, 23 Feb 2024 15:03:19 +0000 Subject: [PATCH 029/151] Update nodes.py --- devito/ir/iet/nodes.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index c0f2a32c4f..8487d89a34 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -28,11 +28,10 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'FuncPtrCall'] + 'CallableBody', 'Transfer', 'Callback'] # First-class IET nodes - class Node(Signer): __metaclass__ = abc.ABCMeta From b45b17c1e9073cb205a63e7bdd71a7d4e17f9b53 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Fri, 23 Feb 2024 15:07:11 +0000 Subject: [PATCH 030/151] Update nodes.py --- devito/ir/iet/nodes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 8487d89a34..728e88afbe 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -32,6 +32,7 @@ # First-class IET nodes + class Node(Signer): __metaclass__ = abc.ABCMeta From 40699ae9bd2206c31f38c5d8e5c2addb1af3722a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 10:46:24 +0000 Subject: [PATCH 031/151] compiler: Add FuncPtrCall --- devito/ir/iet/nodes.py | 13 +++++++++++++ devito/ir/iet/visitors.py | 6 ++++++ tests/test_iet.py | 12 ++++++++++++ 3 files changed, 31 insertions(+) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 728e88afbe..459a0b8bdf 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -358,6 +358,19 @@ def writes(self): return self._writes +class FuncPtrCall(Call): + + """ + Function Pointer Argument to a Call. + """ + def __init__(self, name, return_type, parameter_type, **kwargs): + + super().__init__(name=name) + + self.return_type = return_type + self.parameter_type = parameter_type + + class Expression(ExprStmt, Node): """ diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index d2f5d9db3b..715d1c0d6c 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -520,6 +520,12 @@ def visit_Call(self, o, nested_call=False): else: return c.Initializer(c.Value(rettype, retobj._C_name), call) + def visit_FuncPtrCall(self, o, nested_call=False): + name = o.name + return_type = o.return_type + parameter_type = o.parameter_type + return FuncPtrArg(name, return_type, parameter_type) + def visit_Conditional(self, o): try: then_body, else_body = self._blankline_logic(o.children) diff --git a/tests/test_iet.py b/tests/test_iet.py index 1c36830084..81e40b9c7f 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -129,16 +129,28 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(expected) +<<<<<<< HEAD def test_callback_cgen(): +======= +def test_funcptrcall_cgen(): +>>>>>>> 8677a933f (compiler: Add FuncPtrCall) a = Symbol('a') b = Symbol('b') foo0 = Callable('foo0', Definition(a), 'void', parameters=[b]) +<<<<<<< HEAD foo0_arg = Callback(foo0.name, foo0.retval, 'int') code0 = CGen().visit(foo0_arg) assert str(code0) == '(void (*)(int))foo0' # Test nested calls with a Callback as an argument. +======= + foo0_arg = FuncPtrCall(foo0.name, foo0.retval, 'int') + code0 = CGen().visit(foo0_arg) + assert str(code0) == '(void (*)(int))foo0' + + # test nested calls with a FuncPtrCall as an argument +>>>>>>> 8677a933f (compiler: Add FuncPtrCall) call = Call('foo1', [ Call('foo2', [foo0_arg]) ]) From 0ad362583be1c891ae8599a326eac9fad6d14999 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 11:04:27 +0000 Subject: [PATCH 032/151] compiler: Edit names, move position of Callback and other small edits --- devito/ir/iet/nodes.py | 13 ------------- devito/ir/iet/visitors.py | 6 ------ tests/test_iet.py | 12 ++++++++++++ 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 459a0b8bdf..728e88afbe 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -358,19 +358,6 @@ def writes(self): return self._writes -class FuncPtrCall(Call): - - """ - Function Pointer Argument to a Call. - """ - def __init__(self, name, return_type, parameter_type, **kwargs): - - super().__init__(name=name) - - self.return_type = return_type - self.parameter_type = parameter_type - - class Expression(ExprStmt, Node): """ diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 715d1c0d6c..d2f5d9db3b 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -520,12 +520,6 @@ def visit_Call(self, o, nested_call=False): else: return c.Initializer(c.Value(rettype, retobj._C_name), call) - def visit_FuncPtrCall(self, o, nested_call=False): - name = o.name - return_type = o.return_type - parameter_type = o.parameter_type - return FuncPtrArg(name, return_type, parameter_type) - def visit_Conditional(self, o): try: then_body, else_body = self._blankline_logic(o.children) diff --git a/tests/test_iet.py b/tests/test_iet.py index 81e40b9c7f..473a3b231f 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -129,15 +129,20 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(expected) +<<<<<<< HEAD <<<<<<< HEAD def test_callback_cgen(): ======= def test_funcptrcall_cgen(): >>>>>>> 8677a933f (compiler: Add FuncPtrCall) +======= +def test_callback_cgen(): +>>>>>>> f15e61d35 (compiler: Edit names, move position of Callback and other small edits) a = Symbol('a') b = Symbol('b') foo0 = Callable('foo0', Definition(a), 'void', parameters=[b]) +<<<<<<< HEAD <<<<<<< HEAD foo0_arg = Callback(foo0.name, foo0.retval, 'int') code0 = CGen().visit(foo0_arg) @@ -151,6 +156,13 @@ def test_funcptrcall_cgen(): # test nested calls with a FuncPtrCall as an argument >>>>>>> 8677a933f (compiler: Add FuncPtrCall) +======= + foo0_arg = Callback(foo0.name, foo0.retval, 'int') + code0 = CGen().visit(foo0_arg) + assert str(code0) == '(void (*)(int))foo0' + + # test nested calls with a Callback as an argument. +>>>>>>> f15e61d35 (compiler: Edit names, move position of Callback and other small edits) call = Call('foo1', [ Call('foo2', [foo0_arg]) ]) From 136921a401d964903aa43c56361ab9a10f6302c6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 28 Feb 2024 13:03:28 +0000 Subject: [PATCH 033/151] types: Update PETScArray to have liveness --- devito/types/petsc.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index d6d07b6e2c..46b0ceec49 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -66,6 +66,15 @@ class PETScArray(ArrayBasic): _data_alignment = False + __rkwargs__ = (ArrayBasic.__rkwargs__ + + ('liveness',)) + + def __init_finalize__(self, *args, **kwargs): + super().__init_finalize__(*args, **kwargs) + + self._liveness = kwargs.get('liveness', 'lazy') + assert self._liveness in ['eager', 'lazy'] + @classmethod def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) @@ -80,6 +89,18 @@ def _C_ctype(self): def _C_name(self): return self.name + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' + def dtype_to_petsctype(dtype): """Map numpy types to PETSc datatypes.""" From c9e9836f65da995a07925411d94039f2d0959dc9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 1 Mar 2024 15:45:07 +0000 Subject: [PATCH 034/151] compiler: Switch get to pop in PETScArray --- devito/types/petsc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 46b0ceec49..55af0fd1ef 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -70,11 +70,11 @@ class PETScArray(ArrayBasic): ('liveness',)) def __init_finalize__(self, *args, **kwargs): - super().__init_finalize__(*args, **kwargs) - - self._liveness = kwargs.get('liveness', 'lazy') + self._liveness = kwargs.pop('liveness', 'lazy') assert self._liveness in ['eager', 'lazy'] + super().__init_finalize__(*args, **kwargs) + @classmethod def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) From 31c8a27702185fafa8aa7437c9352788abfcb82a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 11 Mar 2024 17:25:23 +0000 Subject: [PATCH 035/151] types: Move liveness to ArrayBasic from PETScArray --- devito/types/array.py | 14 ++++++++++++++ devito/types/petsc.py | 21 +-------------------- 2 files changed, 15 insertions(+), 20 deletions(-) diff --git a/devito/types/array.py b/devito/types/array.py index 44d7fabf9d..6f55fa87ed 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -20,6 +20,8 @@ class ArrayBasic(AbstractFunction, LocalType): __rkwargs__ = AbstractFunction.__rkwargs__ + ('is_const', 'liveness') def __init_finalize__(self, *args, **kwargs): + + self._liveness = kwargs.setdefault('liveness', 'lazy') super().__init_finalize__(*args, **kwargs) self._liveness = kwargs.get('liveness', 'lazy') @@ -59,6 +61,18 @@ def shape_allocated(self): @property def is_const(self): return self._is_const + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' class Array(ArrayBasic): diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 55af0fd1ef..dea8a77697 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -66,14 +66,7 @@ class PETScArray(ArrayBasic): _data_alignment = False - __rkwargs__ = (ArrayBasic.__rkwargs__ + - ('liveness',)) - - def __init_finalize__(self, *args, **kwargs): - self._liveness = kwargs.pop('liveness', 'lazy') - assert self._liveness in ['eager', 'lazy'] - - super().__init_finalize__(*args, **kwargs) + __rkwargs__ = ArrayBasic.__rkwargs__ @classmethod def __dtype_setup__(cls, **kwargs): @@ -89,18 +82,6 @@ def _C_ctype(self): def _C_name(self): return self.name - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' - def dtype_to_petsctype(dtype): """Map numpy types to PETSc datatypes.""" From 3d5278e99b7dfd9804f4807f3b2072591cb4ea8f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 11 Mar 2024 17:26:57 +0000 Subject: [PATCH 036/151] misc: clean --- devito/types/array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/array.py b/devito/types/array.py index 6f55fa87ed..f478add0e3 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -61,7 +61,7 @@ def shape_allocated(self): @property def is_const(self): return self._is_const - + @property def liveness(self): return self._liveness From ae01826ad84495aaa5da8fc60a3c426dcb9bfe00 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 11 Mar 2024 19:54:12 +0000 Subject: [PATCH 037/151] types: Make PETScArrays differentiable --- devito/types/__init__.py | 2 +- devito/types/petsc.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/devito/types/__init__.py b/devito/types/__init__.py index ff024f24be..9bc151569a 100644 --- a/devito/types/__init__.py +++ b/devito/types/__init__.py @@ -6,13 +6,13 @@ 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 from .caching import _SymbolCache, CacheManager # noqa from .equation import * # noqa from .constant import * # noqa +from .petsc import * # noqa # Some more internal types which depend on some of the types above from .parallel import * # noqa diff --git a/devito/types/petsc.py b/devito/types/petsc.py index dea8a77697..247564bfee 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -3,6 +3,8 @@ from devito.types.array import ArrayBasic import numpy as np from cached_property import cached_property +from devito.finite_differences import Differentiable +from devito.types.basic import AbstractFunction class DM(LocalObject): @@ -57,7 +59,7 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScArray(ArrayBasic): +class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent a customised variant of ArrayBasic. They are designed to @@ -66,7 +68,8 @@ class PETScArray(ArrayBasic): _data_alignment = False - __rkwargs__ = ArrayBasic.__rkwargs__ + __rkwargs__ = (AbstractFunction.__rkwargs__ + + ('dimensions', 'liveness')) @classmethod def __dtype_setup__(cls, **kwargs): From 222d8040c0845aae11550ed245e90fddbe3025d5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 12 Mar 2024 11:59:50 +0000 Subject: [PATCH 038/151] types: Introduce PETScFunction class to enable differentiablity and therefore compatability with Function substitutions --- devito/types/petsc.py | 11 ++++++++++- tests/test_petsc.py | 23 ++++++++++++++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 247564bfee..f70d092a00 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -59,7 +59,16 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScArray(ArrayBasic, Differentiable): +class PETScFunction(Differentiable): + """ + Base class for all PETSc objects. + Differentiable enables compatability with standard Function objects, + allowing for the use of the `subs` method. + """ + pass + + +class PETScArray(ArrayBasic, PETScFunction): """ PETScArrays are generated by the compiler only and represent a customised variant of ArrayBasic. They are designed to diff --git a/tests/test_petsc.py b/tests/test_petsc.py index fa05a32a3d..ce3873b70f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,4 +1,4 @@ -from devito import Grid +from devito import Grid, Function, Eq 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, @@ -63,3 +63,24 @@ def test_petsc_functions(): assert str(defn3) == 'PetscInt**restrict ptr3;' assert str(defn4) == 'const PetscInt**restrict ptr4;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' + + +def test_petsc_subs(): + """ + Test support for PETScArrays in substitutions. + """ + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=2) + f2 = Function(name='f2', grid=grid, space_order=2) + + arr = PETScArray(name='arr', dimensions=f2.dimensions, dtype=f2.dtype) + + eqn = Eq(f1, f2.laplace) + eqn_subs = eqn.subs(f2, arr) + + assert str(eqn) == 'Eq(f1(x, y), Derivative(f2(x, y), (x, 2))' + \ + ' + Derivative(f2(x, y), (y, 2)))' + + assert str(eqn_subs) == 'Eq(f1(x, y), Derivative(arr(x, y), (x, 2))' + \ + ' + Derivative(arr(x, y), (y, 2)))' From bbf64d257f5fc14615f8aa706cf533b46727dc41 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 14 Mar 2024 17:25:22 +0000 Subject: [PATCH 039/151] types: Update PETScArray --- devito/types/petsc.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f70d092a00..c9307d2a09 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -59,20 +59,15 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class PETScFunction(Differentiable): - """ - Base class for all PETSc objects. - Differentiable enables compatability with standard Function objects, - allowing for the use of the `subs` method. - """ - pass - - -class PETScArray(ArrayBasic, PETScFunction): +class PETScArray(ArrayBasic, Differentiable): """ 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. + Differentiable enables compatability with standard Function objects, + allowing for the use of the `subs` method. + TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) + and then PETScArray(ArrayBasic). """ _data_alignment = False From 9ce5462081b77d40ab8dfb1a0add0c8414938b00 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 14 Mar 2024 17:51:22 +0000 Subject: [PATCH 040/151] types: Edit init file for petsc types --- devito/types/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/__init__.py b/devito/types/__init__.py index 9bc151569a..3a5e211e63 100644 --- a/devito/types/__init__.py +++ b/devito/types/__init__.py @@ -12,10 +12,10 @@ from .caching import _SymbolCache, CacheManager # noqa from .equation import * # noqa from .constant import * # noqa -from .petsc import * # noqa # Some more internal types which depend on some of the types above from .parallel import * # noqa +from .petsc import * # noqa # Needed only outside Devito from .grid import * # noqa From 5e70798db35f3d90369213cc085d94f3a6211401 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 15 Mar 2024 11:36:07 +0000 Subject: [PATCH 041/151] revert file messed up rebase --- tests/test_iet.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/tests/test_iet.py b/tests/test_iet.py index 473a3b231f..1c36830084 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -129,40 +129,16 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(expected) -<<<<<<< HEAD -<<<<<<< HEAD def test_callback_cgen(): -======= -def test_funcptrcall_cgen(): ->>>>>>> 8677a933f (compiler: Add FuncPtrCall) -======= -def test_callback_cgen(): ->>>>>>> f15e61d35 (compiler: Edit names, move position of Callback and other small edits) a = Symbol('a') b = Symbol('b') foo0 = Callable('foo0', Definition(a), 'void', parameters=[b]) -<<<<<<< HEAD -<<<<<<< HEAD foo0_arg = Callback(foo0.name, foo0.retval, 'int') code0 = CGen().visit(foo0_arg) assert str(code0) == '(void (*)(int))foo0' # Test nested calls with a Callback as an argument. -======= - foo0_arg = FuncPtrCall(foo0.name, foo0.retval, 'int') - code0 = CGen().visit(foo0_arg) - assert str(code0) == '(void (*)(int))foo0' - - # test nested calls with a FuncPtrCall as an argument ->>>>>>> 8677a933f (compiler: Add FuncPtrCall) -======= - foo0_arg = Callback(foo0.name, foo0.retval, 'int') - code0 = CGen().visit(foo0_arg) - assert str(code0) == '(void (*)(int))foo0' - - # test nested calls with a Callback as an argument. ->>>>>>> f15e61d35 (compiler: Edit names, move position of Callback and other small edits) call = Call('foo1', [ Call('foo2', [foo0_arg]) ]) From 649b28235dcab6258f6c82914cbbed18591cd0fa Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 15 Mar 2024 15:44:21 +0000 Subject: [PATCH 042/151] types: Add coefficients arg to PETScArray --- devito/types/petsc.py | 21 ++++++++++++++++++++- tests/test_petsc.py | 4 ++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index c9307d2a09..f01eba7b65 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -5,6 +5,7 @@ from cached_property import cached_property from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction +from devito.finite_differences.tools import fd_weights_registry class DM(LocalObject): @@ -72,13 +73,31 @@ class PETScArray(ArrayBasic, Differentiable): _data_alignment = False + # Default method for the finite difference approximation weights computation. + _default_fd = 'taylor' + __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'liveness')) + ('dimensions', 'liveness', 'coefficients')) + + def __init_finalize__(self, *args, **kwargs): + + super().__init_finalize__(*args, **kwargs) + + # Symbolic (finite difference) coefficients + self._coefficients = kwargs.get('coefficients', self._default_fd) + if self._coefficients not in fd_weights_registry: + raise ValueError("coefficients must be one of %s" + " not %s" % (str(fd_weights_registry), self._coefficients)) @classmethod def __dtype_setup__(cls, **kwargs): return kwargs.get('dtype', np.float32) + @property + def coefficients(self): + """Form of the coefficients of the function.""" + return self._coefficients + @cached_property def _C_ctype(self): petsc_type = dtype_to_petsctype(self.dtype) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index ce3873b70f..955f361f7f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -84,3 +84,7 @@ def test_petsc_subs(): assert str(eqn_subs) == 'Eq(f1(x, y), Derivative(arr(x, y), (x, 2))' + \ ' + Derivative(arr(x, y), (y, 2)))' + + assert str(eqn_subs.rhs.evaluate) == '-2.0*arr(x, y)/h_x**2' + \ + ' + arr(x - h_x, y)/h_x**2 + arr(x + h_x, y)/h_x**2 - 2.0*arr(x, y)/h_y**2' + \ + ' + arr(x, y - h_y)/h_y**2 + arr(x, y + h_y)/h_y**2' From 2b1721603e587745c07779595b22fe26221c690a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 15 Mar 2024 17:55:47 +0000 Subject: [PATCH 043/151] types: Add basic PETScEq --- devito/types/petsc.py | 64 +++++++++++++++++++++++++++++++++++++++++++ tests/test_petsc.py | 13 +++++++++ 2 files changed, 77 insertions(+) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f01eba7b65..9d0fbc61cd 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -6,6 +6,7 @@ from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry +from devito import Eq class DM(LocalObject): @@ -118,3 +119,66 @@ def dtype_to_petsctype(dtype): np.int64: 'PetscInt', np.float64: 'PetscScalar' }[dtype] + + +class PETScEq(Eq): + """ + Represents a general equation required by a PETSc solve. + """ + + __rkwargs__ = (Eq.__rkwargs__ + ('target', 'solver_parameters',)) + + defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'jacobi', + 'ksp_rtol': 'PETSC_DEFAULT', + 'ksp_atol': 'PETSC_DEFAULT', + 'ksp_divtol': 'PETSC_DEFAULT', + 'ksp_max_it': 'PETSC_DEFAULT' + } + + def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None, implicit_dims=None, + target=None, solver_parameters=None, **kwargs): + + if solver_parameters is None: + solver_parameters = cls.defaults + else: + for key, val in cls.defaults.items(): + solver_parameters[key] = solver_parameters.get(key, val) + + obj = Eq.__new__(cls, lhs, rhs, subdomain=subdomain, coefficients=coefficients, + implicit_dims=implicit_dims, **kwargs) + obj._target = target + obj._solver_parameters = solver_parameters + + return obj + + @property + def target(self): + return self._target + + @property + def solver_parameters(self): + return self._solver_parameters + + +class Action(PETScEq): + """ + Represents the mathematical expression of applying a linear + operator to a vector. This is a key component + for running matrix-free solvers. + """ + pass + + +def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): + + y_matvec = PETScArray(name='y_matvec_'+str(target.name), dtype=target.dtype, + dimensions=target.dimensions, + shape=target.shape, liveness='eager') + + action = Action(y_matvec, eq.lhs, subdomain=eq.subdomain, + implicit_dims=(target.grid.time_dim,), target=target, + solver_parameters=solver_parameters) + + return [action] \ No newline at end of file diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 955f361f7f..9a672a6643 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -88,3 +88,16 @@ def test_petsc_subs(): assert str(eqn_subs.rhs.evaluate) == '-2.0*arr(x, y)/h_x**2' + \ ' + arr(x - h_x, y)/h_x**2 + arr(x + h_x, y)/h_x**2 - 2.0*arr(x, y)/h_y**2' + \ ' + arr(x, y - h_y)/h_y**2 + arr(x, y + h_y)/h_y**2' + + +# def test_petsc_solve(): + +# grid = Grid((2, 2)) + +# f1 = Function(name='f1', grid=grid, space_order=2) +# f2 = Function(name='f2', grid=grid, space_order=2) + +# arr = PETScArray(name='arr', dimensions=f2.dimensions, dtype=f2.dtype) + +# eqn = Eq(f1, f2.laplace) +# eqn_subs = eqn.subs(f2, arr) From 2c01d4fff609c8c9dfb77dda12a44febcf509b55 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 14:49:37 +0000 Subject: [PATCH 044/151] types: Add PETScSolve and simple iet pass --- devito/ir/clusters/cluster.py | 11 +++++---- devito/ir/equations/equation.py | 27 +++++++++++++++++++--- devito/ir/iet/algorithms.py | 9 +++++++- devito/ir/iet/nodes.py | 22 ++++++++++++++++-- devito/operator/operator.py | 4 +++- devito/passes/iet/__init__.py | 1 + devito/types/petsc.py | 41 ++++++++++++++++++++++++++------- 7 files changed, 95 insertions(+), 20 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index f0b0d718ea..cf5957b907 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -375,11 +375,12 @@ def dspace(self): else: d = i.dim try: - if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: - # It'd mean trying to access a point before the - # left halo (test0) or after the right halo (test1) - oobs.update(d._defines) + if not isinstance(f, PETScArray): + if i.lower < 0 or \ + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + # It'd mean trying to access a point before the + # left halo (test0) or after the right halo (test1) + oobs.update(d._defines) except (KeyError, TypeError): # Unable to detect presence of OOB accesses (e.g., `d` not in # `f._size_halo`, that is typical of indirect accesses `A[B[i]]`) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index ada1c23f22..c6529ea8ec 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -10,6 +10,7 @@ from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min +from devito.types.petsc import Action, RHS __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', 'identity_mapper'] @@ -18,7 +19,8 @@ class IREq(sympy.Eq, Pickable): __rargs__ = ('lhs', 'rhs') - __rkwargs__ = ('ispace', 'conditionals', 'implicit_dims', 'operation') + __rkwargs__ = ('ispace', 'conditionals', 'implicit_dims', 'operation', + 'target', 'solver_parameters') @property def is_Scalar(self): @@ -58,6 +60,14 @@ def state(self): def operation(self): return self._operation + @property + def target(self): + return self._target + + @property + def solver_parameters(self): + return self._solver_parameters + @property def is_Reduction(self): return self.operation in (OpInc, OpMin, OpMax) @@ -102,7 +112,9 @@ def detect(cls, expr): reduction_mapper = { Inc: OpInc, ReduceMax: OpMax, - ReduceMin: OpMin + ReduceMin: OpMin, + Action: OpAction, + RHS: OpRHS, } try: return reduction_mapper[type(expr)] @@ -119,6 +131,8 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') +OpAction = Operation('action') +OpRHS = Operation('rhs') identity_mapper = { @@ -235,7 +249,10 @@ def __new__(cls, *args, **kwargs): expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + expr._target = input_expr.target if hasattr(input_expr, 'target') else None + expr._solver_parameters = input_expr.solver_parameters \ + if hasattr(input_expr, 'solver_parameters') else None + return expr @property @@ -291,6 +308,10 @@ def __new__(cls, *args, **kwargs): expr._conditionals = kwargs.get('conditionals', frozendict()) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) + expr._target = input_expr.target \ + if hasattr(input_expr, 'target') else None + expr._solver_parameters = input_expr.solver_parameters \ + if hasattr(input_expr, 'solver_parameters') else None elif len(args) == 2: # origin: ClusterizedEq(lhs, rhs, **kwargs) expr = sympy.Eq.__new__(cls, *args, evaluate=False) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 0b57b876f7..e041e68a9e 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -1,8 +1,9 @@ from collections import OrderedDict from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, - Section, HaloSpot, ExpressionBundle) + Section, HaloSpot, ExpressionBundle, ActionExpr, RHSExpr) from devito.tools import timed_pass +from devito.ir.equations import OpAction, OpRHS __all__ = ['iet_build'] @@ -24,6 +25,12 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) + elif e.operation is OpAction: + exprs.append(ActionExpr(e, operation=e.operation, target=e.target, + solver_parameters=e.solver_parameters)) + elif e.operation is OpRHS: + exprs.append(RHSExpr(e, operation=e.operation, target=e.target, + solver_parameters=e.solver_parameters)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 728e88afbe..fbbd5b4bf4 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -10,7 +10,7 @@ from sympy import IndexedBase, sympify from devito.data import FULL -from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax +from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax, OpAction, OpRHS from devito.ir.support import (INBOUND, SEQUENTIAL, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, VECTORIZED, AFFINE, Property, Forward, WithLock, PrefetchUpdate, detect_io) @@ -28,7 +28,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'Callback'] + 'CallableBody', 'Transfer', 'Callback', 'ActionExpr', 'RHSExpr'] # First-class IET nodes @@ -484,6 +484,24 @@ def __init__(self, expr, pragmas=None): super().__init__(expr, pragmas=pragmas, operation=OpInc) +class ActionExpr(Expression): + + def __init__(self, expr, pragmas=None, operation=OpAction, + target=None, solver_parameters=None): + super().__init__(expr, pragmas=pragmas, operation=operation) + self.target = target + self.solver_parameters = solver_parameters + + +class RHSExpr(Expression): + + def __init__(self, expr, pragmas=None, operation=OpRHS, + target=None, solver_parameters=None): + super().__init__(expr, pragmas=pragmas, operation=operation) + self.target = target + self.solver_parameters = solver_parameters + + class Iteration(Node): """ diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 69a089e12b..64c8817fe3 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -26,7 +26,7 @@ from devito.parameters import configuration from devito.passes import (Graph, lower_index_derivatives, generate_implicit, generate_macros, minimize_symbols, unevaluate, - error_mapper, is_on_device) + error_mapper, is_on_device, lower_petsc) from devito.symbolics import estimate_cost, subs_op_args from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple, flatten, filter_sorted, frozendict, is_integer, @@ -479,6 +479,8 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): graph = Graph(iet, **kwargs) graph = cls._specialize_iet(graph, **kwargs) + lower_petsc(graph, **kwargs) + # Instrument the IET for C-level profiling # Note: this is postponed until after _specialize_iet because during # specialization further Sections may be introduced diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index c09db00c9b..cf6a35de90 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,3 +8,4 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa +from .petsc import * # noqa diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 9d0fbc61cd..f56936f414 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,13 +1,11 @@ from devito.tools import CustomDtype -from devito.types import LocalObject +from devito.types import LocalObject, Eq from devito.types.array import ArrayBasic import numpy as np from cached_property import cached_property from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry -from devito import Eq - class DM(LocalObject): """ @@ -123,11 +121,12 @@ def dtype_to_petsctype(dtype): class PETScEq(Eq): """ - Represents a general equation required by a PETSc solve. + Represents a general equation required by PETScSolve. """ __rkwargs__ = (Eq.__rkwargs__ + ('target', 'solver_parameters',)) + # TODO: Add more solver parameters defaults = { 'ksp_type': 'gmres', 'pc_type': 'jacobi', @@ -171,14 +170,40 @@ class Action(PETScEq): pass +class RHS(PETScEq): + """ + Represents the mathematical expression of building the + rhs of a linear system. + """ + pass + + def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): + # TODO: This is a placeholder for the actual implementation. To start, + # track a single PETScEq i.e an 'Action' through the Operator. + y_matvec = PETScArray(name='y_matvec_'+str(target.name), dtype=target.dtype, dimensions=target.dimensions, shape=target.shape, liveness='eager') - action = Action(y_matvec, eq.lhs, subdomain=eq.subdomain, - implicit_dims=(target.grid.time_dim,), target=target, - solver_parameters=solver_parameters) + x_matvec = PETScArray(name='x_matvec_'+str(target.name), dtype=target.dtype, + dimensions=target.dimensions, + shape=target.shape, liveness='eager') + + b_tmp = PETScArray(name='b_tmp_'+str(target.name), dtype=target.dtype, + dimensions=target.dimensions, + shape=target.shape, liveness='eager') + + # TODO: Extend to rearrange equation for implicit time stepping. + action_tmp = Action(y_matvec, eq.lhs, subdomain=eq.subdomain, target=target, + solver_parameters=solver_parameters) + + rhs = RHS(b_tmp, eq.rhs, subdomain=eq.subdomain, target=target, + solver_parameters=solver_parameters) + + # Only need symbolic representation of equation in mat-vec action callback. + action = action_tmp.subs(target, x_matvec) + - return [action] \ No newline at end of file + return [action] + [rhs] \ No newline at end of file From a49dfa5098e11f4af7feb91cfe5302b95d078961 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:08:42 +0000 Subject: [PATCH 045/151] tests: Add test for PETScSolve --- tests/test_petsc.py | 39 +++++++++++++++++++++++++++++---------- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 9a672a6643..ab0d749dc4 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,8 +1,9 @@ -from devito import Grid, Function, Eq -from devito.ir.iet import Call, ElementalFunction, Definition, DummyExpr +from devito import Grid, Function, Eq, Operator +from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, + ActionExpr, FindNodes, RHSExpr) from devito.passes.iet.languages.C import CDataManager from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP, - PC, KSPConvergedReason, PETScArray) + PC, KSPConvergedReason, PETScArray, PETScSolve) import numpy as np @@ -90,14 +91,32 @@ def test_petsc_subs(): ' + arr(x, y - h_y)/h_y**2 + arr(x, y + h_y)/h_y**2' -# def test_petsc_solve(): +def test_petsc_solve(): -# grid = Grid((2, 2)) + grid = Grid((2, 2)) + + f = Function(name='f', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + eqn = Eq(f.laplace, g) + + petsc = PETScSolve(eqn, f) -# f1 = Function(name='f1', grid=grid, space_order=2) -# f2 = Function(name='f2', grid=grid, space_order=2) + op = Operator(petsc, opt='noop') -# arr = PETScArray(name='arr', dimensions=f2.dimensions, dtype=f2.dtype) + action_expr = FindNodes(ActionExpr).visit(op) + + rhs_expr = FindNodes(RHSExpr).visit(op) + + assert str(action_expr[-1]) == 'y_matvec_f[x][y] = -2.0F*x_matvec_f[x][y]/pow(h_x, 2)' + \ + ' + x_matvec_f[x - 1][y]/pow(h_x, 2) + x_matvec_f[x + 1][y]/pow(h_x, 2)' + \ + ' - 2.0F*x_matvec_f[x][y]/pow(h_y, 2) + x_matvec_f[x][y - 1]/pow(h_y, 2)' + \ + ' + x_matvec_f[x][y + 1]/pow(h_y, 2);' + + assert str(rhs_expr[-1]) == 'b_tmp_f[x][y] = g[x + 2][y + 2];' -# eqn = Eq(f1, f2.laplace) -# eqn_subs = eqn.subs(f2, arr) + # Check the iteration bounds + assert op.arguments().get('x_m') == 0 + assert op.arguments().get('y_m') == 0 + assert op.arguments().get('y_M') == 1 + assert op.arguments().get('x_M') == 1 From 4d9f98e4dc9e416b57f4448f2707bc8eaf976e5e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:24:03 +0000 Subject: [PATCH 046/151] misc: Clean --- devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 2 +- devito/passes/iet/petsc.py | 13 +++++++++++++ devito/types/petsc.py | 16 ++++++++-------- tests/test_petsc.py | 12 ++++++------ 5 files changed, 29 insertions(+), 16 deletions(-) create mode 100644 devito/passes/iet/petsc.py diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index cf5957b907..08b62f28a3 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -377,7 +377,7 @@ def dspace(self): try: if not isinstance(f, PETScArray): if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index c6529ea8ec..4d308eb9e7 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -252,7 +252,7 @@ def __new__(cls, *args, **kwargs): expr._target = input_expr.target if hasattr(input_expr, 'target') else None expr._solver_parameters = input_expr.solver_parameters \ if hasattr(input_expr, 'solver_parameters') else None - + return expr @property diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py new file mode 100644 index 0000000000..264d019e63 --- /dev/null +++ b/devito/passes/iet/petsc.py @@ -0,0 +1,13 @@ +from devito.passes.iet.engine import iet_pass + +__all__ = ['lower_petsc'] + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # TODO: This is a placeholder for the actual PETSc lowering. + # action_expr = FindNodes(ActionExpr).visit(iet) + # rhs_expr = FindNodes(RHSExpr).visit(iet) + + return iet, {} diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f56936f414..bfc3b1fc97 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -7,6 +7,7 @@ from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry + class DM(LocalObject): """ PETSc Data Management object (DM). @@ -159,7 +160,7 @@ def target(self): @property def solver_parameters(self): return self._solver_parameters - + class Action(PETScEq): """ @@ -186,24 +187,23 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): y_matvec = PETScArray(name='y_matvec_'+str(target.name), dtype=target.dtype, dimensions=target.dimensions, shape=target.shape, liveness='eager') - + x_matvec = PETScArray(name='x_matvec_'+str(target.name), dtype=target.dtype, dimensions=target.dimensions, shape=target.shape, liveness='eager') - + b_tmp = PETScArray(name='b_tmp_'+str(target.name), dtype=target.dtype, dimensions=target.dimensions, shape=target.shape, liveness='eager') - - # TODO: Extend to rearrange equation for implicit time stepping. + + # TODO: Extend to rearrange equation for implicit time stepping. action_tmp = Action(y_matvec, eq.lhs, subdomain=eq.subdomain, target=target, solver_parameters=solver_parameters) - + rhs = RHS(b_tmp, eq.rhs, subdomain=eq.subdomain, target=target, solver_parameters=solver_parameters) # Only need symbolic representation of equation in mat-vec action callback. action = action_tmp.subs(target, x_matvec) - - return [action] + [rhs] \ No newline at end of file + return [action] + [rhs] diff --git a/tests/test_petsc.py b/tests/test_petsc.py index ab0d749dc4..94f50480fb 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -89,7 +89,7 @@ def test_petsc_subs(): assert str(eqn_subs.rhs.evaluate) == '-2.0*arr(x, y)/h_x**2' + \ ' + arr(x - h_x, y)/h_x**2 + arr(x + h_x, y)/h_x**2 - 2.0*arr(x, y)/h_y**2' + \ ' + arr(x, y - h_y)/h_y**2 + arr(x, y + h_y)/h_y**2' - + def test_petsc_solve(): @@ -108,11 +108,11 @@ def test_petsc_solve(): rhs_expr = FindNodes(RHSExpr).visit(op) - assert str(action_expr[-1]) == 'y_matvec_f[x][y] = -2.0F*x_matvec_f[x][y]/pow(h_x, 2)' + \ - ' + x_matvec_f[x - 1][y]/pow(h_x, 2) + x_matvec_f[x + 1][y]/pow(h_x, 2)' + \ - ' - 2.0F*x_matvec_f[x][y]/pow(h_y, 2) + x_matvec_f[x][y - 1]/pow(h_y, 2)' + \ - ' + x_matvec_f[x][y + 1]/pow(h_y, 2);' - + assert str(action_expr[-1]) == 'y_matvec_f[x][y] =' + \ + ' -2.0F*x_matvec_f[x][y]/pow(h_x, 2) + x_matvec_f[x - 1][y]/pow(h_x, 2)' + \ + ' + x_matvec_f[x + 1][y]/pow(h_x, 2) - 2.0F*x_matvec_f[x][y]/pow(h_y, 2)' + \ + ' + x_matvec_f[x][y - 1]/pow(h_y, 2) + x_matvec_f[x][y + 1]/pow(h_y, 2);' + assert str(rhs_expr[-1]) == 'b_tmp_f[x][y] = g[x + 2][y + 2];' # Check the iteration bounds From 944adb08d02a12427bdf447e4f66b976341f03db Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:27:50 +0000 Subject: [PATCH 047/151] misc: Clean up --- devito/ir/clusters/cluster.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 08b62f28a3..25a0e0eb5f 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -375,6 +375,9 @@ def dspace(self): else: d = i.dim try: + # TODO: This is the one check that needs to be turned off + # for PETScArrays otherwise the iteration bounds are incorrect. + # Need to discuss this. if not isinstance(f, PETScArray): if i.lower < 0 or \ i.upper > f._size_nodomain[d].left + f._size_halo[d].right: From 5924933a00a3141ef698be0f456bb16f5b73de90 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:30:39 +0000 Subject: [PATCH 048/151] misc: Add a comment in types/petsc.py --- devito/types/petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index bfc3b1fc97..14e84916f3 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -182,7 +182,7 @@ class RHS(PETScEq): def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # TODO: This is a placeholder for the actual implementation. To start, - # track a single PETScEq i.e an 'Action' through the Operator. + # track different PETScEq's (Action, RHS) through the Operator. y_matvec = PETScArray(name='y_matvec_'+str(target.name), dtype=target.dtype, dimensions=target.dimensions, From 01f099fc1a80c386c4995059c800e1bfc57993b1 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 17:04:41 +0000 Subject: [PATCH 049/151] types: Update PETScEq default solver params --- devito/types/petsc.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 14e84916f3..9ab6f5b667 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -127,14 +127,10 @@ class PETScEq(Eq): __rkwargs__ = (Eq.__rkwargs__ + ('target', 'solver_parameters',)) - # TODO: Add more solver parameters + # TODO: Add more default solver parameters. defaults = { 'ksp_type': 'gmres', - 'pc_type': 'jacobi', - 'ksp_rtol': 'PETSC_DEFAULT', - 'ksp_atol': 'PETSC_DEFAULT', - 'ksp_divtol': 'PETSC_DEFAULT', - 'ksp_max_it': 'PETSC_DEFAULT' + 'pc_type': 'jacobi' } def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None, implicit_dims=None, From 80e9b693c660133c7c0aecaeaef705abb0df7ed4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 27 Mar 2024 14:03:24 +0000 Subject: [PATCH 050/151] types: Remove hasattr from LoweredEq, create common node class for petsc exprs and edit iteration check in clsuter.py --- devito/ir/clusters/cluster.py | 25 +++++++++++-------------- devito/ir/equations/equation.py | 23 ++++++++++++++--------- devito/ir/iet/algorithms.py | 17 ++++++++++------- devito/ir/iet/nodes.py | 25 +++++++++++++++++-------- devito/passes/iet/petsc.py | 4 ++-- devito/types/petsc.py | 22 ++++++++++------------ tests/test_petsc.py | 6 +++--- 7 files changed, 67 insertions(+), 55 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 25a0e0eb5f..d85dd9066c 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -369,25 +369,22 @@ def dspace(self): # OOB accesses oobs = set() for f, v in parts.items(): - for i in v: - if i.dim.is_Sub: - d = i.dim.parent - else: - d = i.dim - try: - # TODO: This is the one check that needs to be turned off - # for PETScArrays otherwise the iteration bounds are incorrect. - # Need to discuss this. - if not isinstance(f, PETScArray): + if not isinstance(f, PETScArray): + for i in v: + if i.dim.is_Sub: + d = i.dim.parent + else: + d = i.dim + try: if i.lower < 0 or \ i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) - except (KeyError, TypeError): - # Unable to detect presence of OOB accesses (e.g., `d` not in - # `f._size_halo`, that is typical of indirect accesses `A[B[i]]`) - pass + except (KeyError, TypeError): + # Unable to detect presence of OOB accesses (e.g., `d` not in + # `f._size_halo`, that is typical of indirect accesses `A[B[i]]`) + pass # Construct the `intervals` of the DataSpace, that is a global, # Dimension-centric view of the data space diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 4d308eb9e7..56e194452c 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -9,8 +9,8 @@ Stencil, detect_io, detect_accesses) from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict -from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min -from devito.types.petsc import Action, RHS +from devito.types import (Eq, Inc, ReduceMax, ReduceMin, + relational_min, MatVecEq, RHSEq, LinearSolveEq) __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', 'identity_mapper'] @@ -113,8 +113,8 @@ def detect(cls, expr): Inc: OpInc, ReduceMax: OpMax, ReduceMin: OpMin, - Action: OpAction, - RHS: OpRHS, + MatVecEq: OpMatVec, + RHSEq: OpRHS, } try: return reduction_mapper[type(expr)] @@ -131,7 +131,11 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') -OpAction = Operation('action') + +# Operations required by a Linear Solve of the form Ax=b: +# Application of linear operator on a vector -> op for matrix-vector multiplication. +OpMatVec = Operation('matvec') +# Building the right-hand side of linear system. OpRHS = Operation('rhs') @@ -249,9 +253,10 @@ def __new__(cls, *args, **kwargs): expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - expr._target = input_expr.target if hasattr(input_expr, 'target') else None + expr._target = input_expr.target \ + if isinstance(input_expr, LinearSolveEq) else None expr._solver_parameters = input_expr.solver_parameters \ - if hasattr(input_expr, 'solver_parameters') else None + if isinstance(input_expr, LinearSolveEq) else None return expr @@ -309,9 +314,9 @@ def __new__(cls, *args, **kwargs): expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) expr._target = input_expr.target \ - if hasattr(input_expr, 'target') else None + if isinstance(input_expr, LinearSolveEq) else None expr._solver_parameters = input_expr.solver_parameters \ - if hasattr(input_expr, 'solver_parameters') else None + if isinstance(input_expr, LinearSolveEq) else None elif len(args) == 2: # origin: ClusterizedEq(lhs, rhs, **kwargs) expr = sympy.Eq.__new__(cls, *args, evaluate=False) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index e041e68a9e..c1b7e68f49 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -1,9 +1,10 @@ from collections import OrderedDict from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, - Section, HaloSpot, ExpressionBundle, ActionExpr, RHSExpr) + Section, HaloSpot, ExpressionBundle, MatVecAction, + RHSLinearSystem) from devito.tools import timed_pass -from devito.ir.equations import OpAction, OpRHS +from devito.ir.equations import OpMatVec, OpRHS __all__ = ['iet_build'] @@ -25,12 +26,14 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) - elif e.operation is OpAction: - exprs.append(ActionExpr(e, operation=e.operation, target=e.target, - solver_parameters=e.solver_parameters)) + elif e.operation is OpMatVec: + exprs.append(MatVecAction(e, operation=e.operation, + target=e.target, + solver_parameters=e.solver_parameters)) elif e.operation is OpRHS: - exprs.append(RHSExpr(e, operation=e.operation, target=e.target, - solver_parameters=e.solver_parameters)) + exprs.append(RHSLinearSystem(e, operation=e.operation, + target=e.target, + solver_parameters=e.solver_parameters)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index fbbd5b4bf4..eee152daed 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -10,7 +10,7 @@ from sympy import IndexedBase, sympify from devito.data import FULL -from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax, OpAction, OpRHS +from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax, OpMatVec, OpRHS from devito.ir.support import (INBOUND, SEQUENTIAL, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, VECTORIZED, AFFINE, Property, Forward, WithLock, PrefetchUpdate, detect_io) @@ -28,7 +28,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'Callback', 'ActionExpr', 'RHSExpr'] + 'CallableBody', 'Transfer', 'Callback', 'MatVecAction', 'RHSLinearSystem'] # First-class IET nodes @@ -484,22 +484,31 @@ def __init__(self, expr, pragmas=None): super().__init__(expr, pragmas=pragmas, operation=OpInc) -class ActionExpr(Expression): +class LinearSolverExpression(Expression): + """General expression required by a matrix-free linear solve of the + form Ax=b.""" - def __init__(self, expr, pragmas=None, operation=OpAction, + def __init__(self, expr, pragmas=None, operation=None, target=None, solver_parameters=None): super().__init__(expr, pragmas=pragmas, operation=operation) self.target = target self.solver_parameters = solver_parameters -class RHSExpr(Expression): +class MatVecAction(LinearSolverExpression): + + def __init__(self, expr, pragmas=None, operation=OpMatVec, + target=None, solver_parameters=None): + super().__init__(expr, pragmas=pragmas, operation=operation, + target=target, solver_parameters=solver_parameters) + + +class RHSLinearSystem(LinearSolverExpression): def __init__(self, expr, pragmas=None, operation=OpRHS, target=None, solver_parameters=None): - super().__init__(expr, pragmas=pragmas, operation=operation) - self.target = target - self.solver_parameters = solver_parameters + super().__init__(expr, pragmas=pragmas, operation=operation, + target=target, solver_parameters=solver_parameters) class Iteration(Node): diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index 264d019e63..a4eecedacb 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -1,5 +1,6 @@ from devito.passes.iet.engine import iet_pass + __all__ = ['lower_petsc'] @@ -7,7 +8,6 @@ def lower_petsc(iet, **kwargs): # TODO: This is a placeholder for the actual PETSc lowering. - # action_expr = FindNodes(ActionExpr).visit(iet) + # action_expr = FindNodes(MatVecAction).visit(iet) # rhs_expr = FindNodes(RHSExpr).visit(iet) - return iet, {} diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 9ab6f5b667..34e43df9cb 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -120,7 +120,7 @@ def dtype_to_petsctype(dtype): }[dtype] -class PETScEq(Eq): +class LinearSolveEq(Eq): """ Represents a general equation required by PETScSolve. """ @@ -158,7 +158,7 @@ def solver_parameters(self): return self._solver_parameters -class Action(PETScEq): +class MatVecEq(LinearSolveEq): """ Represents the mathematical expression of applying a linear operator to a vector. This is a key component @@ -167,7 +167,7 @@ class Action(PETScEq): pass -class RHS(PETScEq): +class RHSEq(LinearSolveEq): """ Represents the mathematical expression of building the rhs of a linear system. @@ -178,7 +178,7 @@ class RHS(PETScEq): def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # TODO: This is a placeholder for the actual implementation. To start, - # track different PETScEq's (Action, RHS) through the Operator. + # track different PETScEq's (MatVecAction, RHS) through the Operator. y_matvec = PETScArray(name='y_matvec_'+str(target.name), dtype=target.dtype, dimensions=target.dimensions, @@ -193,13 +193,11 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): shape=target.shape, liveness='eager') # TODO: Extend to rearrange equation for implicit time stepping. - action_tmp = Action(y_matvec, eq.lhs, subdomain=eq.subdomain, target=target, - solver_parameters=solver_parameters) + matvecaction = MatVecEq(y_matvec, eq.lhs.subs(target, x_matvec), + subdomain=eq.subdomain, target=target, + solver_parameters=solver_parameters) - rhs = RHS(b_tmp, eq.rhs, subdomain=eq.subdomain, target=target, - solver_parameters=solver_parameters) + rhs = RHSEq(b_tmp, eq.rhs, subdomain=eq.subdomain, target=target, + solver_parameters=solver_parameters) - # Only need symbolic representation of equation in mat-vec action callback. - action = action_tmp.subs(target, x_matvec) - - return [action] + [rhs] + return [matvecaction] + [rhs] diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 94f50480fb..92798e1024 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,6 +1,6 @@ from devito import Grid, Function, Eq, Operator from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, - ActionExpr, FindNodes, RHSExpr) + MatVecAction, FindNodes, RHSLinearSystem) from devito.passes.iet.languages.C import CDataManager from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, PETScSolve) @@ -104,9 +104,9 @@ def test_petsc_solve(): op = Operator(petsc, opt='noop') - action_expr = FindNodes(ActionExpr).visit(op) + action_expr = FindNodes(MatVecAction).visit(op) - rhs_expr = FindNodes(RHSExpr).visit(op) + rhs_expr = FindNodes(RHSLinearSystem).visit(op) assert str(action_expr[-1]) == 'y_matvec_f[x][y] =' + \ ' -2.0F*x_matvec_f[x][y]/pow(h_x, 2) + x_matvec_f[x - 1][y]/pow(h_x, 2)' + \ From 807e93ce3fcb30ae4d560d65114edb792b125c80 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 27 Mar 2024 14:07:15 +0000 Subject: [PATCH 051/151] compiler: Edit iteration bound check in cluster.py i.e change PETScArray isinstance check to ArrayBasic --- devito/ir/clusters/cluster.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index d85dd9066c..473707c0d9 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -369,7 +369,7 @@ def dspace(self): # OOB accesses oobs = set() for f, v in parts.items(): - if not isinstance(f, PETScArray): + if not isinstance(f, ArrayBasic): for i in v: if i.dim.is_Sub: d = i.dim.parent From 8f8c3c07fd5d656d6378fc881a3898708ea9429b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 09:27:24 +0000 Subject: [PATCH 052/151] compiler: Remove deep nesting in cluster.py isinstance check for ArrayBasic --- devito/ir/clusters/cluster.py | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 473707c0d9..2547c1ca94 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -369,22 +369,23 @@ def dspace(self): # OOB accesses oobs = set() for f, v in parts.items(): - if not isinstance(f, ArrayBasic): - for i in v: - if i.dim.is_Sub: - d = i.dim.parent - else: - d = i.dim - try: - if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: - # It'd mean trying to access a point before the - # left halo (test0) or after the right halo (test1) - oobs.update(d._defines) - except (KeyError, TypeError): - # Unable to detect presence of OOB accesses (e.g., `d` not in - # `f._size_halo`, that is typical of indirect accesses `A[B[i]]`) - pass + if isinstance(f, ArrayBasic): + continue + for i in v: + if i.dim.is_Sub: + d = i.dim.parent + else: + d = i.dim + try: + if i.lower < 0 or \ + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + # It'd mean trying to access a point before the + # left halo (test0) or after the right halo (test1) + oobs.update(d._defines) + except (KeyError, TypeError): + # Unable to detect presence of OOB accesses (e.g., `d` not in + # `f._size_halo`, that is typical of indirect accesses `A[B[i]]`) + pass # Construct the `intervals` of the DataSpace, that is a global, # Dimension-centric view of the data space From efc228918f1cd3356118f714d3f368cac61dc0ef Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 12:55:01 +0000 Subject: [PATCH 053/151] compiler: Replace additions in LoweredEq with PETScRHS class which is a subclass of sympy.Expr - used to track target/solver_params through operator lowering --- devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 25 ++------ devito/ir/iet/algorithms.py | 8 +-- devito/operator/operator.py | 2 +- devito/passes/iet/petsc.py | 4 +- devito/symbolics/printer.py | 3 + devito/types/petsc.py | 106 ++++++++++++++++++-------------- tests/test_petsc.py | 10 +++ 8 files changed, 82 insertions(+), 78 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 2547c1ca94..67e231ec70 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -378,7 +378,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 56e194452c..008448c29b 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -10,7 +10,7 @@ from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict from devito.types import (Eq, Inc, ReduceMax, ReduceMin, - relational_min, MatVecEq, RHSEq, LinearSolveEq) + relational_min, MatVecEq, RHSEq) __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', 'identity_mapper'] @@ -19,8 +19,7 @@ class IREq(sympy.Eq, Pickable): __rargs__ = ('lhs', 'rhs') - __rkwargs__ = ('ispace', 'conditionals', 'implicit_dims', 'operation', - 'target', 'solver_parameters') + __rkwargs__ = ('ispace', 'conditionals', 'implicit_dims', 'operation') @property def is_Scalar(self): @@ -60,14 +59,6 @@ def state(self): def operation(self): return self._operation - @property - def target(self): - return self._target - - @property - def solver_parameters(self): - return self._solver_parameters - @property def is_Reduction(self): return self.operation in (OpInc, OpMin, OpMax) @@ -241,22 +232,18 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - + # from IPython import embed; embed() # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - + # from IPython import embed; embed() expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - expr._target = input_expr.target \ - if isinstance(input_expr, LinearSolveEq) else None - expr._solver_parameters = input_expr.solver_parameters \ - if isinstance(input_expr, LinearSolveEq) else None return expr @@ -313,10 +300,6 @@ def __new__(cls, *args, **kwargs): expr._conditionals = kwargs.get('conditionals', frozendict()) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - expr._target = input_expr.target \ - if isinstance(input_expr, LinearSolveEq) else None - expr._solver_parameters = input_expr.solver_parameters \ - if isinstance(input_expr, LinearSolveEq) else None elif len(args) == 2: # origin: ClusterizedEq(lhs, rhs, **kwargs) expr = sympy.Eq.__new__(cls, *args, evaluate=False) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index c1b7e68f49..54459fada5 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -27,13 +27,9 @@ def iet_build(stree): if e.is_Increment: exprs.append(Increment(e)) elif e.operation is OpMatVec: - exprs.append(MatVecAction(e, operation=e.operation, - target=e.target, - solver_parameters=e.solver_parameters)) + exprs.append(MatVecAction(e, operation=e.operation)) elif e.operation is OpRHS: - exprs.append(RHSLinearSystem(e, operation=e.operation, - target=e.target, - solver_parameters=e.solver_parameters)) + exprs.append(RHSLinearSystem(e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 64c8817fe3..038a9b84cb 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -353,7 +353,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - + # from IPython import embed; embed() return processed # Compilation -- Cluster level diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index a4eecedacb..b998d1691e 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -1,6 +1,5 @@ from devito.passes.iet.engine import iet_pass - __all__ = ['lower_petsc'] @@ -8,6 +7,5 @@ def lower_petsc(iet, **kwargs): # TODO: This is a placeholder for the actual PETSc lowering. - # action_expr = FindNodes(MatVecAction).visit(iet) - # rhs_expr = FindNodes(RHSExpr).visit(iet) + return iet, {} diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index 2c366a389e..e7b395d436 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,6 +242,9 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) + def _print_PETScRHS(self, expr): + return "%s" % self._print(expr._expr) + _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 34e43df9cb..f26819fb28 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -6,6 +6,8 @@ from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry +import sympy +from devito.tools import Reconstructable class DM(LocalObject): @@ -120,45 +122,7 @@ def dtype_to_petsctype(dtype): }[dtype] -class LinearSolveEq(Eq): - """ - Represents a general equation required by PETScSolve. - """ - - __rkwargs__ = (Eq.__rkwargs__ + ('target', 'solver_parameters',)) - - # TODO: Add more default solver parameters. - defaults = { - 'ksp_type': 'gmres', - 'pc_type': 'jacobi' - } - - def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None, implicit_dims=None, - target=None, solver_parameters=None, **kwargs): - - if solver_parameters is None: - solver_parameters = cls.defaults - else: - for key, val in cls.defaults.items(): - solver_parameters[key] = solver_parameters.get(key, val) - - obj = Eq.__new__(cls, lhs, rhs, subdomain=subdomain, coefficients=coefficients, - implicit_dims=implicit_dims, **kwargs) - obj._target = target - obj._solver_parameters = solver_parameters - - return obj - - @property - def target(self): - return self._target - - @property - def solver_parameters(self): - return self._solver_parameters - - -class MatVecEq(LinearSolveEq): +class MatVecEq(Eq): """ Represents the mathematical expression of applying a linear operator to a vector. This is a key component @@ -167,7 +131,7 @@ class MatVecEq(LinearSolveEq): pass -class RHSEq(LinearSolveEq): +class RHSEq(Eq): """ Represents the mathematical expression of building the rhs of a linear system. @@ -192,12 +156,62 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): dimensions=target.dimensions, shape=target.shape, liveness='eager') - # TODO: Extend to rearrange equation for implicit time stepping. - matvecaction = MatVecEq(y_matvec, eq.lhs.subs(target, x_matvec), - subdomain=eq.subdomain, target=target, - solver_parameters=solver_parameters) + # # TODO: Extend to rearrange equation for implicit time stepping. + matvecaction = MatVecEq(y_matvec, PETScRHS(eq.lhs.subs(target, x_matvec), + target=target, solver_parameters=solver_parameters), + subdomain=eq.subdomain) - rhs = RHSEq(b_tmp, eq.rhs, subdomain=eq.subdomain, target=target, - solver_parameters=solver_parameters) + rhs = RHSEq(b_tmp, PETScRHS(eq.rhs, target=target, + solver_parameters=solver_parameters), subdomain=eq.subdomain) return [matvecaction] + [rhs] + + +class PETScRHS(sympy.Expr, Reconstructable): + + __rargs__ = ('expr', 'target', 'solver_parameters',) + + defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'jacobi' + } + + def __new__(cls, expr, target, solver_parameters, **kwargs): + + if solver_parameters is None: + solver_parameters = cls.defaults + else: + for key, val in cls.defaults.items(): + solver_parameters[key] = solver_parameters.get(key, val) + + obj = sympy.Expr.__new__(cls, expr) + obj._expr = expr + obj._target = target + obj._solver_parameters = solver_parameters + + return obj + + def __repr__(self): + return "%s" % self.expr + + __str__ = __repr__ + + def _sympystr(self, printer): + return str(self) + + def _hashable_content(self): + return super()._hashable_content() + (self.target,) + + @property + def expr(self): + return self._expr + + @property + def target(self): + return self._target + + @property + def solver_parameters(self): + return self._solver_parameters + + func = Reconstructable._rebuild diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 92798e1024..5f8798b63c 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -120,3 +120,13 @@ def test_petsc_solve(): assert op.arguments().get('y_m') == 0 assert op.arguments().get('y_M') == 1 assert op.arguments().get('x_M') == 1 + + # Check the target + assert rhs_expr[-1].expr.rhs.target == f + assert action_expr[-1].expr.rhs.target == f + + # Check the solver parameters + assert rhs_expr[-1].expr.rhs.solver_parameters == \ + {'ksp_type': 'gmres', 'pc_type': 'jacobi'} + assert action_expr[-1].expr.rhs.solver_parameters == \ + {'ksp_type': 'gmres', 'pc_type': 'jacobi'} From 7ecaf4a0a72051527c7c41fefe42d8957cde287d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 12:57:59 +0000 Subject: [PATCH 054/151] misc: Clean --- devito/ir/equations/equation.py | 4 ++-- devito/operator/operator.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 008448c29b..6c3fad732b 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -232,13 +232,13 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - # from IPython import embed; embed() + # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - # from IPython import embed; embed() + expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 038a9b84cb..64c8817fe3 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -353,7 +353,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - # from IPython import embed; embed() + return processed # Compilation -- Cluster level From d99d672581cb04386202a8742589f68d3973fac2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 13:01:40 +0000 Subject: [PATCH 055/151] compiler: Remove target and solver params from LinearSolverExpression --- devito/ir/iet/nodes.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index eee152daed..1e9a226f00 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -488,27 +488,20 @@ class LinearSolverExpression(Expression): """General expression required by a matrix-free linear solve of the form Ax=b.""" - def __init__(self, expr, pragmas=None, operation=None, - target=None, solver_parameters=None): + def __init__(self, expr, pragmas=None, operation=None): super().__init__(expr, pragmas=pragmas, operation=operation) - self.target = target - self.solver_parameters = solver_parameters class MatVecAction(LinearSolverExpression): - def __init__(self, expr, pragmas=None, operation=OpMatVec, - target=None, solver_parameters=None): - super().__init__(expr, pragmas=pragmas, operation=operation, - target=target, solver_parameters=solver_parameters) + def __init__(self, expr, pragmas=None, operation=OpMatVec): + super().__init__(expr, pragmas=pragmas, operation=operation) class RHSLinearSystem(LinearSolverExpression): - def __init__(self, expr, pragmas=None, operation=OpRHS, - target=None, solver_parameters=None): - super().__init__(expr, pragmas=pragmas, operation=operation, - target=target, solver_parameters=solver_parameters) + def __init__(self, expr, pragmas=None, operation=OpRHS): + super().__init__(expr, pragmas=pragmas, operation=operation) class Iteration(Node): From f62a12b090869d1de53c5084b02be86c1d974325 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 13:36:15 +0000 Subject: [PATCH 056/151] compiler: Reduce elif statements inside iet/algorithms by using a PETSc mapper --- devito/ir/iet/algorithms.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 54459fada5..8fd3a06a3e 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -4,6 +4,7 @@ Section, HaloSpot, ExpressionBundle, MatVecAction, RHSLinearSystem) from devito.tools import timed_pass +from devito.types import PETScRHS from devito.ir.equations import OpMatVec, OpRHS __all__ = ['iet_build'] @@ -26,10 +27,8 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) - elif e.operation is OpMatVec: - exprs.append(MatVecAction(e, operation=e.operation)) - elif e.operation is OpRHS: - exprs.append(RHSLinearSystem(e, operation=e.operation)) + elif isinstance(e.rhs, PETScRHS): + exprs.append(Op_to_Expr(e.operation)(e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) @@ -61,3 +60,12 @@ def iet_build(stree): queues.setdefault(i.parent, []).append(body) assert False + + +def Op_to_Expr(operation): + """Map Eq operation to IET Expression type.""" + + return { + OpMatVec: MatVecAction, + OpRHS: RHSLinearSystem + }[operation] From fb6c9365e2633a5c43ec2defd2aa50b9e29d280e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 13:51:06 +0000 Subject: [PATCH 057/151] types: Change name from PETScRHS to LinearSolveExpr --- devito/ir/iet/algorithms.py | 4 ++-- devito/symbolics/printer.py | 2 +- devito/types/petsc.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 8fd3a06a3e..cb56a2fc45 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -4,7 +4,7 @@ Section, HaloSpot, ExpressionBundle, MatVecAction, RHSLinearSystem) from devito.tools import timed_pass -from devito.types import PETScRHS +from devito.types import LinearSolveExpr from devito.ir.equations import OpMatVec, OpRHS __all__ = ['iet_build'] @@ -27,7 +27,7 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) - elif isinstance(e.rhs, PETScRHS): + elif isinstance(e.rhs, LinearSolveExpr): exprs.append(Op_to_Expr(e.operation)(e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index e7b395d436..3769e3d229 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,7 +242,7 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) - def _print_PETScRHS(self, expr): + def _print_LinearSolveExpr(self, expr): return "%s" % self._print(expr._expr) _print_EvalDerivative = C99CodePrinter._print_Add diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f26819fb28..89b2096bcf 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -157,17 +157,17 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): shape=target.shape, liveness='eager') # # TODO: Extend to rearrange equation for implicit time stepping. - matvecaction = MatVecEq(y_matvec, PETScRHS(eq.lhs.subs(target, x_matvec), + matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) - rhs = RHSEq(b_tmp, PETScRHS(eq.rhs, target=target, + rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) return [matvecaction] + [rhs] -class PETScRHS(sympy.Expr, Reconstructable): +class LinearSolveExpr(sympy.Expr, Reconstructable): __rargs__ = ('expr', 'target', 'solver_parameters',) From d232332fd648360f5ee3a5237771923c7471ed4a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 14:33:54 +0000 Subject: [PATCH 058/151] types: Edit LinearSolveExpr --- devito/types/petsc.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 89b2096bcf..c29576daf0 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -167,16 +167,17 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): return [matvecaction] + [rhs] -class LinearSolveExpr(sympy.Expr, Reconstructable): +class LinearSolveExpr(sympy.Function, Reconstructable): - __rargs__ = ('expr', 'target', 'solver_parameters',) + __rargs__ = ('expr',) + __rkwargs__ = ('target', 'solver_parameters',) defaults = { 'ksp_type': 'gmres', 'pc_type': 'jacobi' } - def __new__(cls, expr, target, solver_parameters, **kwargs): + def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -184,15 +185,14 @@ def __new__(cls, expr, target, solver_parameters, **kwargs): for key, val in cls.defaults.items(): solver_parameters[key] = solver_parameters.get(key, val) - obj = sympy.Expr.__new__(cls, expr) + obj = sympy.Function.__new__(cls, expr) obj._expr = expr obj._target = target obj._solver_parameters = solver_parameters - return obj def __repr__(self): - return "%s" % self.expr + return "%s(%s)" % (self.__class__.__name__, self.expr) __str__ = __repr__ @@ -200,7 +200,7 @@ def _sympystr(self, printer): return str(self) def _hashable_content(self): - return super()._hashable_content() + (self.target,) + return super()._hashable_content() + (self.expr, self.target) @property def expr(self): From 9fd5a8b2eead9f6657a933c8fab0695b8225e83c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 14:36:40 +0000 Subject: [PATCH 059/151] compiler: Small edit in LinearSolveExpr --- devito/types/petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index c29576daf0..5648e36eb3 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -185,7 +185,7 @@ def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): for key, val in cls.defaults.items(): solver_parameters[key] = solver_parameters.get(key, val) - obj = sympy.Function.__new__(cls, expr) + obj = super().__new__(cls, expr) obj._expr = expr obj._target = target obj._solver_parameters = solver_parameters From 527269e8a4a373c90be11f0c08bffc90e37a3845 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 14:51:59 +0000 Subject: [PATCH 060/151] compiler: Small edit to hash in LinearSolveExpr --- devito/types/petsc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 5648e36eb3..9abc170546 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -199,8 +199,8 @@ def __repr__(self): def _sympystr(self, printer): return str(self) - def _hashable_content(self): - return super()._hashable_content() + (self.expr, self.target) + def __hash__(self): + return hash(self.target) @property def expr(self): From ce3ebba1a49e0a04472671f643988ab8707c2697 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 16:42:32 +0000 Subject: [PATCH 061/151] compiler: Remove function from iet/algorithms and replace with just the mapper --- devito/ir/iet/algorithms.py | 13 +++++-------- devito/ir/iet/nodes.py | 15 +++++++++++++-- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index cb56a2fc45..0177e7dde2 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -28,7 +28,7 @@ def iet_build(stree): if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): - exprs.append(Op_to_Expr(e.operation)(e, operation=e.operation)) + exprs.append(mapper[e.operation](e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) @@ -62,10 +62,7 @@ def iet_build(stree): assert False -def Op_to_Expr(operation): - """Map Eq operation to IET Expression type.""" - - return { - OpMatVec: MatVecAction, - OpRHS: RHSLinearSystem - }[operation] +# Mapping special Eq operations to their corresponding IET Expression subclass types. +# These operations correspond to subclasses of Eq utilised within PETScSolve. +mapper = {OpMatVec: MatVecAction, + OpRHS: RHSLinearSystem} diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 1e9a226f00..4f099143b1 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -485,8 +485,11 @@ def __init__(self, expr, pragmas=None): class LinearSolverExpression(Expression): - """General expression required by a matrix-free linear solve of the - form Ax=b.""" + + """ + General expression required by a matrix-free linear solve of the + form Ax=b. + """ def __init__(self, expr, pragmas=None, operation=None): super().__init__(expr, pragmas=pragmas, operation=operation) @@ -494,12 +497,20 @@ def __init__(self, expr, pragmas=None, operation=None): class MatVecAction(LinearSolverExpression): + """ + Expression representing matrix-vector multiplication. + """ + def __init__(self, expr, pragmas=None, operation=OpMatVec): super().__init__(expr, pragmas=pragmas, operation=operation) class RHSLinearSystem(LinearSolverExpression): + """ + Expression to build the RHS of a linear system. + """ + def __init__(self, expr, pragmas=None, operation=OpRHS): super().__init__(expr, pragmas=pragmas, operation=operation) From d8a57bbd9c89fef50d1897e710eb78946d607d3c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Apr 2024 13:44:34 +0100 Subject: [PATCH 062/151] compiler: Remove unecessary __init__ from LinearSolverExpression --- devito/ir/iet/nodes.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 4f099143b1..9ab5289c66 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -490,9 +490,7 @@ class LinearSolverExpression(Expression): General expression required by a matrix-free linear solve of the form Ax=b. """ - - def __init__(self, expr, pragmas=None, operation=None): - super().__init__(expr, pragmas=pragmas, operation=operation) + pass class MatVecAction(LinearSolverExpression): From f2d05a11fd6333724b363eaa4203d8f3854d11f5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Apr 2024 18:42:57 +0100 Subject: [PATCH 063/151] compiler: Remove printer for LinearSolveExpr --- devito/ir/iet/algorithms.py | 6 +++--- devito/ir/iet/nodes.py | 4 ++-- devito/passes/iet/petsc.py | 3 +++ devito/symbolics/printer.py | 3 --- devito/types/petsc.py | 23 +++++++++-------------- tests/test_petsc.py | 24 ++++++++++++++++-------- 6 files changed, 33 insertions(+), 30 deletions(-) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 0177e7dde2..1a2a9c504b 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -28,7 +28,7 @@ def iet_build(stree): if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): - exprs.append(mapper[e.operation](e, operation=e.operation)) + exprs.append(linsolve_mapper[e.operation](e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) @@ -64,5 +64,5 @@ def iet_build(stree): # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. -mapper = {OpMatVec: MatVecAction, - OpRHS: RHSLinearSystem} +linsolve_mapper = {OpMatVec: MatVecAction, + OpRHS: RHSLinearSystem} diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 9ab5289c66..09d56465a8 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -487,8 +487,8 @@ def __init__(self, expr, pragmas=None): class LinearSolverExpression(Expression): """ - General expression required by a matrix-free linear solve of the - form Ax=b. + Base class for general expressions required by a + matrix-free linear solve of the form Ax=b. """ pass diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index b998d1691e..bc89118f42 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -8,4 +8,7 @@ def lower_petsc(iet, **kwargs): # TODO: This is a placeholder for the actual PETSc lowering. + # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't + # appear in ccode + return iet, {} diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index 3769e3d229..2c366a389e 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,9 +242,6 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) - def _print_LinearSolveExpr(self, expr): - return "%s" % self._print(expr._expr) - _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 9abc170546..7916c37c97 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,12 +1,12 @@ +import sympy +import numpy as np from devito.tools import CustomDtype from devito.types import LocalObject, Eq from devito.types.array import ArrayBasic -import numpy as np from cached_property import cached_property from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry -import sympy from devito.tools import Reconstructable @@ -144,17 +144,12 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # TODO: This is a placeholder for the actual implementation. To start, # track different PETScEq's (MatVecAction, RHS) through the Operator. - y_matvec = PETScArray(name='y_matvec_'+str(target.name), dtype=target.dtype, - dimensions=target.dimensions, - shape=target.shape, liveness='eager') - - x_matvec = PETScArray(name='x_matvec_'+str(target.name), dtype=target.dtype, - dimensions=target.dimensions, - shape=target.shape, liveness='eager') - - b_tmp = PETScArray(name='b_tmp_'+str(target.name), dtype=target.dtype, - dimensions=target.dimensions, - shape=target.shape, liveness='eager') + y_matvec, x_matvec, b_tmp = [ + PETScArray(name=f'{prefix}_{target.name}', + dtype=target.dtype, + dimensions=target.dimensions, + shape=target.shape, liveness='eager') + for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] # # TODO: Extend to rearrange equation for implicit time stepping. matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), @@ -170,7 +165,7 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): class LinearSolveExpr(sympy.Function, Reconstructable): __rargs__ = ('expr',) - __rkwargs__ = ('target', 'solver_parameters',) + __rkwargs__ = ('target', 'solver_parameters') defaults = { 'ksp_type': 'gmres', diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 5f8798b63c..d9c03ec6f8 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,10 +1,10 @@ +import numpy as np from devito import Grid, Function, Eq, Operator from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, MatVecAction, FindNodes, RHSLinearSystem) from devito.passes.iet.languages.C import CDataManager from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, PETScSolve) -import numpy as np def test_petsc_local_object(): @@ -92,7 +92,9 @@ def test_petsc_subs(): def test_petsc_solve(): - + """ + Test PETScSolve. + """ grid = Grid((2, 2)) f = Function(name='f', grid=grid, space_order=2) @@ -108,14 +110,20 @@ def test_petsc_solve(): rhs_expr = FindNodes(RHSLinearSystem).visit(op) - assert str(action_expr[-1]) == 'y_matvec_f[x][y] =' + \ - ' -2.0F*x_matvec_f[x][y]/pow(h_x, 2) + x_matvec_f[x - 1][y]/pow(h_x, 2)' + \ - ' + x_matvec_f[x + 1][y]/pow(h_x, 2) - 2.0F*x_matvec_f[x][y]/pow(h_y, 2)' + \ - ' + x_matvec_f[x][y - 1]/pow(h_y, 2) + x_matvec_f[x][y + 1]/pow(h_y, 2);' + # Verify that the action expression has not been shifted by the + # computational domain since the halo is abstracted within DMDA. + assert str(action_expr[-1].expr.rhs.args[0]) == \ + '-2.0*x_matvec_f[x, y]/h_x**2 + x_matvec_f[x - 1, y]/h_x**2' + \ + ' + x_matvec_f[x + 1, y]/h_x**2 - 2.0*x_matvec_f[x, y]/h_y**2' + \ + ' + x_matvec_f[x, y - 1]/h_y**2 + x_matvec_f[x, y + 1]/h_y**2' - assert str(rhs_expr[-1]) == 'b_tmp_f[x][y] = g[x + 2][y + 2];' + # Verify that the RHS expression has been shifted according to the + # computational domain. This is necessary because the RHS of the + # linear system originates Devito allocated Function objects, not PETSc objects. + assert str(rhs_expr[-1].expr.rhs.args[0]) == 'g[x + 2, y + 2]' - # Check the iteration bounds + # Check the iteration bounds have not changed since PETSc DMDA handles + # negative indexing (the halo is abstracted). assert op.arguments().get('x_m') == 0 assert op.arguments().get('y_m') == 0 assert op.arguments().get('y_M') == 1 From 0cbae6fef622ae577eb2acdc9d10b7b9e18c2aee Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Apr 2024 18:44:31 +0100 Subject: [PATCH 064/151] misc: Clean up --- tests/test_petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d9c03ec6f8..3c40e46b36 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -119,7 +119,7 @@ def test_petsc_solve(): # Verify that the RHS expression has been shifted according to the # computational domain. This is necessary because the RHS of the - # linear system originates Devito allocated Function objects, not PETSc objects. + # linear system is built from Devito allocated Function objects, not PETSc objects. assert str(rhs_expr[-1].expr.rhs.args[0]) == 'g[x + 2, y + 2]' # Check the iteration bounds have not changed since PETSc DMDA handles From eafda55575bed39605c6de10aad3df8d19d364c5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 17 Apr 2024 21:55:51 +0100 Subject: [PATCH 065/151] compiler: Edit PETScArrays to support devito multi-dim indexing --- devito/ir/clusters/cluster.py | 2 -- devito/passes/iet/definitions.py | 4 ++- devito/types/basic.py | 5 ++- devito/types/petsc.py | 34 ++++++++++++++++---- tests/test_petsc.py | 54 ++++++++++++++++++++++---------- 5 files changed, 73 insertions(+), 26 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 67e231ec70..f0b0d718ea 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -369,8 +369,6 @@ def dspace(self): # OOB accesses oobs = set() for f, v in parts.items(): - if isinstance(f, ArrayBasic): - continue for i in v: if i.dim.is_Sub: d = i.dim.parent diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 490b3a4ea8..b692b169af 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -451,7 +451,9 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [i for i in bases if i.name != i.function._C_name] + bases = [ + i for i in bases + if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/types/basic.py b/devito/types/basic.py index bd04be8564..14e3949fb8 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -14,7 +14,7 @@ from devito.data import default_allocator from devito.parameters import configuration from devito.tools import (Pickable, as_tuple, ctypes_to_cstr, dtype_to_ctype, - frozendict, memoized_meth, sympy_mutex) + frozendict, memoized_meth, sympy_mutex, CustomDtype) from devito.types.args import ArgProvider from devito.types.caching import Cached, Uncached from devito.types.lazy import Evaluable @@ -84,6 +84,9 @@ def _C_typedata(self): The type of the object in the generated code as a `str`. """ _type = self._C_ctype + if isinstance(_type, CustomDtype): + return ctypes_to_cstr(_type) + while issubclass(_type, _Pointer): _type = _type._type_ diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 7916c37c97..1b9b5dacfe 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -2,12 +2,14 @@ import numpy as np from devito.tools import CustomDtype from devito.types import LocalObject, Eq +from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic from cached_property import cached_property from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry from devito.tools import Reconstructable +from devito.symbolics import FieldFromComposite class DM(LocalObject): @@ -62,6 +64,14 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') +class DMDALocalInfo(LocalObject): + """ + PETSc object - C struct containing information + about the local grid. + """ + dtype = CustomDtype('DMDALocalInfo') + + class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent @@ -102,13 +112,24 @@ def coefficients(self): @cached_property def _C_ctype(self): - petsc_type = dtype_to_petsctype(self.dtype) - modifier = '*' * len(self.dimensions) - return CustomDtype(petsc_type, modifier=modifier) + return CustomDtype(self.petsc_type, modifier='*') + + @property + def petsc_type(self): + return dtype_to_petsctype(self._dtype) + + @property + def dtype(self): + return CustomDtype(self.petsc_type) @property - def _C_name(self): - return self.name + def symbolic_shape(self): + info = DMDALocalInfo(name='info') + locals = ['gxm', 'gym', 'gzm'] + field_from_composites = [ + FieldFromComposite(lgp, info) for lgp in locals[:len(self.dimensions)]] + ret = tuple(i for i in field_from_composites) + return DimensionTuple(*ret, getters=self.dimensions) def dtype_to_petsctype(dtype): @@ -148,7 +169,8 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): PETScArray(name=f'{prefix}_{target.name}', dtype=target.dtype, dimensions=target.dimensions, - shape=target.shape, liveness='eager') + shape=target.shape, liveness='eager', + halo=target.halo) for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] # # TODO: Extend to rearrange equation for implicit time stepping. diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 3c40e46b36..8cfb94e8cf 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,7 +1,8 @@ import numpy as np from devito import Grid, Function, Eq, Operator from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, - MatVecAction, FindNodes, RHSLinearSystem) + MatVecAction, FindNodes, RHSLinearSystem, + PointerCast) from devito.passes.iet.languages.C import CDataManager from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, PETScSolve) @@ -58,11 +59,11 @@ def test_petsc_functions(): 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(defn0) == 'PetscScalar*restrict ptr0_vec;' + assert str(defn1) == 'const PetscScalar*restrict ptr1_vec;' + assert str(defn2) == 'const PetscScalar*restrict ptr2_vec;' + assert str(defn3) == 'PetscInt*restrict ptr3_vec;' + assert str(defn4) == 'const PetscInt*restrict ptr4_vec;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' @@ -110,20 +111,14 @@ def test_petsc_solve(): rhs_expr = FindNodes(RHSLinearSystem).visit(op) - # Verify that the action expression has not been shifted by the - # computational domain since the halo is abstracted within DMDA. assert str(action_expr[-1].expr.rhs.args[0]) == \ - '-2.0*x_matvec_f[x, y]/h_x**2 + x_matvec_f[x - 1, y]/h_x**2' + \ - ' + x_matvec_f[x + 1, y]/h_x**2 - 2.0*x_matvec_f[x, y]/h_y**2' + \ - ' + x_matvec_f[x, y - 1]/h_y**2 + x_matvec_f[x, y + 1]/h_y**2' + 'x_matvec_f[x + 1, y + 2]/h_x**2 - 2.0*x_matvec_f[x + 2, y + 2]/h_x**2' + \ + ' + x_matvec_f[x + 3, y + 2]/h_x**2 + x_matvec_f[x + 2, y + 1]/h_y**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/h_y**2 + x_matvec_f[x + 2, y + 3]/h_y**2' - # Verify that the RHS expression has been shifted according to the - # computational domain. This is necessary because the RHS of the - # linear system is built from Devito allocated Function objects, not PETSc objects. assert str(rhs_expr[-1].expr.rhs.args[0]) == 'g[x + 2, y + 2]' - # Check the iteration bounds have not changed since PETSc DMDA handles - # negative indexing (the halo is abstracted). + # Check the iteration bounds are correct. assert op.arguments().get('x_m') == 0 assert op.arguments().get('y_m') == 0 assert op.arguments().get('y_M') == 1 @@ -138,3 +133,30 @@ def test_petsc_solve(): {'ksp_type': 'gmres', 'pc_type': 'jacobi'} assert action_expr[-1].expr.rhs.solver_parameters == \ {'ksp_type': 'gmres', 'pc_type': 'jacobi'} + + +def test_petsc_cast(): + """ + Test casting of PETScArray. + """ + g0 = Grid((2)) + g1 = Grid((2, 2)) + g2 = Grid((2, 2, 2)) + + arr0 = PETScArray(name='arr0', dimensions=g0.dimensions, shape=g0.shape) + arr1 = PETScArray(name='arr1', dimensions=g1.dimensions, shape=g1.shape) + arr2 = PETScArray(name='arr2', dimensions=g2.dimensions, shape=g2.shape) + + # Casts will be explictly generated and placed at specific locations in the C code, + # specifically after various other PETSc calls have been executed. + cast0 = PointerCast(arr0) + cast1 = PointerCast(arr1) + cast2 = PointerCast(arr2) + + assert str(cast0) == \ + 'PetscScalar (*restrict arr0) = (PetscScalar (*)) arr0_vec;' + assert str(cast1) == \ + 'PetscScalar (*restrict arr1)[info.gym] = (PetscScalar (*)[info.gym]) arr1_vec;' + assert str(cast2) == \ + 'PetscScalar (*restrict arr2)[info.gym][info.gzm] = ' + \ + '(PetscScalar (*)[info.gym][info.gzm]) arr2_vec;' From bcb892fa5ee2eb591b1afa7b60e7733e122fc578 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 17 Apr 2024 22:11:49 +0100 Subject: [PATCH 066/151] misc: clean --- devito/types/petsc.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 1b9b5dacfe..ffc55bce29 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -125,6 +125,9 @@ def dtype(self): @property def symbolic_shape(self): info = DMDALocalInfo(name='info') + # To access the local grid info via the DM in + # PETSc you use DMDAGetLocalInfo(da, &info) + # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz locals = ['gxm', 'gym', 'gzm'] field_from_composites = [ FieldFromComposite(lgp, info) for lgp in locals[:len(self.dimensions)]] From 133343234c59fe664b90ea9e48e13dea9934fc76 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 09:52:45 +0100 Subject: [PATCH 067/151] tests: Add test for no automatic cast generation with PETScArrays --- tests/test_petsc.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8cfb94e8cf..5b4cf24e63 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -160,3 +160,22 @@ def test_petsc_cast(): assert str(cast2) == \ 'PetscScalar (*restrict arr2)[info.gym][info.gzm] = ' + \ '(PetscScalar (*)[info.gym][info.gzm]) arr2_vec;' + + +def test_no_automatic_cast(): + """ + Verify that the compiler does not automatically generate casts for PETScArrays. + They will be generated at specific points within the C code, particularly after + other PETSc calls, rather than necessarily at the top of the Kernel. + """ + grid = Grid((2, 2)) + + f = Function(name='f', grid=grid, space_order=2) + + arr = PETScArray(name='arr', dimensions=f.dimensions, shape=f.shape) + + eqn = Eq(arr, f.laplace) + + op = Operator(eqn, opt='noop') + + assert len(op.body.casts) == 1 From 754eb0fc01d792d0381f132a2e175aac23cf3ef0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 09:57:36 +0100 Subject: [PATCH 068/151] misc: Docstring edit in test petsc --- tests/test_petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 5b4cf24e63..6e889d8abb 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -164,7 +164,7 @@ def test_petsc_cast(): def test_no_automatic_cast(): """ - Verify that the compiler does not automatically generate casts for PETScArrays. + Verify that the compiler doesn't automatically generate casts for PETScArrays. They will be generated at specific points within the C code, particularly after other PETSc calls, rather than necessarily at the top of the Kernel. """ From 4057c2ce31da49eaf4f7149a5410f48ec622bbcf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 13:23:08 +0100 Subject: [PATCH 069/151] compiler: Add PETScIndexedData --- devito/passes/iet/definitions.py | 11 +++++------ devito/types/petsc.py | 10 ++++++++++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index b692b169af..7141e7bd54 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -437,8 +437,8 @@ def place_casts(self, iet, **kwargs): The input Iteration/Expression tree. """ # Candidates - indexeds = FindSymbols('indexeds|indexedbases').visit(iet) - + # indexeds = FindSymbols('indexeds|indexedbases').visit(iet) + indexeds = FindSymbols('indexeds').visit(iet) # Create Function -> n-dimensional array casts # E.g. `float (*u)[.] = (float (*)[.]) u_vec->data` # NOTE: a cast is needed only if the underlying data object isn't already @@ -446,14 +446,13 @@ def place_casts(self, iet, **kwargs): # (i) Dereferencing a PointerArray, e.g., `float (*r0)[.] = (float(*)[.]) pr0[.]` # (ii) Declaring a raw pointer, e.g., `float * r0 = NULL; *malloc(&(r0), ...) defines = set(FindSymbols('defines|globals').visit(iet)) - bases = sorted({i.base for i in indexeds}, key=lambda i: i.name) + # indexeds = [i for i in indexeds if isinstance(i.base, IndexedData)] + bases = sorted({i.base for i in indexeds if isinstance(i.base, IndexedData)}, key=lambda i: i.name) # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [ - i for i in bases - if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] + bases = [i for i in bases if i.name != i.function._C_name] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/types/petsc.py b/devito/types/petsc.py index ffc55bce29..5ac335e0ff 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -10,6 +10,7 @@ from devito.finite_differences.tools import fd_weights_registry from devito.tools import Reconstructable from devito.symbolics import FieldFromComposite +from devito.types.basic import IndexedBase class DM(LocalObject): @@ -133,6 +134,15 @@ def symbolic_shape(self): FieldFromComposite(lgp, info) for lgp in locals[:len(self.dimensions)]] ret = tuple(i for i in field_from_composites) return DimensionTuple(*ret, getters=self.dimensions) + + @cached_property + def indexed(self): + """The wrapped IndexedData object.""" + return PETScIndexedData(self.name, shape=self._shape, function=self.function) + + +class PETScIndexedData(IndexedBase): + pass def dtype_to_petsctype(dtype): From eb51bd0abf924e5dac300621635ecf35e5d58056 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 13:27:36 +0100 Subject: [PATCH 070/151] misc: Clean --- devito/passes/iet/definitions.py | 7 +++---- devito/types/petsc.py | 5 ++--- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 7141e7bd54..55ea27fe33 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -19,7 +19,7 @@ SizeOf, VOID, Keyword, pow_to_mul) from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap, - DeviceRM, Eq, Symbol) + DeviceRM, Eq, Symbol, IndexedData) __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage'] @@ -437,7 +437,6 @@ def place_casts(self, iet, **kwargs): The input Iteration/Expression tree. """ # Candidates - # indexeds = FindSymbols('indexeds|indexedbases').visit(iet) indexeds = FindSymbols('indexeds').visit(iet) # Create Function -> n-dimensional array casts # E.g. `float (*u)[.] = (float (*)[.]) u_vec->data` @@ -446,8 +445,8 @@ def place_casts(self, iet, **kwargs): # (i) Dereferencing a PointerArray, e.g., `float (*r0)[.] = (float(*)[.]) pr0[.]` # (ii) Declaring a raw pointer, e.g., `float * r0 = NULL; *malloc(&(r0), ...) defines = set(FindSymbols('defines|globals').visit(iet)) - # indexeds = [i for i in indexeds if isinstance(i.base, IndexedData)] - bases = sorted({i.base for i in indexeds if isinstance(i.base, IndexedData)}, key=lambda i: i.name) + bases = sorted({i.base for i in indexeds + if isinstance(i.base, IndexedData)}, key=lambda i: i.name) # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 5ac335e0ff..d6e8132008 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -134,12 +134,11 @@ def symbolic_shape(self): FieldFromComposite(lgp, info) for lgp in locals[:len(self.dimensions)]] ret = tuple(i for i in field_from_composites) return DimensionTuple(*ret, getters=self.dimensions) - + @cached_property def indexed(self): - """The wrapped IndexedData object.""" return PETScIndexedData(self.name, shape=self._shape, function=self.function) - + class PETScIndexedData(IndexedBase): pass From bdaba5e8c6583f3f08228adfe88271c6880714a5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 13:42:54 +0100 Subject: [PATCH 071/151] compiler: Add indexedbases back and simplify PETScArray symbolic shape --- devito/passes/iet/definitions.py | 2 +- devito/types/petsc.py | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 55ea27fe33..284e04516c 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -437,7 +437,7 @@ def place_casts(self, iet, **kwargs): The input Iteration/Expression tree. """ # Candidates - indexeds = FindSymbols('indexeds').visit(iet) + indexeds = FindSymbols('indexeds|indexedbases').visit(iet) # Create Function -> n-dimensional array casts # E.g. `float (*u)[.] = (float (*)[.]) u_vec->data` # NOTE: a cast is needed only if the underlying data object isn't already diff --git a/devito/types/petsc.py b/devito/types/petsc.py index d6e8132008..ca11ffd48f 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -129,11 +129,9 @@ def symbolic_shape(self): # To access the local grid info via the DM in # PETSc you use DMDAGetLocalInfo(da, &info) # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz - locals = ['gxm', 'gym', 'gzm'] field_from_composites = [ - FieldFromComposite(lgp, info) for lgp in locals[:len(self.dimensions)]] - ret = tuple(i for i in field_from_composites) - return DimensionTuple(*ret, getters=self.dimensions) + FieldFromComposite('g%sm' % d.name, info) for d in self.dimensions] + return DimensionTuple(*field_from_composites, getters=self.dimensions) @cached_property def indexed(self): From a488156d3a9fc41619f9ecf17d27a42b9aaf92f6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 25 Apr 2024 16:28:04 +0100 Subject: [PATCH 072/151] types: Reverse dmda and therefore casting to match Devito --- devito/types/petsc.py | 3 ++- tests/test_petsc.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index ca11ffd48f..8ff0c69b76 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -131,7 +131,8 @@ def symbolic_shape(self): # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz field_from_composites = [ FieldFromComposite('g%sm' % d.name, info) for d in self.dimensions] - return DimensionTuple(*field_from_composites, getters=self.dimensions) + # Reverse it since DMDA is setup backwards to Devito dimensions. + return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) @cached_property def indexed(self): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 6e889d8abb..8317204d8f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -156,10 +156,10 @@ def test_petsc_cast(): assert str(cast0) == \ 'PetscScalar (*restrict arr0) = (PetscScalar (*)) arr0_vec;' assert str(cast1) == \ - 'PetscScalar (*restrict arr1)[info.gym] = (PetscScalar (*)[info.gym]) arr1_vec;' + 'PetscScalar (*restrict arr1)[info.gxm] = (PetscScalar (*)[info.gxm]) arr1_vec;' assert str(cast2) == \ - 'PetscScalar (*restrict arr2)[info.gym][info.gzm] = ' + \ - '(PetscScalar (*)[info.gym][info.gzm]) arr2_vec;' + 'PetscScalar (*restrict arr2)[info.gym][info.gxm] = ' + \ + '(PetscScalar (*)[info.gym][info.gxm]) arr2_vec;' def test_no_automatic_cast(): From 59148523dc3a9bb1da4d6862b75c175543b559a5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 26 Apr 2024 13:37:11 +0100 Subject: [PATCH 073/151] misc: Doc --- tests/test_petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8317204d8f..ca71343fa5 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -164,7 +164,7 @@ def test_petsc_cast(): def test_no_automatic_cast(): """ - Verify that the compiler doesn't automatically generate casts for PETScArrays. + Verify that the compiler does not automatically generate casts for PETScArrays. They will be generated at specific points within the C code, particularly after other PETSc calls, rather than necessarily at the top of the Kernel. """ From 41c7ff4cab1d6c4d1955896870e4dab15e965575 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 1 May 2024 10:53:57 +0100 Subject: [PATCH 074/151] misc: Add line back --- devito/passes/iet/definitions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 284e04516c..7f0350fec3 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -438,6 +438,7 @@ def place_casts(self, iet, **kwargs): """ # Candidates indexeds = FindSymbols('indexeds|indexedbases').visit(iet) + # Create Function -> n-dimensional array casts # E.g. `float (*u)[.] = (float (*)[.]) u_vec->data` # NOTE: a cast is needed only if the underlying data object isn't already From 3e0288d4b98a1bf029333e2787787213db5edc93 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 01:12:18 +0100 Subject: [PATCH 075/151] compiler: Add Mock/placeholder expressions to create distinct iteration loops for each component of a linear solve --- devito/ir/equations/equation.py | 5 ++++- devito/ir/iet/algorithms.py | 8 +++++--- devito/ir/iet/nodes.py | 19 +++++++++++++++++-- devito/operator/operator.py | 1 + devito/passes/iet/petsc.py | 30 ++++++++++++++++++++++++++++-- devito/types/petsc.py | 28 +++++++++++++++++++++++++--- 6 files changed, 80 insertions(+), 11 deletions(-) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 6c3fad732b..cc1da7f04a 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -10,7 +10,7 @@ from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict from devito.types import (Eq, Inc, ReduceMax, ReduceMin, - relational_min, MatVecEq, RHSEq) + relational_min, MatVecEq, RHSEq, MockEq) __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', 'identity_mapper'] @@ -106,6 +106,7 @@ def detect(cls, expr): ReduceMin: OpMin, MatVecEq: OpMatVec, RHSEq: OpRHS, + MockEq: OpMock, } try: return reduction_mapper[type(expr)] @@ -128,6 +129,8 @@ def detect(cls, expr): OpMatVec = Operation('matvec') # Building the right-hand side of linear system. OpRHS = Operation('rhs') +# Operation linked to MockEq, placeholders to be removed at the IET level. +OpMock = Operation('mock') identity_mapper = { diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 1a2a9c504b..73495a1287 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -2,10 +2,10 @@ from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, Section, HaloSpot, ExpressionBundle, MatVecAction, - RHSLinearSystem) + RHSLinearSystem, LinSolveMock) from devito.tools import timed_pass from devito.types import LinearSolveExpr -from devito.ir.equations import OpMatVec, OpRHS +from devito.ir.equations import OpMatVec, OpRHS, OpMock __all__ = ['iet_build'] @@ -25,6 +25,7 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: + # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): @@ -65,4 +66,5 @@ def iet_build(stree): # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. linsolve_mapper = {OpMatVec: MatVecAction, - OpRHS: RHSLinearSystem} + OpRHS: RHSLinearSystem, + OpMock: LinSolveMock} diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 09d56465a8..68a3c28380 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -10,7 +10,8 @@ from sympy import IndexedBase, sympify from devito.data import FULL -from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax, OpMatVec, OpRHS +from devito.ir.equations import (DummyEq, OpInc, OpMin, OpMax, OpMatVec, OpRHS, + OpMock) from devito.ir.support import (INBOUND, SEQUENTIAL, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, VECTORIZED, AFFINE, Property, Forward, WithLock, PrefetchUpdate, detect_io) @@ -28,7 +29,8 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'Callback', 'MatVecAction', 'RHSLinearSystem'] + 'CallableBody', 'Transfer', 'Callback', 'MatVecAction', 'RHSLinearSystem', + 'LinSolveMock'] # First-class IET nodes @@ -513,6 +515,19 @@ def __init__(self, expr, pragmas=None, operation=OpRHS): super().__init__(expr, pragmas=pragmas, operation=operation) +class LinSolveMock(LinearSolverExpression): + + """ + Placeholder expression to wrap MockEqs, which are dropped + at the IET level. + """ + # NOTE: The requirement for init=False otherwise there are issues + # inside specialize_iet? + + def __init__(self, expr, init=False, pragmas=None, operation=OpMock): + super().__init__(expr, init=init, pragmas=pragmas, operation=operation) + + class Iteration(Node): """ diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 64c8817fe3..775697d2a0 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -477,6 +477,7 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Lower IET to a target-specific IET graph = Graph(iet, **kwargs) + # from IPython import embed; embed() graph = cls._specialize_iet(graph, **kwargs) lower_petsc(graph, **kwargs) diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index bc89118f42..18250d9db9 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -1,4 +1,6 @@ from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, LinSolveMock, retrieve_iteration_tree, + filter_iterations, Transformer) __all__ = ['lower_petsc'] @@ -6,9 +8,33 @@ @iet_pass def lower_petsc(iet, **kwargs): - # TODO: This is a placeholder for the actual PETSc lowering. - # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't # appear in ccode + # Drop all placeholder expressions previously used to create distinct + # iteration loops for each component of the linear solve. + iet = drop_mocks(iet) + return iet, {} + + +def drop_mocks(iet): + """ + Drop the spatial iteration loop containing each LinSolveMock. + """ + + mapper = {} + + for tree in retrieve_iteration_tree(iet): + + # Eventually, when using implicit dims etc do not want to drop + # the time loop. + root = filter_iterations(tree, key=lambda i: i.dim.is_Space) + mock = FindNodes(LinSolveMock).visit(root) + + if mock: + mapper.update({root[0]: None}) + + iet = Transformer(mapper, nested=True).visit(iet) + + return iet diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 8ff0c69b76..0b133f55cb 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -1,5 +1,6 @@ import sympy import numpy as np + from devito.tools import CustomDtype from devito.types import LocalObject, Eq from devito.types.utils import DimensionTuple @@ -171,10 +172,21 @@ class RHSEq(Eq): pass +class MockEq(Eq): + """ + Represents a mock/placeholder equation to ensure distinct iteration loops. + + For example, the mat-vec action iteration loop is to be isolated from the + expression loop used to build the RHS of the linear system. This separation + facilitates the utilisation of the mat-vec iteration loop in callback functions + created at the IET level. + """ + pass + + def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - # TODO: This is a placeholder for the actual implementation. To start, - # track different PETScEq's (MatVecAction, RHS) through the Operator. + # TODO: Add check for time dimensions and utilise implicit dimensions. y_matvec, x_matvec, b_tmp = [ PETScArray(name=f'{prefix}_{target.name}', @@ -192,7 +204,17 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) - return [matvecaction] + [rhs] + # Create mock equations to ensure distinct iteration loops for each component + # of the linear solve. + indices = tuple(d + 1 for d in target.dimensions) + s0 = Scalar(name='s0') + s1 = Scalar(name='s1') + + # Wrapped rhs in LinearSolveExpr for simplicity in iet_build pass. + mock_action = MockEq(s0, LinearSolveExpr(y_matvec.indexify(indices=indices))) + mock_rhs = MockEq(s1, LinearSolveExpr(b_tmp.indexify(indices=indices))) + + return [matvecaction, mock_action] + [rhs, mock_rhs] class LinearSolveExpr(sympy.Function, Reconstructable): From 17193ea7a9fd21b45b2b3e8468f5808b3db49b3e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 01:18:36 +0100 Subject: [PATCH 076/151] misc: Clean up --- devito/ir/iet/algorithms.py | 1 - devito/ir/iet/nodes.py | 2 +- devito/operator/operator.py | 1 - devito/passes/iet/petsc.py | 3 --- devito/types/petsc.py | 4 ++-- 5 files changed, 3 insertions(+), 8 deletions(-) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 73495a1287..fc9b72740c 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -25,7 +25,6 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: - # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 68a3c28380..3fdafaab26 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -522,7 +522,7 @@ class LinSolveMock(LinearSolverExpression): at the IET level. """ # NOTE: The requirement for init=False otherwise there are issues - # inside specialize_iet? + # inside specialize_iet. def __init__(self, expr, init=False, pragmas=None, operation=OpMock): super().__init__(expr, init=init, pragmas=pragmas, operation=operation) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 775697d2a0..64c8817fe3 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -477,7 +477,6 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Lower IET to a target-specific IET graph = Graph(iet, **kwargs) - # from IPython import embed; embed() graph = cls._specialize_iet(graph, **kwargs) lower_petsc(graph, **kwargs) diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index 18250d9db9..006ebab8ff 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -22,11 +22,8 @@ def drop_mocks(iet): """ Drop the spatial iteration loop containing each LinSolveMock. """ - mapper = {} - for tree in retrieve_iteration_tree(iet): - # Eventually, when using implicit dims etc do not want to drop # the time loop. root = filter_iterations(tree, key=lambda i: i.dim.is_Space) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 0b133f55cb..a11db983a4 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -207,8 +207,8 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # Create mock equations to ensure distinct iteration loops for each component # of the linear solve. indices = tuple(d + 1 for d in target.dimensions) - s0 = Scalar(name='s0') - s1 = Scalar(name='s1') + s0 = Symbol(name='s0') + s1 = Symbol(name='s1') # Wrapped rhs in LinearSolveExpr for simplicity in iet_build pass. mock_action = MockEq(s0, LinearSolveExpr(y_matvec.indexify(indices=indices))) From d55a2d04acacca4d11ad645502f4a782665989ad Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 10:08:12 +0100 Subject: [PATCH 077/151] bench: Remove null profiled sections in Kernel for openmp --- devito/operator/operator.py | 3 ++- devito/operator/profiling.py | 10 +++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 64c8817fe3..b83c49ea04 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -477,10 +477,11 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Lower IET to a target-specific IET graph = Graph(iet, **kwargs) - graph = cls._specialize_iet(graph, **kwargs) lower_petsc(graph, **kwargs) + graph = cls._specialize_iet(graph, **kwargs) + # Instrument the IET for C-level profiling # Note: this is postponed until after _specialize_iet because during # specialization further Sections may be introduced diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 32742489cd..1061d2c61d 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -135,7 +135,15 @@ def instrument(self, iet, timer): for i in sections: n = i.name assert n in timer.fields - mapper[i] = i._rebuild(body=TimedList(timer=timer, lname=n, body=i.body)) + # TODO: Need an improved method for profiling with PETSc..? + # The Mock expressions (along with their corresponding iteration loops) + # are still detected as Sections with an empty body (even though they + # are dropped). If they are dropped here, there is a mismatch between + # the no.of Sections present in the Kernel and no.of Sections displayed + # in the profiler struct. + if i.body != (): + mapper[i] = i._rebuild(body=TimedList(timer=timer, + lname=n, body=i.body)) return Transformer(mapper, nested=True).visit(iet) else: return iet From 76d05c7dbf6aede1a568cdc9f88abd1f8d65053c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 10:12:05 +0100 Subject: [PATCH 078/151] compiler: Remove init=False from LinSolveMock --- devito/ir/iet/nodes.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 3fdafaab26..8d4b8659ad 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -521,11 +521,9 @@ class LinSolveMock(LinearSolverExpression): Placeholder expression to wrap MockEqs, which are dropped at the IET level. """ - # NOTE: The requirement for init=False otherwise there are issues - # inside specialize_iet. - def __init__(self, expr, init=False, pragmas=None, operation=OpMock): - super().__init__(expr, init=init, pragmas=pragmas, operation=operation) + def __init__(self, expr, pragmas=None, operation=OpMock): + super().__init__(expr, pragmas=pragmas, operation=operation) class Iteration(Node): From 5555b674f461aace778d75e4362466d90c43ecee Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 11:14:26 +0100 Subject: [PATCH 079/151] compiler: Simplify by just having a Mock RHS which drops corresponding Cluster before IET level --- devito/ir/equations/equation.py | 7 ++----- devito/ir/iet/algorithms.py | 7 +++---- devito/ir/iet/nodes.py | 17 ++--------------- devito/ir/stree/algorithms.py | 1 + devito/operator/operator.py | 1 + devito/operator/profiling.py | 10 +--------- devito/passes/iet/petsc.py | 25 ------------------------- devito/types/petsc.py | 29 +++++++++++++++++++++++++++-- 8 files changed, 37 insertions(+), 60 deletions(-) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index cc1da7f04a..7d86bcc85c 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -10,7 +10,7 @@ from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict from devito.types import (Eq, Inc, ReduceMax, ReduceMin, - relational_min, MatVecEq, RHSEq, MockEq) + relational_min, MatVecEq, RHSEq) __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', 'identity_mapper'] @@ -105,8 +105,7 @@ def detect(cls, expr): ReduceMax: OpMax, ReduceMin: OpMin, MatVecEq: OpMatVec, - RHSEq: OpRHS, - MockEq: OpMock, + RHSEq: OpRHS } try: return reduction_mapper[type(expr)] @@ -129,8 +128,6 @@ def detect(cls, expr): OpMatVec = Operation('matvec') # Building the right-hand side of linear system. OpRHS = Operation('rhs') -# Operation linked to MockEq, placeholders to be removed at the IET level. -OpMock = Operation('mock') identity_mapper = { diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index fc9b72740c..1a2a9c504b 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -2,10 +2,10 @@ from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, Section, HaloSpot, ExpressionBundle, MatVecAction, - RHSLinearSystem, LinSolveMock) + RHSLinearSystem) from devito.tools import timed_pass from devito.types import LinearSolveExpr -from devito.ir.equations import OpMatVec, OpRHS, OpMock +from devito.ir.equations import OpMatVec, OpRHS __all__ = ['iet_build'] @@ -65,5 +65,4 @@ def iet_build(stree): # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. linsolve_mapper = {OpMatVec: MatVecAction, - OpRHS: RHSLinearSystem, - OpMock: LinSolveMock} + OpRHS: RHSLinearSystem} diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 8d4b8659ad..09d56465a8 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -10,8 +10,7 @@ from sympy import IndexedBase, sympify from devito.data import FULL -from devito.ir.equations import (DummyEq, OpInc, OpMin, OpMax, OpMatVec, OpRHS, - OpMock) +from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax, OpMatVec, OpRHS from devito.ir.support import (INBOUND, SEQUENTIAL, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, VECTORIZED, AFFINE, Property, Forward, WithLock, PrefetchUpdate, detect_io) @@ -29,8 +28,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'Callback', 'MatVecAction', 'RHSLinearSystem', - 'LinSolveMock'] + 'CallableBody', 'Transfer', 'Callback', 'MatVecAction', 'RHSLinearSystem'] # First-class IET nodes @@ -515,17 +513,6 @@ def __init__(self, expr, pragmas=None, operation=OpRHS): super().__init__(expr, pragmas=pragmas, operation=operation) -class LinSolveMock(LinearSolverExpression): - - """ - Placeholder expression to wrap MockEqs, which are dropped - at the IET level. - """ - - def __init__(self, expr, pragmas=None, operation=OpMock): - super().__init__(expr, pragmas=pragmas, operation=operation) - - class Iteration(Node): """ diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index a85b93460d..c738db645a 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,6 +163,7 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: + # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index b83c49ea04..b136a0370d 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -424,6 +424,7 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters + # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 1061d2c61d..32742489cd 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -135,15 +135,7 @@ def instrument(self, iet, timer): for i in sections: n = i.name assert n in timer.fields - # TODO: Need an improved method for profiling with PETSc..? - # The Mock expressions (along with their corresponding iteration loops) - # are still detected as Sections with an empty body (even though they - # are dropped). If they are dropped here, there is a mismatch between - # the no.of Sections present in the Kernel and no.of Sections displayed - # in the profiler struct. - if i.body != (): - mapper[i] = i._rebuild(body=TimedList(timer=timer, - lname=n, body=i.body)) + mapper[i] = i._rebuild(body=TimedList(timer=timer, lname=n, body=i.body)) return Transformer(mapper, nested=True).visit(iet) else: return iet diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py index 006ebab8ff..c050dbbb80 100644 --- a/devito/passes/iet/petsc.py +++ b/devito/passes/iet/petsc.py @@ -1,6 +1,4 @@ from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, LinSolveMock, retrieve_iteration_tree, - filter_iterations, Transformer) __all__ = ['lower_petsc'] @@ -11,27 +9,4 @@ def lower_petsc(iet, **kwargs): # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't # appear in ccode - # Drop all placeholder expressions previously used to create distinct - # iteration loops for each component of the linear solve. - iet = drop_mocks(iet) - return iet, {} - - -def drop_mocks(iet): - """ - Drop the spatial iteration loop containing each LinSolveMock. - """ - mapper = {} - for tree in retrieve_iteration_tree(iet): - # Eventually, when using implicit dims etc do not want to drop - # the time loop. - root = filter_iterations(tree, key=lambda i: i.dim.is_Space) - mock = FindNodes(LinSolveMock).visit(root) - - if mock: - mapper.update({root[0]: None}) - - iet = Transformer(mapper, nested=True).visit(iet) - - return iet diff --git a/devito/types/petsc.py b/devito/types/petsc.py index a11db983a4..560d920526 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -211,8 +211,8 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): s1 = Symbol(name='s1') # Wrapped rhs in LinearSolveExpr for simplicity in iet_build pass. - mock_action = MockEq(s0, LinearSolveExpr(y_matvec.indexify(indices=indices))) - mock_rhs = MockEq(s1, LinearSolveExpr(b_tmp.indexify(indices=indices))) + mock_action = Eq(s0, Mock(y_matvec.indexify(indices=indices))) + mock_rhs = Eq(s1, Mock(b_tmp.indexify(indices=indices))) return [matvecaction, mock_action] + [rhs, mock_rhs] @@ -265,3 +265,28 @@ def solver_parameters(self): return self._solver_parameters func = Reconstructable._rebuild + + +class Mock(sympy.Function, Reconstructable): + + __rargs__ = ('expr',) + + def __new__(cls, expr, **kwargs): + + obj = super().__new__(cls, expr) + obj._expr = expr + return obj + + def __repr__(self): + return "%s(%s)" % (self.__class__.__name__, self.expr) + + __str__ = __repr__ + + def _sympystr(self, printer): + return str(self) + + @property + def expr(self): + return self._expr + + func = Reconstructable._rebuild From c4ec0b731be575701ef9807bf6e2829c179cc2b7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 11:17:51 +0100 Subject: [PATCH 080/151] misc: Clean up --- devito/ir/stree/algorithms.py | 1 - devito/operator/operator.py | 1 - devito/types/petsc.py | 21 +++++++++------------ 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index c738db645a..a85b93460d 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,7 +163,6 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: - # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index b136a0370d..b83c49ea04 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -424,7 +424,6 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters - # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 560d920526..0d3873f72d 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -172,18 +172,6 @@ class RHSEq(Eq): pass -class MockEq(Eq): - """ - Represents a mock/placeholder equation to ensure distinct iteration loops. - - For example, the mat-vec action iteration loop is to be isolated from the - expression loop used to build the RHS of the linear system. This separation - facilitates the utilisation of the mat-vec iteration loop in callback functions - created at the IET level. - """ - pass - - def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # TODO: Add check for time dimensions and utilise implicit dimensions. @@ -269,6 +257,15 @@ def solver_parameters(self): class Mock(sympy.Function, Reconstructable): + """ + Represents a mock/placeholder RHS to ensure distinct iteration loops. + + For example, the mat-vec action iteration loop is to be isolated from the + expression loop used to build the RHS of the linear system. This separation + facilitates the utilisation of the mat-vec iteration loop in callback functions + created at the IET level. + """ + __rargs__ = ('expr',) def __new__(cls, expr, **kwargs): From ff4cf07e4205f6fd32c70802803c5a6b5715bc9e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 11:30:41 +0100 Subject: [PATCH 081/151] misc: Clean --- devito/types/petsc.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 0d3873f72d..1006b3dbe5 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -197,8 +197,6 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): indices = tuple(d + 1 for d in target.dimensions) s0 = Symbol(name='s0') s1 = Symbol(name='s1') - - # Wrapped rhs in LinearSolveExpr for simplicity in iet_build pass. mock_action = Eq(s0, Mock(y_matvec.indexify(indices=indices))) mock_rhs = Eq(s1, Mock(b_tmp.indexify(indices=indices))) From 309a3f9f3628aa9ec8f37533fcb09b478995fb2a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 12:18:20 +0100 Subject: [PATCH 082/151] misc: Comments --- devito/types/petsc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 1006b3dbe5..50383d485f 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -192,7 +192,7 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) - # Create mock equations to ensure distinct iteration loops for each component + # Create 'mock' dependencies to ensure distinct iteration loops for each component # of the linear solve. indices = tuple(d + 1 for d in target.dimensions) s0 = Symbol(name='s0') @@ -256,7 +256,7 @@ def solver_parameters(self): class Mock(sympy.Function, Reconstructable): """ - Represents a mock/placeholder RHS to ensure distinct iteration loops. + Represents a 'mock' dependency ensure distinct iteration loops. For example, the mat-vec action iteration loop is to be isolated from the expression loop used to build the RHS of the linear system. This separation From 452a21360e67f1c60309d122daf34a0a0d9bb7cc Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 12:20:08 +0100 Subject: [PATCH 083/151] misc: Typo --- devito/types/petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 50383d485f..1933edc052 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -256,7 +256,7 @@ def solver_parameters(self): class Mock(sympy.Function, Reconstructable): """ - Represents a 'mock' dependency ensure distinct iteration loops. + Represents a 'mock' dependency to ensure distinct iteration loops. For example, the mat-vec action iteration loop is to be isolated from the expression loop used to build the RHS of the linear system. This separation From 21d6af2fb8b2b08b2147fb69767931b387956683 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 10:40:15 +0100 Subject: [PATCH 084/151] compiler: Lift iteration space instead of using new Mock object --- devito/ir/clusters/algorithms.py | 4 ++++ devito/types/petsc.py | 11 ++--------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index 5323539dd4..8dc471fdc1 100644 --- a/devito/ir/clusters/algorithms.py +++ b/devito/ir/clusters/algorithms.py @@ -21,6 +21,7 @@ is_integer, split, timed_pass, toposort) from devito.types import Array, Eq, Symbol from devito.types.dimension import BOTTOM, ModuloDimension +from devito.types import LinearSolveExpr __all__ = ['clusterize'] @@ -72,6 +73,9 @@ def impose_total_ordering(clusters): relations = c.ispace.relations ispace = c.ispace.reorder(relations=relations, mode='total') + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = ispace.lift(c.exprs[0].rhs.target.dimensions) + processed.append(c.rebuild(ispace=ispace)) return processed diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 1933edc052..a5526a91c6 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -184,7 +184,7 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): halo=target.halo) for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] - # # TODO: Extend to rearrange equation for implicit time stepping. + # TODO: Extend to rearrange equation for implicit time stepping. matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) @@ -192,15 +192,8 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) - # Create 'mock' dependencies to ensure distinct iteration loops for each component - # of the linear solve. - indices = tuple(d + 1 for d in target.dimensions) - s0 = Symbol(name='s0') - s1 = Symbol(name='s1') - mock_action = Eq(s0, Mock(y_matvec.indexify(indices=indices))) - mock_rhs = Eq(s1, Mock(b_tmp.indexify(indices=indices))) - return [matvecaction, mock_action] + [rhs, mock_rhs] + return [matvecaction] + [rhs] class LinearSolveExpr(sympy.Function, Reconstructable): From 5b4a79c6d5bd1edc30a245a4a6367bb129767e39 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 10:43:37 +0100 Subject: [PATCH 085/151] misc: Clean --- devito/types/petsc.py | 35 ----------------------------------- 1 file changed, 35 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index a5526a91c6..9a52392f4d 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -192,7 +192,6 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) - return [matvecaction] + [rhs] @@ -244,37 +243,3 @@ def solver_parameters(self): return self._solver_parameters func = Reconstructable._rebuild - - -class Mock(sympy.Function, Reconstructable): - - """ - Represents a 'mock' dependency to ensure distinct iteration loops. - - For example, the mat-vec action iteration loop is to be isolated from the - expression loop used to build the RHS of the linear system. This separation - facilitates the utilisation of the mat-vec iteration loop in callback functions - created at the IET level. - """ - - __rargs__ = ('expr',) - - def __new__(cls, expr, **kwargs): - - obj = super().__new__(cls, expr) - obj._expr = expr - return obj - - def __repr__(self): - return "%s(%s)" % (self.__class__.__name__, self.expr) - - __str__ = __repr__ - - def _sympystr(self, printer): - return str(self) - - @property - def expr(self): - return self._expr - - func = Reconstructable._rebuild From a1dc4a440d77507ad97292886afb0c591db460a3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 11:13:02 +0100 Subject: [PATCH 086/151] misc: Clean --- tests/test_petsc.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index ca71343fa5..883fc322d9 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -2,7 +2,7 @@ from devito import Grid, Function, Eq, Operator from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, MatVecAction, FindNodes, RHSLinearSystem, - PointerCast) + PointerCast, retrieve_iteration_tree) from devito.passes.iet.languages.C import CDataManager from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, PETScSolve) @@ -134,6 +134,10 @@ def test_petsc_solve(): assert action_expr[-1].expr.rhs.solver_parameters == \ {'ksp_type': 'gmres', 'pc_type': 'jacobi'} + # Check the matvec action and rhs have distinct iteration loops i.e + # each iteration space was "lifted" properly. + assert len(retrieve_iteration_tree(op)) == 2 + def test_petsc_cast(): """ From 3e574126ef7422aa5d1e5b989979fe804ae66adb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 11:48:07 +0100 Subject: [PATCH 087/151] misc: Clean --- devito/types/petsc.py | 1 - 1 file changed, 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 9a52392f4d..25a3d8b596 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -173,7 +173,6 @@ class RHSEq(Eq): def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - # TODO: Add check for time dimensions and utilise implicit dimensions. y_matvec, x_matvec, b_tmp = [ From c5615731629e1aba1c14ea8fbb3f3f051774dc7e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 12:30:38 +0100 Subject: [PATCH 088/151] compiler: Move petsc lift to petsc only file --- devito/ir/clusters/algorithms.py | 5 +---- devito/types/petsc.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index 8dc471fdc1..abe65e22df 100644 --- a/devito/ir/clusters/algorithms.py +++ b/devito/ir/clusters/algorithms.py @@ -21,7 +21,6 @@ is_integer, split, timed_pass, toposort) from devito.types import Array, Eq, Symbol from devito.types.dimension import BOTTOM, ModuloDimension -from devito.types import LinearSolveExpr __all__ = ['clusterize'] @@ -35,6 +34,7 @@ def clusterize(exprs, **kwargs): # Setup the IterationSpaces based on data dependence analysis clusters = impose_total_ordering(clusters) + clusters = petsc_lift(clusters) clusters = Schedule().process(clusters) # Handle SteppingDimensions @@ -73,9 +73,6 @@ def impose_total_ordering(clusters): relations = c.ispace.relations ispace = c.ispace.reorder(relations=relations, mode='total') - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = ispace.lift(c.exprs[0].rhs.target.dimensions) - processed.append(c.rebuild(ispace=ispace)) return processed diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 25a3d8b596..5cfc43cb16 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -14,6 +14,11 @@ from devito.types.basic import IndexedBase +__all__ = ['DM', 'Mat', 'Vec', 'PetscMPIInt', 'KSP', 'PC', 'KSPConvergedReason', + 'DMDALocalInfo', 'PETScArray', 'MatVecEq', 'RHSEq', 'LinearSolveExpr', + 'PETScSolve', 'petsc_lift'] + + class DM(LocalObject): """ PETSc Data Management object (DM). @@ -242,3 +247,21 @@ def solver_parameters(self): return self._solver_parameters func = Reconstructable._rebuild + + +def petsc_lift(clusters): + """ + Lift the iteration space associated with each PETSc equation. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions. + """ + processed = [] + for c in clusters: + + ispace = c.ispace + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + + processed.append(c.rebuild(ispace=ispace)) + + return processed From 2721817e587bc9958dc5bb28af8bedd97c64fffb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 12:33:49 +0100 Subject: [PATCH 089/151] types: Simplify petsc_lift --- devito/types/petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 5cfc43cb16..2a1f3e4b95 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -260,7 +260,7 @@ def petsc_lift(clusters): ispace = c.ispace if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + ispace = ispace.lift(c.exprs[0].rhs.target.dimensions) processed.append(c.rebuild(ispace=ispace)) From 06a9c4c53c9aa59761ee445646e31e246617cea8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 14:01:28 +0100 Subject: [PATCH 090/151] types: Add else to petsc_lift --- devito/types/petsc.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 2a1f3e4b95..80e5a50682 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -258,10 +258,10 @@ def petsc_lift(clusters): processed = [] for c in clusters: - ispace = c.ispace if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = ispace.lift(c.exprs[0].rhs.target.dimensions) - - processed.append(c.rebuild(ispace=ispace)) + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) return processed From 11954703dfb0f56f3c7128cee256f6359736fd16 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 15:45:44 +0100 Subject: [PATCH 091/151] misc: Move petsc_lift to passes/clusters file --- devito/passes/clusters/petsc.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 devito/passes/clusters/petsc.py diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py new file mode 100644 index 0000000000..4acb74862d --- /dev/null +++ b/devito/passes/clusters/petsc.py @@ -0,0 +1,23 @@ +from devito.tools import timed_pass +from devito.types import LinearSolveExpr + +__all__ = ['petsc_lift'] + + +@timed_pass() +def petsc_lift(clusters): + """ + Lift the iteration space associated with each PETSc equation. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions. + """ + processed = [] + for c in clusters: + + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed From 1c2083379b00f9fbc2c4fcd3534bf2cf647746dd Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 16:02:19 +0100 Subject: [PATCH 092/151] compiler: Move petsc_lift to after clusterize inside lower_clusters --- devito/ir/clusters/algorithms.py | 1 - devito/operator/operator.py | 6 +++++- devito/passes/clusters/__init__.py | 1 + devito/passes/clusters/petsc.py | 6 ++++-- devito/types/petsc.py | 20 +------------------- 5 files changed, 11 insertions(+), 23 deletions(-) diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index abe65e22df..5323539dd4 100644 --- a/devito/ir/clusters/algorithms.py +++ b/devito/ir/clusters/algorithms.py @@ -34,7 +34,6 @@ def clusterize(exprs, **kwargs): # Setup the IterationSpaces based on data dependence analysis clusters = impose_total_ordering(clusters) - clusters = petsc_lift(clusters) clusters = Schedule().process(clusters) # Handle SteppingDimensions diff --git a/devito/operator/operator.py b/devito/operator/operator.py index b83c49ea04..02243a1ccb 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -26,7 +26,7 @@ from devito.parameters import configuration from devito.passes import (Graph, lower_index_derivatives, generate_implicit, generate_macros, minimize_symbols, unevaluate, - error_mapper, is_on_device, lower_petsc) + error_mapper, is_on_device, petsc_lift, lower_petsc) from devito.symbolics import estimate_cost, subs_op_args from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple, flatten, filter_sorted, frozendict, is_integer, @@ -380,6 +380,10 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Build a sequence of Clusters from a sequence of Eqs clusters = clusterize(expressions, **kwargs) + # Lift iteration spaces surrounding PETSc equations to produce + # distinct iteration loops. + clusters = petsc_lift(clusters) + # Operation count before specialization init_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) diff --git a/devito/passes/clusters/__init__.py b/devito/passes/clusters/__init__.py index c41a628e06..64a4784062 100644 --- a/devito/passes/clusters/__init__.py +++ b/devito/passes/clusters/__init__.py @@ -9,3 +9,4 @@ from .misc import * # noqa from .derivatives import * # noqa from .unevaluate import * # noqa +from .petsc import * # noqa diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py index 4acb74862d..133ac5df66 100644 --- a/devito/passes/clusters/petsc.py +++ b/devito/passes/clusters/petsc.py @@ -7,9 +7,11 @@ @timed_pass() def petsc_lift(clusters): """ - Lift the iteration space associated with each PETSc equation. + Lift the iteration space surrounding each PETSc equation to create + distinct iteration loops. This simplifys the movement of the loops + into specific callback functions generated at the IET level. TODO: Potentially only need to lift the PETSc equations required - by the callback functions. + by the callback functions, not the ones that stay inside the main kernel. """ processed = [] for c in clusters: diff --git a/devito/types/petsc.py b/devito/types/petsc.py index 80e5a50682..f56b096950 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -16,7 +16,7 @@ __all__ = ['DM', 'Mat', 'Vec', 'PetscMPIInt', 'KSP', 'PC', 'KSPConvergedReason', 'DMDALocalInfo', 'PETScArray', 'MatVecEq', 'RHSEq', 'LinearSolveExpr', - 'PETScSolve', 'petsc_lift'] + 'PETScSolve'] class DM(LocalObject): @@ -247,21 +247,3 @@ def solver_parameters(self): return self._solver_parameters func = Reconstructable._rebuild - - -def petsc_lift(clusters): - """ - Lift the iteration space associated with each PETSc equation. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions. - """ - processed = [] - for c in clusters: - - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - - return processed From fdd3d67921d60611e97d8d0dd711907127f42e04 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 9 Jul 2024 23:02:26 +0100 Subject: [PATCH 093/151] misc: Fix functools --- devito/types/petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/petsc.py b/devito/types/petsc.py index f56b096950..abd97e2220 100644 --- a/devito/types/petsc.py +++ b/devito/types/petsc.py @@ -5,7 +5,7 @@ from devito.types import LocalObject, Eq from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic -from cached_property import cached_property +from functools import cached_property from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry From 1cce1af564789dc7c42a4b1b1ff00c2a0d52d44a Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 094/151] misc: Split Dockerfile.cpu into .cpu and .intel --- docker/Dockerfile.intel | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index e013fccc03..1fab7147d6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,6 +35,7 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu +<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -50,6 +51,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +======= +# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal +RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg +RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list + +RUN apt-get update -y && apt-get dist-upgrade -y && \ + apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ + intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ + libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -76,7 +87,11 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD # ICX OpenMP image +======= +# ICX image +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -93,6 +108,35 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD +======= +# ICX hpc image +############################################################## +FROM oneapi as icx-hpc + +# Install both icc and icx to avoid missing dependencies +RUN apt-get update -y && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic + +# Missig components +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager +RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ + chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ + rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# Devito config +ENV DEVITO_ARCH="icx" +ENV DEVITO_LANGUAGE="openmp" +# MPICC compiler for mpi4py +ENV MPICC=mpiicc +ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' + +############################################################## +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -107,6 +151,12 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl +<<<<<<< HEAD +======= +# NOTE: the name of this file ends with ".cpu" but this is a GPU image. +# It then feels a bit akward, so some restructuring might be needed + +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 7ad9bc977140b7a8d73b83fcc828d4998945feba Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 095/151] misc: Update file due to incorrect rebase --- .github/workflows/pytest-petsc.yml | 79 ++++ conftest.py | 16 +- devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 2 +- devito/ir/iet/algorithms.py | 16 +- devito/ir/iet/nodes.py | 38 +- devito/ir/stree/algorithms.py | 1 + devito/operator/operator.py | 5 +- devito/passes/clusters/__init__.py | 3 +- devito/passes/iet/__init__.py | 1 - devito/passes/iet/definitions.py | 4 +- devito/passes/iet/engine.py | 1 + devito/passes/iet/petsc.py | 12 - devito/petsc/__init__.py | 1 + .../clusters/petsc.py => petsc/clusters.py} | 6 +- devito/petsc/iet/__init__.py | 1 + devito/petsc/iet/nodes.py | 31 ++ devito/petsc/iet/passes.py | 418 ++++++++++++++++++ devito/petsc/solve.py | 160 +++++++ devito/{types/petsc.py => petsc/types.py} | 175 +++++++- devito/petsc/utils.py | 72 +++ devito/symbolics/printer.py | 3 + devito/types/__init__.py | 1 - docker/Dockerfile.cpu | 3 +- docker/Dockerfile.devito | 23 +- docker/Dockerfile.intel | 50 --- tests/test_petsc.py | 277 +++++++++++- 27 files changed, 1236 insertions(+), 165 deletions(-) create mode 100644 .github/workflows/pytest-petsc.yml delete mode 100644 devito/passes/iet/petsc.py create mode 100644 devito/petsc/__init__.py rename devito/{passes/clusters/petsc.py => petsc/clusters.py} (83%) create mode 100644 devito/petsc/iet/__init__.py create mode 100644 devito/petsc/iet/nodes.py create mode 100644 devito/petsc/iet/passes.py create mode 100644 devito/petsc/solve.py rename devito/{types/petsc.py => petsc/types.py} (53%) create mode 100644 devito/petsc/utils.py diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml new file mode 100644 index 0000000000..08da1220fa --- /dev/null +++ b/.github/workflows/pytest-petsc.yml @@ -0,0 +1,79 @@ +name: CI-petsc + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +on: + # Trigger the workflow on push or pull request, + # but only for the master branch + push: + branches: + - master + # Temporary + - separate-eqn + pull_request: + branches: + - master + # Temporary + - separate-eqn + - centre-stencil + +jobs: + pytest: + name: ${{ matrix.name }}-${{ matrix.set }} + runs-on: "${{ matrix.os }}" + + env: + DOCKER_BUILDKIT: "1" + DEVITO_ARCH: "${{ matrix.arch }}" + DEVITO_LANGUAGE: ${{ matrix.language }} + + strategy: + # Prevent all build to stop if a single one fails + fail-fast: false + + matrix: + name: [ + pytest-docker-py39-gcc-noomp + ] + include: + - name: pytest-docker-py39-gcc-noomp + python-version: '3.9' + os: ubuntu-latest + arch: "gcc" + language: "C" + sympy: "1.12" + + steps: + - name: Checkout devito + uses: actions/checkout@v4 + + - name: Build docker image + run: | + docker build . --file docker/Dockerfile.devito --tag devito_img --build-arg base=zoeleibowitz/bases:cpu-${{ matrix.arch }} --build-arg petscinstall=petsc + + - name: Set run prefix + run: | + echo "RUN_CMD=docker run --rm -t -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} --name testrun devito_img" >> $GITHUB_ENV + id: set-run + + - name: Set tests + run : | + echo "TESTS=tests/test_petsc.py" >> $GITHUB_ENV + id: set-tests + + - name: Check configuration + run: | + ${{ env.RUN_CMD }} python3 -c "from devito import configuration; print(''.join(['%s: %s \n' % (k, v) for (k, v) in configuration.items()]))" + + - name: Test with pytest + run: | + ${{ env.RUN_CMD }} pytest -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + + - name: Upload coverage to Codecov + if: "!contains(matrix.name, 'docker')" + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + name: ${{ matrix.name }} \ No newline at end of file diff --git a/conftest.py b/conftest.py index 4bb0629327..b804aaebc7 100644 --- a/conftest.py +++ b/conftest.py @@ -14,6 +14,7 @@ from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock, retrieve_iteration_tree) from devito.tools import as_tuple +from devito.petsc.utils import get_petsc_dir, get_petsc_arch try: from mpi4py import MPI # noqa @@ -33,7 +34,7 @@ def skipif(items, whole_module=False): accepted = set() accepted.update({'device', 'device-C', 'device-openmp', 'device-openacc', 'device-aomp', 'cpu64-icc', 'cpu64-icx', 'cpu64-nvc', 'cpu64-arm', - 'cpu64-icpx', 'chkpnt'}) + 'cpu64-icpx', 'chkpnt', 'petsc'}) accepted.update({'nodevice'}) unknown = sorted(set(items) - accepted) if unknown: @@ -87,6 +88,19 @@ def skipif(items, whole_module=False): if i == 'chkpnt' and Revolver is NoopRevolver: skipit = "pyrevolve not installed" break + if i == 'petsc': + petsc_dir = get_petsc_dir() + petsc_arch = get_petsc_arch() + if petsc_dir is None or petsc_arch is None: + skipit = "PETSC_DIR or PETSC_ARCH are not set" + break + else: + petsc_installed = os.path.join( + petsc_dir, petsc_arch, 'include', 'petscconf.h' + ) + if not os.path.isfile(petsc_installed): + skipit = "PETSc is not installed" + break if skipit is False: return pytest.mark.skipif(False, reason='') diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index f0b0d718ea..5aa71a0529 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -376,7 +376,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 7d86bcc85c..ff524a1f90 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -244,7 +244,7 @@ def __new__(cls, *args, **kwargs): expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + return expr @property diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 1a2a9c504b..af079f25d8 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -1,11 +1,10 @@ from collections import OrderedDict from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, - Section, HaloSpot, ExpressionBundle, MatVecAction, - RHSLinearSystem) + Section, HaloSpot, ExpressionBundle) from devito.tools import timed_pass -from devito.types import LinearSolveExpr -from devito.ir.equations import OpMatVec, OpRHS +from devito.petsc.types import LinearSolveExpr +from devito.petsc.utils import petsc_iet_mapper __all__ = ['iet_build'] @@ -25,10 +24,11 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: + # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): - exprs.append(linsolve_mapper[e.operation](e, operation=e.operation)) + exprs.append(petsc_iet_mapper[e.operation](e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs) @@ -60,9 +60,3 @@ def iet_build(stree): queues.setdefault(i.parent, []).append(body) assert False - - -# Mapping special Eq operations to their corresponding IET Expression subclass types. -# These operations correspond to subclasses of Eq utilised within PETScSolve. -linsolve_mapper = {OpMatVec: MatVecAction, - OpRHS: RHSLinearSystem} diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 09d56465a8..1739731fdc 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -10,7 +10,7 @@ from sympy import IndexedBase, sympify from devito.data import FULL -from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax, OpMatVec, OpRHS +from devito.ir.equations import DummyEq, OpInc, OpMin, OpMax from devito.ir.support import (INBOUND, SEQUENTIAL, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, VECTORIZED, AFFINE, Property, Forward, WithLock, PrefetchUpdate, detect_io) @@ -28,7 +28,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'Callback', 'MatVecAction', 'RHSLinearSystem'] + 'CallableBody', 'Transfer', 'Callback'] # First-class IET nodes @@ -484,35 +484,6 @@ def __init__(self, expr, pragmas=None): super().__init__(expr, pragmas=pragmas, operation=OpInc) -class LinearSolverExpression(Expression): - - """ - Base class for general expressions required by a - matrix-free linear solve of the form Ax=b. - """ - pass - - -class MatVecAction(LinearSolverExpression): - - """ - Expression representing matrix-vector multiplication. - """ - - def __init__(self, expr, pragmas=None, operation=OpMatVec): - super().__init__(expr, pragmas=pragmas, operation=operation) - - -class RHSLinearSystem(LinearSolverExpression): - - """ - Expression to build the RHS of a linear system. - """ - - def __init__(self, expr, pragmas=None, operation=OpRHS): - super().__init__(expr, pragmas=pragmas, operation=operation) - - class Iteration(Node): """ @@ -1182,8 +1153,9 @@ class Callback(Call): engine fails to bind the callback to a specific Call. Consequently, errors occur during the creation of the call graph. """ - - def __init__(self, name, retval, param_types): + # TODO: Create a common base class for Call and Callback to avoid + # having arguments=None here + def __init__(self, name, retval, param_types, arguments=None): super().__init__(name=name) self.retval = retval self.param_types = as_tuple(param_types) diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index a85b93460d..c738db645a 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,6 +163,7 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: + # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 02243a1ccb..acc5524c5d 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -26,7 +26,7 @@ from devito.parameters import configuration from devito.passes import (Graph, lower_index_derivatives, generate_implicit, generate_macros, minimize_symbols, unevaluate, - error_mapper, is_on_device, petsc_lift, lower_petsc) + error_mapper, is_on_device) from devito.symbolics import estimate_cost, subs_op_args from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple, flatten, filter_sorted, frozendict, is_integer, @@ -34,6 +34,8 @@ from devito.types import (Buffer, Grid, Evaluable, host_layer, device_layer, disk_layer) from devito.types.dimension import Thickness +from devito.petsc.iet.passes import lower_petsc +from devito.petsc.clusters import petsc_preprocess __all__ = ['Operator'] @@ -428,6 +430,7 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters + # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) diff --git a/devito/passes/clusters/__init__.py b/devito/passes/clusters/__init__.py index 64a4784062..1b14ce1aad 100644 --- a/devito/passes/clusters/__init__.py +++ b/devito/passes/clusters/__init__.py @@ -8,5 +8,4 @@ from .implicit import * # noqa from .misc import * # noqa from .derivatives import * # noqa -from .unevaluate import * # noqa -from .petsc import * # noqa +from .unevaluate import * # noqa \ No newline at end of file diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index cf6a35de90..c09db00c9b 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,4 +8,3 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa -from .petsc import * # noqa diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 7f0350fec3..62c24b9c78 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -452,7 +452,9 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [i for i in bases if i.name != i.function._C_name] + bases = [ + i for i in bases + if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index b9b5bf15d4..cf71b7773c 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -129,6 +129,7 @@ def apply(self, func, **kwargs): compiler.add_libraries(as_tuple(metadata.get('libs'))) compiler.add_library_dirs(as_tuple(metadata.get('lib_dirs')), rpath=metadata.get('rpath', False)) + compiler.add_ldflags(as_tuple(metadata.get('ldflags'))) except KeyError: pass diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py deleted file mode 100644 index c050dbbb80..0000000000 --- a/devito/passes/iet/petsc.py +++ /dev/null @@ -1,12 +0,0 @@ -from devito.passes.iet.engine import iet_pass - -__all__ = ['lower_petsc'] - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't - # appear in ccode - - return iet, {} diff --git a/devito/petsc/__init__.py b/devito/petsc/__init__.py new file mode 100644 index 0000000000..2927bc960e --- /dev/null +++ b/devito/petsc/__init__.py @@ -0,0 +1 @@ +from devito.petsc.solve import * # noqa diff --git a/devito/passes/clusters/petsc.py b/devito/petsc/clusters.py similarity index 83% rename from devito/passes/clusters/petsc.py rename to devito/petsc/clusters.py index 133ac5df66..8aca92b035 100644 --- a/devito/passes/clusters/petsc.py +++ b/devito/petsc/clusters.py @@ -1,7 +1,5 @@ from devito.tools import timed_pass -from devito.types import LinearSolveExpr - -__all__ = ['petsc_lift'] +from devito.petsc.types import LinearSolveExpr @timed_pass() @@ -17,7 +15,7 @@ def petsc_lift(clusters): for c in clusters: if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + ispace = c.ispace.lift(c.exprs[0].rhs.target.space_dimensions) processed.append(c.rebuild(ispace=ispace)) else: processed.append(c) diff --git a/devito/petsc/iet/__init__.py b/devito/petsc/iet/__init__.py new file mode 100644 index 0000000000..7ce0b8c3f1 --- /dev/null +++ b/devito/petsc/iet/__init__.py @@ -0,0 +1 @@ +from devito.petsc.iet import * # noqa \ No newline at end of file diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py new file mode 100644 index 0000000000..4682ea91d2 --- /dev/null +++ b/devito/petsc/iet/nodes.py @@ -0,0 +1,31 @@ +from devito.ir.iet import Expression +from devito.ir.equations import OpMatVec, OpRHS + + +class LinearSolverExpression(Expression): + + """ + Base class for general expressions required by a + matrix-free linear solve of the form Ax=b. + """ + pass + + +class MatVecAction(LinearSolverExpression): + + """ + Expression representing matrix-vector multiplication. + """ + + def __init__(self, expr, pragmas=None, operation=OpMatVec): + super().__init__(expr, pragmas=pragmas, operation=operation) + + +class RHSLinearSystem(LinearSolverExpression): + + """ + Expression to build the RHS of a linear system. + """ + + def __init__(self, expr, pragmas=None, operation=OpRHS): + super().__init__(expr, pragmas=pragmas, operation=operation) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py new file mode 100644 index 0000000000..3bc1e9a29b --- /dev/null +++ b/devito/petsc/iet/passes.py @@ -0,0 +1,418 @@ +from collections import OrderedDict +import cgen as c + +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, + Transformer, FindSymbols, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition, BlankLine, PointerCast, filter_iterations, + retrieve_iteration_tree) +from devito.symbolics import Byref, Macro, FieldFromPointer +from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, + Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) +from devito.petsc.iet.nodes import MatVecAction, LinearSolverExpression +from devito.petsc.utils import (solver_mapper, petsc_call, petsc_call_mpi, + core_metadata) + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # Check if PETScSolve was used + petsc_nodes = FindNodes(MatVecAction).visit(iet) + + if not petsc_nodes: + return iet, {} + + # Collect all petsc solution fields + unique_targets = list({i.expr.rhs.target for i in petsc_nodes}) + init = init_petsc(**kwargs) + + # Assumption is that all targets have the same grid so can use any target here + objs = build_core_objects(unique_targets[-1], **kwargs) + objs['struct'] = build_petsc_struct(iet) + + # Create core PETSc calls (not specific to each PETScSolve) + core = make_core_petsc_calls(objs, **kwargs) + + # Create matvec mapper from the spatial iteration loops (exclude time loop if present) + spatial_body = [] + for tree in retrieve_iteration_tree(iet): + root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] + spatial_body.append(root) + matvec_mapper = MapNodes(Iteration, MatVecAction, + 'groupby').visit(List(body=spatial_body)) + + setup = [] + + # Create a different DMDA for each target with a unique space order + unique_dmdas = create_dmda_objs(unique_targets) + objs.update(unique_dmdas) + for dmda in unique_dmdas.values(): + setup.extend(create_dmda_calls(dmda, objs)) + + subs = {} + efuncs = OrderedDict() + + # Create the PETSc calls which are specific to each target + for target in unique_targets: + solver_objs = build_solver_objs(target) + + # Generate solver setup for target + for iter, (matvec,) in matvec_mapper.items(): + # Skip the MatVecAction if it is not associated with the target + # There will most likely be more than one MatVecAction + # associated with the target e.g interior matvec + BC matvecs + if matvec.expr.rhs.target != target: + continue + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + break + + # Create the body of the matrix-vector callback for target + matvec_body_list = [] + for iter, (matvec,) in matvec_mapper.items(): + if matvec.expr.rhs.target != target: + continue + matvec_body_list.append(iter[0]) + # Remove the iteration loop from the main kernel encapsulating + # the matvec equations since they are moved into the callback + subs.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, List(body=matvec_body_list), solver_objs, objs + ) + + setup.extend([matvec_op, BlankLine]) + efuncs[matvec_callback.name] = matvec_callback + + # Remove the LinSolveExpr's from iet and efuncs + subs.update(rebuild_expr_mapper(iet)) + iet = Transformer(subs).visit(iet) + efuncs = transform_efuncs(efuncs, objs['struct']) + + body = iet.body._rebuild(init=init, body=core+tuple(setup)+iet.body.body) + iet = iet._rebuild(body=body) + + metadata = core_metadata() + metadata.update({'efuncs': tuple(efuncs.values())}) + + return iet, metadata + + +def init_petsc(**kwargs): + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve + # TODO: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options + initialize = petsc_call('PetscInitialize', [Null, Null, Null, Null]) + + return petsc_func_begin_user, initialize + + +def build_petsc_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct + usr_ctx = [] + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def make_core_petsc_calls(objs, **kwargs): + call_mpi = petsc_call_mpi('MPI_Comm_size', [objs['comm'], Byref(objs['size'])]) + + return call_mpi, BlankLine + + +def build_core_objects(target, **kwargs): + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return { + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'err': PetscErrorCode(name='err'), + 'grid': target.grid + } + + +def create_dmda_objs(unique_targets): + unique_dmdas = {} + for target in unique_targets: + name = 'da_so_%s' % target.space_order + unique_dmdas[name] = DM(name=name, liveness='eager', + stencil_width=target.space_order) + return unique_dmdas + + +def create_dmda_calls(dmda, objs): + dmda_create = create_dmda(dmda, objs) + dm_setup = petsc_call('DMSetUp', [dmda]) + dm_app_ctx = petsc_call('DMSetApplicationContext', [dmda, objs['struct']]) + dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) + + return dmda_create, dm_setup, dm_app_ctx, dm_mat_type, BlankLine + + +def create_dmda(dmda, objs): + no_of_space_dims = len(objs['grid'].dimensions) + + # MPI communicator + args = [objs['comm']] + + # Type of ghost nodes + args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(no_of_space_dims)]) + + # Stencil type + if no_of_space_dims > 1: + args.append('DMDA_STENCIL_BOX') + + # Global dimensions + args.extend(list(objs['grid'].shape)[::-1]) + # No.of processors in each dimension + if no_of_space_dims > 1: + args.extend(list(objs['grid'].distributor.topology)[::-1]) + + # Number of degrees of freedom per node + args.append(1) + # "Stencil width" -> size of overlap + args.append(dmda.stencil_width) + args.extend([Null for _ in range(no_of_space_dims)]) + + # The distributed array object + args.append(Byref(dmda)) + + # The PETSc call used to create the DMDA + dmda = petsc_call('DMDACreate%sd' % no_of_space_dims, args) + + return dmda + + +def build_solver_objs(target): + name = target.name + return { + 'Jac': Mat(name='J_%s' % name), + 'x_global': Vec(name='x_global_%s' % name), + 'x_local': Vec(name='x_local_%s' % name, liveness='eager'), + 'b_global': Vec(name='b_global_%s' % name), + 'b_local': Vec(name='b_local_%s' % name, liveness='eager'), + 'ksp': KSP(name='ksp_%s' % name), + 'pc': PC(name='pc_%s' % name), + 'snes': SNES(name='snes_%s' % name), + 'X_global': Vec(name='X_global_%s' % name), + 'Y_global': Vec(name='Y_global_%s' % name), + 'X_local': Vec(name='X_local_%s' % name, liveness='eager'), + 'Y_local': Vec(name='Y_local_%s' % name, liveness='eager') + } + + +def generate_solver_calls(solver_objs, objs, matvec, target): + dmda = objs['da_so_%s' % target.space_order] + + solver_params = matvec.expr.rhs.solver_parameters + + snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) + + snes_set_dm = petsc_call('SNESSetDM', [solver_objs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(solver_objs['Jac'])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) + + global_x = petsc_call('DMCreateGlobalVector', + [dmda, Byref(solver_objs['x_global'])]) + + local_x = petsc_call('DMCreateLocalVector', + [dmda, Byref(solver_objs['x_local'])]) + + global_b = petsc_call('DMCreateGlobalVector', + [dmda, Byref(solver_objs['b_global'])]) + + local_b = petsc_call('DMCreateLocalVector', + [dmda, Byref(solver_objs['b_local'])]) + + snes_get_ksp = petsc_call('SNESGetKSP', + [solver_objs['snes'], Byref(solver_objs['ksp'])]) + + vec_replace_array = petsc_call( + 'VecReplaceArray', [solver_objs['x_local'], + FieldFromPointer(target._C_field_data, target._C_symbol)] + ) + + ksp_set_tols = petsc_call( + 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], + solver_params['ksp_atol'], solver_params['ksp_divtol'], + solver_params['ksp_max_it']] + ) + + ksp_set_type = petsc_call( + 'KSPSetType', [solver_objs['ksp'], solver_mapper[solver_params['ksp_type']]] + ) + + ksp_get_pc = petsc_call('KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])]) + + pc_set_type = petsc_call( + 'PCSetType', [solver_objs['pc'], solver_mapper[solver_params['pc_type']]] + ) + + ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) + + return ( + snes_create, + snes_set_dm, + create_matrix, + snes_set_type, + global_x, + local_x, + global_b, + local_b, + snes_get_ksp, + vec_replace_array, + ksp_set_tols, + ksp_set_type, + ksp_get_pc, + pc_set_type, + ksp_set_from_ops + ) + + +def create_matvec_callback(target, body, solver_objs, objs): + dmda = objs['da_so_%s' % target.space_order] + + # There will be 2 PETScArrays within the body + petsc_arrays = [i for i in FindSymbols('indexedbases').visit(body) + if isinstance(i.function, PETScArray)] + + # There will only be one PETScArray that is written to within this body and + # one PETScArray which corresponds to the 'seed' vector + petsc_arr_write, = FindSymbols('writes').visit(body) + petsc_arr_seed, = [i.function for i in petsc_arrays + if i.function != petsc_arr_write.function] + + define_arrays = [Definition(i.function) for i in petsc_arrays] + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + define_struct = Definition(objs['struct']) + + mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(objs['struct']._C_symbol)] + ) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], + 'INSERT_VALUES', solver_objs['X_local']] + ) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + ) + + vec_get_array_y = petsc_call( + 'VecGetArray', [solver_objs['Y_local'], Byref(petsc_arr_write._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [solver_objs['X_local'], Byref(petsc_arr_seed._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(petsc_arr_seed.function.dmda_info)] + ) + + casts = [PointerCast(i.function) for i in petsc_arrays] + + vec_restore_array_y = petsc_call( + 'VecRestoreArray', [solver_objs['Y_local'], Byref(petsc_arr_write._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [solver_objs['X_local'], Byref(petsc_arr_seed._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + ]) + + func_return = Call('PetscFunctionReturn', arguments=[0]) + + matvec_body = List(body=[ + petsc_func_begin_user, + define_arrays, + define_struct, + mat_get_dm, + dm_get_app_context, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + vec_get_array_y, + vec_get_array_x, + dm_get_local_info, + casts, + BlankLine, + body, + vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end, + func_return]) + + matvec_callback = Callable( + 'MyMatShellMult_%s' % target.name, matvec_body, retval=objs['err'], + parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global']) + ) + + matvec_operation = petsc_call( + 'MatShellSetOperation', [solver_objs['Jac'], 'MATOP_MULT', + Callback(matvec_callback.name, void, void)] + ) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(callable): + # This mapper removes LinSolveExpr instances from the callable + # These expressions were previously used in lower_petc to carry metadata, + # such as solver_parameters + nodes = FindNodes(LinearSolverExpression).visit(callable) + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for expr in nodes} + + +def transform_efuncs(efuncs, struct): + subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + for efunc in efuncs.values(): + transformed_efunc = Transformer(rebuild_expr_mapper(efunc)).visit(efunc) + transformed_efunc = Uxreplace(subs).visit(transformed_efunc) + efuncs[efunc.name] = transformed_efunc + return efuncs + + +Null = Macro('NULL') +void = 'void' + +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py new file mode 100644 index 0000000000..5bdad59b43 --- /dev/null +++ b/devito/petsc/solve.py @@ -0,0 +1,160 @@ +from functools import singledispatch + +from sympy import simplify + +from devito.finite_differences.differentiable import Add, Mul, EvalDerivative, diffify +from devito.finite_differences.derivative import Derivative +from devito.types import Eq +from devito.operations.solve import eval_time_derivatives +from devito.symbolics import retrieve_functions +from devito.symbolics import uxreplace +from devito.petsc.types import PETScArray, LinearSolveExpr, MatVecEq, RHSEq + +__all__ = ['PETScSolve'] + + +def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): + # TODO: Add check for time dimensions and utilise implicit dimensions. + + is_time_dep = any(dim.is_Time for dim in target.dimensions) + # TODO: Current assumption is rhs is part of pde that remains + # constant at each timestep. Need to insert function to extract this from eq. + y_matvec, x_matvec, b_tmp = [ + PETScArray(name=f'{prefix}_{target.name}', + dtype=target.dtype, + dimensions=target.space_dimensions, + shape=target.grid.shape, + liveness='eager', + halo=target.halo[1:] if is_time_dep else target.halo) + for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] + + b, F_target = separate_eqn(eq, target) + + # Args were updated so need to update target to enable uxreplace on F_target + new_target = {f for f in retrieve_functions(F_target) if + f.function == target.function} + assert len(new_target) == 1 # Sanity check: only one target expected + new_target = new_target.pop() + + # TODO: Current assumption is that problem is linear and user has not provided + # a jacobian. Hence, we can use F_target to form the jac-vec product + matvecaction = MatVecEq( + y_matvec, LinearSolveExpr(uxreplace(F_target, {new_target: x_matvec}), + target=target, solver_parameters=solver_parameters), + subdomain=eq.subdomain) + + # Part of pde that remains constant at each timestep + rhs = RHSEq(b_tmp, LinearSolveExpr(b, target=target, + solver_parameters=solver_parameters), subdomain=eq.subdomain) + + if not bcs: + return [matvecaction, rhs] + + bcs_for_matvec = [] + for bc in bcs: + # TODO: Insert code to distiguish between essential and natural + # boundary conditions since these are treated differently within + # the solver + # NOTE: May eventually remove the essential bcs from the solve + # (and move to rhs) but for now, they are included since this + # is not trivial to implement when using DMDA + # NOTE: Below is temporary -> Just using this as a palceholder for + # the actual BC implementation for the matvec callback + new_rhs = bc.rhs.subs(target, x_matvec) + bc_rhs = LinearSolveExpr( + new_rhs, target=target, solver_parameters=solver_parameters + ) + bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) + + return [matvecaction] + bcs_for_matvec + [rhs] + + +def separate_eqn(eqn, target): + """ + Separate the equation into two separate expressions, + where F(target) = b. + """ + zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) + tmp = eval_time_derivatives(zeroed_eqn.lhs) + b = remove_target(tmp, target) + F_target = diffify(simplify(tmp - b)) + return -b, F_target + + +@singledispatch +def remove_target(expr, target): + return 0 if expr == target else expr + + +@remove_target.register(Add) +@remove_target.register(EvalDerivative) +def _(expr, target): + if not expr.has(target): + return expr + + args = [remove_target(a, target) for a in expr.args] + return expr.func(*args, evaluate=False) + + +@remove_target.register(Mul) +def _(expr, target): + if not expr.has(target): + return expr + + args = [] + for a in expr.args: + if not a.has(target): + args.append(a) + else: + a1 = remove_target(a, target) + args.append(a1) + + return expr.func(*args, evaluate=False) + + +@remove_target.register(Derivative) +def _(expr, target): + return 0 if expr.has(target) else expr + + +@singledispatch +def centre_stencil(expr, target): + """ + Extract the centre stencil from an expression. Its coefficient is what + would appear on the diagonal of the matrix system if the matrix were + formed explicitly. + """ + return expr if expr == target else 0 + + +@centre_stencil.register(Add) +@centre_stencil.register(EvalDerivative) +def _(expr, target): + if not expr.has(target): + return 0 + + args = [centre_stencil(a, target) for a in expr.args] + return expr.func(*args, evaluate=False) + + +@centre_stencil.register(Mul) +def _(expr, target): + if not expr.has(target): + return 0 + + args = [] + for a in expr.args: + if not a.has(target): + args.append(a) + else: + args.append(centre_stencil(a, target)) + + return expr.func(*args, evaluate=False) + + +@centre_stencil.register(Derivative) +def _(expr, target): + if not expr.has(target): + return 0 + args = [centre_stencil(a, target) for a in expr.evaluate.args] + return expr.evaluate.func(*args) diff --git a/devito/types/petsc.py b/devito/petsc/types.py similarity index 53% rename from devito/types/petsc.py rename to devito/petsc/types.py index abd97e2220..62c69f666a 100644 --- a/devito/types/petsc.py +++ b/devito/petsc/types.py @@ -1,30 +1,33 @@ import sympy +from functools import cached_property import numpy as np from devito.tools import CustomDtype -from devito.types import LocalObject, Eq +from devito.types import LocalObject, Eq, CompositeObject from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic -from functools import cached_property from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction +from devito.types.basic import AbstractFunction, Symbol from devito.finite_differences.tools import fd_weights_registry -from devito.tools import Reconstructable +from devito.tools import Reconstructable, dtype_to_ctype from devito.symbolics import FieldFromComposite from devito.types.basic import IndexedBase -__all__ = ['DM', 'Mat', 'Vec', 'PetscMPIInt', 'KSP', 'PC', 'KSPConvergedReason', - 'DMDALocalInfo', 'PETScArray', 'MatVecEq', 'RHSEq', 'LinearSolveExpr', - 'PETScSolve'] - - class DM(LocalObject): """ PETSc Data Management object (DM). """ dtype = CustomDtype('DM') + def __init__(self, *args, stencil_width=None, **kwargs): + super().__init__(*args, **kwargs) + self._stencil_width = stencil_width + + @property + def stencil_width(self): + return self._stencil_width + class Mat(LocalObject): """ @@ -56,6 +59,13 @@ class KSP(LocalObject): dtype = CustomDtype('KSP') +class SNES(LocalObject): + """ + PETSc SNES : Non-Linear Systems Solvers. + """ + dtype = CustomDtype('SNES') + + class PC(LocalObject): """ PETSc object that manages all preconditioners (PC). @@ -79,6 +89,14 @@ class DMDALocalInfo(LocalObject): dtype = CustomDtype('DMDALocalInfo') +class PetscErrorCode(LocalObject): + """ + PETSc datatype used to return PETSc error codes. + https://petsc.org/release/manualpages/Sys/PetscErrorCode/ + """ + dtype = CustomDtype('PetscErrorCode') + + class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent @@ -96,7 +114,7 @@ class PETScArray(ArrayBasic, Differentiable): _default_fd = 'taylor' __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'liveness', 'coefficients')) + ('dimensions', 'shape', 'liveness', 'coefficients')) def __init_finalize__(self, *args, **kwargs): @@ -107,6 +125,7 @@ def __init_finalize__(self, *args, **kwargs): if self._coefficients not in fd_weights_registry: raise ValueError("coefficients must be one of %s" " not %s" % (str(fd_weights_registry), self._coefficients)) + self._shape = kwargs.get('shape') @classmethod def __dtype_setup__(cls, **kwargs): @@ -117,6 +136,40 @@ def coefficients(self): """Form of the coefficients of the function.""" return self._coefficients + @property + def shape(self): + return self._shape + + @cached_property + def _shape_with_inhalo(self): + """ + Shape of the domain+inhalo region. The inhalo region comprises the + outhalo as well as any additional "ghost" layers for MPI halo + exchanges. Data in the inhalo region are exchanged when running + Operators to maintain consistent values as in sequential runs. + + Notes + ----- + Typically, this property won't be used in user code, but it may come + in handy for testing or debugging + """ + return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) + + @cached_property + def shape_allocated(self): + """ + Shape of the allocated data of the Function type object from which + this PETScArray was derived. It includes the domain and inhalo regions, + as well as any additional padding surrounding the halo. + + Notes + ----- + In an MPI context, this is the *local* with_halo region shape. + """ + return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, + self._padding)], + getters=self.dimensions) + @cached_property def _C_ctype(self): return CustomDtype(self.petsc_type, modifier='*') @@ -131,15 +184,18 @@ def dtype(self): @property def symbolic_shape(self): - info = DMDALocalInfo(name='info') - # To access the local grid info via the DM in - # PETSc you use DMDAGetLocalInfo(da, &info) - # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz field_from_composites = [ - FieldFromComposite('g%sm' % d.name, info) for d in self.dimensions] + FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] # Reverse it since DMDA is setup backwards to Devito dimensions. return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + @cached_property + def dmda_info(self): + # To access the local grid info via the DM in + # PETSc you use DMDAGetLocalInfo(da, &info) + # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz + return DMDALocalInfo(name='info', liveness='eager') + @cached_property def indexed(self): return PETScIndexedData(self.name, shape=self._shape, function=self.function) @@ -177,15 +233,31 @@ class RHSEq(Eq): pass +class MockEq(Eq): + """ + Represents a mock/placeholder equation to ensure distinct iteration loops. + + For example, the mat-vec action iteration loop is to be isolated from the + expression loop used to build the RHS of the linear system. This separation + facilitates the utilisation of the mat-vec iteration loop in callback functions + created at the IET level. + """ + pass + + def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # TODO: Add check for time dimensions and utilise implicit dimensions. + is_time_dep = any(dim.is_Time for dim in target.dimensions) + # TODO: Current assumption is rhs is part of pde that remains + # constant at each timestep. Need to insert function to extract this from eq. y_matvec, x_matvec, b_tmp = [ PETScArray(name=f'{prefix}_{target.name}', dtype=target.dtype, - dimensions=target.dimensions, - shape=target.shape, liveness='eager', - halo=target.halo) + dimensions=target.space_dimensions, + shape=target.grid.shape, + liveness='eager', + halo=target.halo[1:] if is_time_dep else target.halo) for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] # TODO: Extend to rearrange equation for implicit time stepping. @@ -193,10 +265,30 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) + # Part of pde that remains constant at each timestep rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, solver_parameters=solver_parameters), subdomain=eq.subdomain) - return [matvecaction] + [rhs] + if not bcs: + return [matvecaction, rhs] + + bcs_for_matvec = [] + for bc in bcs: + # TODO: Insert code to distiguish between essential and natural + # boundary conditions since these are treated differently within + # the solver + # NOTE: May eventually remove the essential bcs from the solve + # (and move to rhs) but for now, they are included since this + # is not trivial to implement when using DMDA + # NOTE: Below is temporary -> Just using this as a palceholder for + # the actual BC implementation for the matvec callback + new_rhs = bc.rhs.subs(target, x_matvec) + bc_rhs = LinearSolveExpr( + new_rhs, target=target, solver_parameters=solver_parameters + ) + bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) + + return [matvecaction] + bcs_for_matvec + [rhs] class LinearSolveExpr(sympy.Function, Reconstructable): @@ -206,7 +298,11 @@ class LinearSolveExpr(sympy.Function, Reconstructable): defaults = { 'ksp_type': 'gmres', - 'pc_type': 'jacobi' + 'pc_type': 'jacobi', + 'ksp_rtol': 1e-7, # Relative tolerance + 'ksp_atol': 1e-50, # Absolute tolerance + 'ksp_divtol': 1e4, # Divergence tolerance + 'ksp_max_it': 10000 # Maximum iterations } def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): @@ -247,3 +343,42 @@ def solver_parameters(self): return self._solver_parameters func = Reconstructable._rebuild + + +def petsc_lift(clusters): + """ + Lift the iteration space associated with each PETSc equation. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions. + """ + processed = [] + for c in clusters: + + ispace = c.ispace + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + + processed.append(c.rebuild(ispace=ispace)) + + return processed + + +class PETScStruct(CompositeObject): + + __rargs__ = ('name', 'usr_ctx') + + def __init__(self, name, usr_ctx): + pfields = [(i._C_name, dtype_to_ctype(i.dtype)) + for i in usr_ctx if isinstance(i, Symbol)] + self._usr_ctx = usr_ctx + super().__init__(name, 'MatContext', pfields) + + @property + def usr_ctx(self): + return self._usr_ctx + + def _arg_values(self, **kwargs): + values = super()._arg_values(**kwargs) + for i in self.fields: + setattr(values[self.name]._obj, i, kwargs['args'][i]) + return values diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py new file mode 100644 index 0000000000..1d171aad2b --- /dev/null +++ b/devito/petsc/utils.py @@ -0,0 +1,72 @@ +import os + +from devito.ir.equations import OpMatVec, OpRHS +from devito.tools import memoized_func +from devito.ir.iet import Call +from devito.petsc.iet.nodes import MatVecAction, RHSLinearSystem + +# Mapping special Eq operations to their corresponding IET Expression subclass types. +# These operations correspond to subclasses of Eq utilised within PETScSolve. +petsc_iet_mapper = {OpMatVec: MatVecAction, + OpRHS: RHSLinearSystem} + + +solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} + + +@memoized_func +def get_petsc_dir(): + # *** First try: via commonly used environment variables + for i in ['PETSC_DIR']: + petsc_dir = os.environ.get(i) + if petsc_dir: + return petsc_dir + + # TODO: Raise error if PETSC_DIR is not set + return None + + +@memoized_func +def get_petsc_arch(): + # *** First try: via commonly used environment variables + for i in ['PETSC_ARCH']: + petsc_arch = os.environ.get(i) + if petsc_arch: + return petsc_arch + # TODO: Raise error if PETSC_ARCH is not set + return None + + +def core_metadata(): + petsc_dir = get_petsc_dir() + petsc_arch = get_petsc_arch() + + # Include directories + global_include = os.path.join(petsc_dir, 'include') + config_specific_include = os.path.join(petsc_dir, f'{petsc_arch}', 'include') + include_dirs = (global_include, config_specific_include) + + # Lib directories + lib_dir = os.path.join(petsc_dir, f'{petsc_arch}', 'lib') + + return { + 'includes': ('petscksp.h', 'petscsnes.h', 'petscdmda.h'), + 'include_dirs': include_dirs, + 'libs': ('petsc'), + 'lib_dirs': lib_dir, + 'ldflags': ('-Wl,-rpath,%s' % lib_dir) + } + + +def petsc_call(specific_call, call_args): + general_call = 'PetscCall' + return Call(general_call, [Call(specific_call, arguments=call_args)]) + + +def petsc_call_mpi(specific_call, call_args): + general_call = 'PetscCallMPI' + return Call(general_call, [Call(specific_call, arguments=call_args)]) diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index 2c366a389e..e7b395d436 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,6 +242,9 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) + def _print_PETScRHS(self, expr): + return "%s" % self._print(expr._expr) + _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/devito/types/__init__.py b/devito/types/__init__.py index 3a5e211e63..6ec8bdfd16 100644 --- a/devito/types/__init__.py +++ b/devito/types/__init__.py @@ -15,7 +15,6 @@ # Some more internal types which depend on some of the types above from .parallel import * # noqa -from .petsc import * # noqa # Needed only outside Devito from .grid import * # noqa diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 8f036badb9..9e51eb9814 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -17,7 +17,8 @@ RUN apt-get update && \ # Install for basic base not containing it RUN apt-get install -y wget flex libnuma-dev hwloc curl cmake git \ - autoconf libtool build-essential procps software-properties-common + autoconf libtool build-essential procps software-properties-common \ + gfortran pkgconf libopenblas-serial-dev # Install gcc RUN if [ -n "$gcc" ]; then \ diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index f167c36e54..ae70971c10 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -4,8 +4,26 @@ # Base image with compilers ARG base=devitocodes/bases:cpu-gcc +ARG petscinstall="" -FROM $base AS builder +FROM $base AS copybase + +################## Install PETSc ############################################ +FROM copybase AS petsccopybase + +RUN apt-get update && apt-get install -y git && \ + python3 -m venv /venv && \ + /venv/bin/pip install --no-cache-dir --upgrade pip && \ + /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ + mkdir -p /opt/petsc && \ + cd /opt/petsc && \ + git clone -b release https://gitlab.com/petsc/petsc.git petsc && \ + cd petsc && \ + ./configure --with-fortran-bindings=0 --with-mpi-dir=/opt/openmpi --with-openblas-include=$(pkg-config --variable=includedir openblas) --with-openblas-lib=$(pkg-config --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build && \ + make all + +ARG petscinstall="" +FROM ${petscinstall}copybase AS builder # User/Group Ids ARG USER_ID=1000 @@ -69,6 +87,9 @@ ARG GROUP_ID=1000 ENV HOME=/app ENV APP_HOME=/app +ENV PETSC_ARCH="devito_build" +ENV PETSC_DIR="/opt/petsc/petsc" + # Create the home directory for the new app user. # Create an app user so our program doesn't run as root. # Chown all the files to the app user. diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index 1fab7147d6..e013fccc03 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,7 +35,6 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu -<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -51,16 +50,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev -======= -# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal -RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg -RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list - -RUN apt-get update -y && apt-get dist-upgrade -y && \ - apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ - intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ - libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -87,11 +76,7 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD # ICX OpenMP image -======= -# ICX image ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -108,35 +93,6 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD -======= -# ICX hpc image -############################################################## -FROM oneapi as icx-hpc - -# Install both icc and icx to avoid missing dependencies -RUN apt-get update -y && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic - -# Missig components -# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager -RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ - chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ - rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -# Devito config -ENV DEVITO_ARCH="icx" -ENV DEVITO_LANGUAGE="openmp" -# MPICC compiler for mpi4py -ENV MPICC=mpiicc -ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' - -############################################################## ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -151,12 +107,6 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl -<<<<<<< HEAD -======= -# NOTE: the name of this file ends with ".cpu" but this is a GPU image. -# It then feels a bit akward, so some restructuring might be needed - ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 883fc322d9..9ce8b209e6 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,18 +1,26 @@ import numpy as np -from devito import Grid, Function, Eq, Operator +import os +import pytest + +from conftest import skipif +from devito import Grid, Function, TimeFunction, Eq, Operator, switchconfig from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, - MatVecAction, FindNodes, RHSLinearSystem, + FindNodes, PointerCast, retrieve_iteration_tree) from devito.passes.iet.languages.C import CDataManager -from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP, - PC, KSPConvergedReason, PETScArray, PETScSolve) +from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, + PC, KSPConvergedReason, PETScArray, + LinearSolveExpr, PETScStruct) +from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil +from devito.petsc.iet.nodes import MatVecAction, RHSLinearSystem +@skipif('petsc') def test_petsc_local_object(): """ Test C++ support for PETSc LocalObjects. """ - lo0 = DM('da') + lo0 = DM('da', stencil_width=1) lo1 = Mat('A') lo2 = Vec('x') lo3 = PetscMPIInt('size') @@ -35,6 +43,7 @@ def test_petsc_local_object(): assert 'KSPConvergedReason reason;' in str(iet) +@skipif('petsc') def test_petsc_functions(): """ Test C++ support for PETScArrays. @@ -67,6 +76,7 @@ def test_petsc_functions(): assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' +@skipif('petsc') def test_petsc_subs(): """ Test support for PETScArrays in substitutions. @@ -92,6 +102,7 @@ def test_petsc_subs(): ' + arr(x, y - h_y)/h_y**2 + arr(x, y + h_y)/h_y**2' +@skipif('petsc') def test_petsc_solve(): """ Test PETScSolve. @@ -105,18 +116,25 @@ def test_petsc_solve(): petsc = PETScSolve(eqn, f) - op = Operator(petsc, opt='noop') + with switchconfig(openmp=False): + op = Operator(petsc, opt='noop') + + callable_roots = [meta_call.root for meta_call in op._func_table.values()] - action_expr = FindNodes(MatVecAction).visit(op) + matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_f'] + action_expr = FindNodes(MatVecAction).visit(matvec_callback[0]) rhs_expr = FindNodes(RHSLinearSystem).visit(op) - assert str(action_expr[-1].expr.rhs.args[0]) == \ - 'x_matvec_f[x + 1, y + 2]/h_x**2 - 2.0*x_matvec_f[x + 2, y + 2]/h_x**2' + \ - ' + x_matvec_f[x + 3, y + 2]/h_x**2 + x_matvec_f[x + 2, y + 1]/h_y**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/h_y**2 + x_matvec_f[x + 2, y + 3]/h_y**2' + assert str(action_expr[-1].expr.rhs) == \ + 'ctx->h_x**(-2)*x_matvec_f[x + 1, y + 2]' + \ + ' - 2.0*ctx->h_x**(-2)*x_matvec_f[x + 2, y + 2]' + \ + ' + ctx->h_x**(-2)*x_matvec_f[x + 3, y + 2]' + \ + ' + ctx->h_y**(-2)*x_matvec_f[x + 2, y + 1]' + \ + ' - 2.0*ctx->h_y**(-2)*x_matvec_f[x + 2, y + 2]' + \ + ' + ctx->h_y**(-2)*x_matvec_f[x + 2, y + 3]' - assert str(rhs_expr[-1].expr.rhs.args[0]) == 'g[x + 2, y + 2]' + assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' # Check the iteration bounds are correct. assert op.arguments().get('x_m') == 0 @@ -124,21 +142,46 @@ def test_petsc_solve(): assert op.arguments().get('y_M') == 1 assert op.arguments().get('x_M') == 1 - # Check the target - assert rhs_expr[-1].expr.rhs.target == f - assert action_expr[-1].expr.rhs.target == f + assert len(retrieve_iteration_tree(op)) == 1 - # Check the solver parameters - assert rhs_expr[-1].expr.rhs.solver_parameters == \ - {'ksp_type': 'gmres', 'pc_type': 'jacobi'} - assert action_expr[-1].expr.rhs.solver_parameters == \ - {'ksp_type': 'gmres', 'pc_type': 'jacobi'} + # TODO: Remove pragmas from PETSc callback functions + assert len(matvec_callback[0].parameters) == 3 + + +@skipif('petsc') +def test_multiple_petsc_solves(): + """ + Test multiple PETScSolves. + """ + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=2) + g1 = Function(name='g1', grid=grid, space_order=2) + + f2 = Function(name='f2', grid=grid, space_order=2) + g2 = Function(name='g2', grid=grid, space_order=2) + + eqn1 = Eq(f1.laplace, g1) + eqn2 = Eq(f2.laplace, g2) + + petsc1 = PETScSolve(eqn1, f1) + petsc2 = PETScSolve(eqn2, f2) + + with switchconfig(openmp=False): + op = Operator(petsc1+petsc2, opt='noop') + + callable_roots = [meta_call.root for meta_call in op._func_table.values()] - # Check the matvec action and rhs have distinct iteration loops i.e - # each iteration space was "lifted" properly. - assert len(retrieve_iteration_tree(op)) == 2 + assert len(callable_roots) == 2 + structs = [i for i in op.parameters if isinstance(i, PETScStruct)] + # Only create 1 struct per Grid/DMDA + assert len(structs) == 1 + assert len(structs[0].fields) == 6 + + +@skipif('petsc') def test_petsc_cast(): """ Test casting of PETScArray. @@ -166,6 +209,7 @@ def test_petsc_cast(): '(PetscScalar (*)[info.gym][info.gxm]) arr2_vec;' +@skipif('petsc') def test_no_automatic_cast(): """ Verify that the compiler does not automatically generate casts for PETScArrays. @@ -180,6 +224,191 @@ def test_no_automatic_cast(): eqn = Eq(arr, f.laplace) - op = Operator(eqn, opt='noop') + with switchconfig(openmp=False): + op = Operator(eqn, opt='noop') assert len(op.body.casts) == 1 + + +@skipif('petsc') +def test_LinearSolveExpr(): + + grid = Grid((2, 2)) + + f = Function(name='f', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + eqn = Eq(f, g.laplace) + + linsolveexpr = LinearSolveExpr(eqn.rhs, target=f) + + # Check the target + assert linsolveexpr.target == f + # Check the solver parameters + assert linsolveexpr.solver_parameters == \ + {'ksp_type': 'gmres', 'pc_type': 'jacobi', 'ksp_rtol': 1e-07, + 'ksp_atol': 1e-50, 'ksp_divtol': 10000.0, 'ksp_max_it': 10000} + + +@skipif('petsc') +def test_dmda_create(): + + grid1 = Grid((2)) + grid2 = Grid((2, 2)) + grid3 = Grid((4, 5, 6)) + + f1 = Function(name='f1', grid=grid1, space_order=2) + f2 = Function(name='f2', grid=grid2, space_order=4) + f3 = Function(name='f3', grid=grid3, space_order=6) + + eqn1 = Eq(f1.laplace, 10) + eqn2 = Eq(f2.laplace, 10) + eqn3 = Eq(f3.laplace, 10) + + petsc1 = PETScSolve(eqn1, f1) + petsc2 = PETScSolve(eqn2, f2) + petsc3 = PETScSolve(eqn3, f3) + + with switchconfig(openmp=False): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + op3 = Operator(petsc3, opt='noop') + + assert 'PetscCall(DMDACreate1d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + '2,1,2,NULL,&(da_so_2)));' in str(op1) + + assert 'PetscCall(DMDACreate2d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_so_4)));' \ + in str(op2) + + assert 'PetscCall(DMDACreate3d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ + ',1,1,1,1,6,NULL,NULL,NULL,&(da_so_6)));' in str(op3) + + # Check unique DMDA is created per grid, per space_order + f4 = Function(name='f4', grid=grid2, space_order=6) + eqn4 = Eq(f4.laplace, 10) + petsc4 = PETScSolve(eqn4, f4) + with switchconfig(openmp=False): + op4 = Operator(petsc2+petsc2+petsc4, opt='noop') + assert str(op4).count('DMDACreate2d') == 2 + + +@skipif('petsc') +def test_cinterface_petsc_struct(): + + grid = Grid(shape=(11, 11)) + f = Function(name='f', grid=grid, space_order=2) + eq = Eq(f.laplace, 10) + petsc = PETScSolve(eq, f) + + name = "foo" + with switchconfig(openmp=False): + op = Operator(petsc, name=name) + + # Trigger the generation of a .c and a .h files + ccode, hcode = op.cinterface(force=True) + + dirname = op._compiler.get_jit_dir() + assert os.path.isfile(os.path.join(dirname, "%s.c" % name)) + assert os.path.isfile(os.path.join(dirname, "%s.h" % name)) + + ccode = str(ccode) + hcode = str(hcode) + + assert 'include "%s.h"' % name in ccode + + # The public `struct MatContext` only appears in the header file + assert any(isinstance(i, PETScStruct) for i in op.parameters) + assert 'struct MatContext\n{' not in ccode + assert 'struct MatContext\n{' in hcode + + +@skipif('petsc') +def test_separate_eqn(): + """ + Test the separate_eqn function. + + This function is called within PETScSolve to decompose the equation + into the form F(x) = b. This is necessary to utilise the SNES + interface in PETSc. + """ + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=2) + g1 = Function(name='g1', grid=grid, space_order=2) + eq1 = Eq(f1.laplace, g1) + b1, F1 = separate_eqn(eq1, f1) + assert str(b1) == 'g1(x, y)' + assert str(F1) == 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))' + + f2 = TimeFunction(name='f2', grid=grid, space_order=2) + g2 = TimeFunction(name='g2', grid=grid, space_order=2) + eq2 = Eq(f2.dt, f2.laplace + g2) + b2, F2 = separate_eqn(eq2, f2.forward) + assert str(b2) == 'g2(t, x, y) + Derivative(f2(t, x, y), (x, 2))' + \ + ' + Derivative(f2(t, x, y), (y, 2)) + f2(t, x, y)/dt' + assert str(F2) == 'f2(t + dt, x, y)/dt' + + # Implicit Time Stepping + eqn3 = Eq(f2.dt, f2.forward.laplace + g2) + b3, F3 = separate_eqn(eqn3, f2.forward) + assert str(b3) == 'g2(t, x, y) + f2(t, x, y)/dt' + assert str(F3) == '-Derivative(f2(t + dt, x, y), (x, 2)) - ' + \ + 'Derivative(f2(t + dt, x, y), (y, 2)) + f2(t + dt, x, y)/dt' + + # Semi-implicit Time Stepping + eqn4 = Eq(f2.dt, f2.forward.laplace + f2.laplace + g2) + b4, F4 = separate_eqn(eqn4, f2.forward) + assert str(b4) == 'g2(t, x, y) + Derivative(f2(t, x, y), (x, 2))' + \ + ' + Derivative(f2(t, x, y), (y, 2)) + f2(t, x, y)/dt' + assert str(F4) == '-Derivative(f2(t + dt, x, y), (x, 2)) -' + \ + ' Derivative(f2(t + dt, x, y), (y, 2)) + f2(t + dt, x, y)/dt' + + +@skipif('petsc') +@pytest.mark.parametrize('expr, so, target, expected', [ + ('f1.laplace', 2, 'f1', '-2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'), + ('f1 + f1.laplace', 2, 'f1', + 'f1(x, y) - 2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'), + ('g1.dx + f1.dx', 2, 'f1', '-f1(x, y)/h_x'), + ('10 + f1.dx2', 2, 'g1', '0'), + ('(f1 * g1.dx).dy', 2, 'f1', + '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)'), + ('(f1 * g1.dx).dy', 2, 'g1', '-(-1/h_y)*f1(x, y)*g1(x, y)/h_x'), + ('f2.laplace', 2, 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'), + ('f2*g2', 2, 'f2', 'f2(t, x, y)*g2(t, x, y)'), + ('g2*f2.laplace', 2, 'f2', + '(-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2)*g2(t, x, y)'), + ('f2.forward', 2, 'f2.forward', 'f2(t + dt, x, y)'), + ('f2.forward.laplace', 2, 'f2.forward', + '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 2, 'f2.forward', + '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 2, + 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'), + ('f2.laplace', 4, 'f2', '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 4, 'f2.forward', + '-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 4, 'f2', + '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'), + ('f2.forward*f2.forward.laplace', 4, 'f2.forward', + '(-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2)*f2(t + dt, x, y)') +]) +def test_centre_stencil(expr, so, target, expected): + """ + Test extraction of centre stencil from an equation. + """ + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + h1 = Function(name='h1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + h2 = TimeFunction(name='h2', grid=grid, space_order=so) # noqa + + centre = centre_stencil(eval(expr), eval(target)) + + assert str(centre) == expected From dd753958eaef48fb5661517ac139e75d9e5abe9c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 10 Jul 2024 11:31:30 +0100 Subject: [PATCH 096/151] dsl: Dispatch to sympy.Add not both Add and EvalDerivative --- .github/workflows/examples-mpi.yml | 1 + .github/workflows/examples.yml | 1 + .github/workflows/flake8.yml | 1 + .github/workflows/pytest-core-mpi.yml | 1 + .github/workflows/pytest-core-nompi.yml | 1 + .github/workflows/pytest-petsc.yml | 10 +- .github/workflows/tutorials.yml | 1 + devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 22 +- devito/ir/iet/nodes.py | 17 +- devito/ir/iet/visitors.py | 11 +- devito/ir/stree/algorithms.py | 1 - devito/operator/operator.py | 15 +- devito/passes/clusters/__init__.py | 2 +- devito/passes/clusters/petsc.py | 25 ++ devito/passes/iet/__init__.py | 1 + devito/passes/iet/petsc.py | 298 ++++++++++++++++++ devito/petsc/clusters.py | 7 +- devito/petsc/iet.py | 372 ++++++++++++++++++++++ devito/petsc/iet/__init__.py | 2 +- devito/petsc/iet/nodes.py | 40 ++- devito/petsc/iet/passes.py | 292 +++++++---------- devito/petsc/iet/routines.py | 401 ++++++++++++++++++++++++ devito/petsc/solve.py | 126 ++++---- devito/petsc/types.py | 303 ++++++++++-------- devito/petsc/utils.py | 21 +- docker/Dockerfile.devito | 1 - docker/Dockerfile.petsc | 49 +++ tests/test_iet.py | 19 +- tests/test_petsc.py | 311 +++++++++++++----- 30 files changed, 1816 insertions(+), 538 deletions(-) create mode 100644 devito/passes/clusters/petsc.py create mode 100644 devito/passes/iet/petsc.py create mode 100644 devito/petsc/iet.py create mode 100644 devito/petsc/iet/routines.py create mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 3bbebabb1c..9df9251991 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,6 +17,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 026331c881..4490f47547 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 95a94f6e79..b61e777454 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 4f2853e46e..cc4b8ab683 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 54b5b05b1f..802aa3afb3 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 08da1220fa..ffb140819e 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -10,14 +10,10 @@ on: push: branches: - master - # Temporary - - separate-eqn + - enable-opt pull_request: branches: - master - # Temporary - - separate-eqn - - centre-stencil jobs: pytest: @@ -69,11 +65,11 @@ jobs: - name: Test with pytest run: | - ${{ env.RUN_CMD }} pytest -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + ${{ env.RUN_CMD }} pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" uses: codecov/codecov-action@v4 with: token: ${{ secrets.CODECOV_TOKEN }} - name: ${{ matrix.name }} \ No newline at end of file + name: ${{ matrix.name }} diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index d27af78fcf..af7f17d255 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 5aa71a0529..f0b0d718ea 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -376,7 +376,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index ff524a1f90..6585663c7a 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -9,11 +9,11 @@ Stencil, detect_io, detect_accesses) from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict -from devito.types import (Eq, Inc, ReduceMax, ReduceMin, - relational_min, MatVecEq, RHSEq) +from devito.types import Eq, Inc, ReduceMax, ReduceMin +from devito.petsc.types import InjectSolveEq __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', - 'identity_mapper'] + 'identity_mapper', 'OpInjectSolve'] class IREq(sympy.Eq, Pickable): @@ -104,8 +104,7 @@ def detect(cls, expr): Inc: OpInc, ReduceMax: OpMax, ReduceMin: OpMin, - MatVecEq: OpMatVec, - RHSEq: OpRHS + InjectSolveEq: OpInjectSolve } try: return reduction_mapper[type(expr)] @@ -122,12 +121,7 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') - -# Operations required by a Linear Solve of the form Ax=b: -# Application of linear operator on a vector -> op for matrix-vector multiplication. -OpMatVec = Operation('matvec') -# Building the right-hand side of linear system. -OpRHS = Operation('rhs') +OpInjectSolve = Operation('solve') identity_mapper = { @@ -232,19 +226,19 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - + # from IPython import embed; embed() # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - + # from IPython import embed; embed() expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + return expr @property diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 1739731fdc..c3569fa60b 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1135,7 +1135,7 @@ def defines(self): class Callback(Call): """ - Callback as a function pointer. + Base class for special callback types. Parameters ---------- @@ -1148,18 +1148,29 @@ class Callback(Call): Notes ----- - The reason Callback is an IET type rather than a SymPy type is + - The reason Callback is an IET type rather than a SymPy type is due to the fact that, when represented at the SymPy level, the IET engine fails to bind the callback to a specific Call. Consequently, errors occur during the creation of the call graph. """ # TODO: Create a common base class for Call and Callback to avoid # having arguments=None here - def __init__(self, name, retval, param_types, arguments=None): + def __init__(self, name, retval=None, param_types=None, arguments=None): super().__init__(name=name) self.retval = retval self.param_types = as_tuple(param_types) + @property + def callback_form(self): + """ + A string representation of the callback form. + + Notes + ----- + To be overridden by subclasses. + """ + return + class Section(List): diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index d2f5d9db3b..58d69430e9 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -625,7 +625,7 @@ def visit_Lambda(self, o): return LambdaCollection([top, c.Block(body)]) def visit_Callback(self, o, nested_call=False): - return CallbackArg(o.name, o.retval, o.param_types) + return CallbackArg(o) def visit_HaloSpot(self, o): body = flatten(self._visit(i) for i in o.children) @@ -1434,11 +1434,8 @@ def sorted_efuncs(efuncs): class CallbackArg(c.Generable): - def __init__(self, name, retval, param_types): - self.name = name - self.retval = retval - self.param_types = param_types + def __init__(self, callback): + self.callback = callback def generate(self): - param_types_str = ', '.join([str(t) for t in self.param_types]) - yield "(%s (*)(%s))%s" % (self.retval, param_types_str, self.name) + yield self.callback.callback_form diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index c738db645a..a85b93460d 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,7 +163,6 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: - # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index acc5524c5d..39277fabbc 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -355,7 +355,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - + # from IPython import embed; embed() return processed # Compilation -- Cluster level @@ -500,6 +500,9 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Target-independent optimizations minimize_symbols(graph) + # If necessary, sort frees into a specific order + sort_frees(graph) + return graph.root, graph # Read-only properties exposed to the outside world @@ -518,7 +521,12 @@ def dimensions(self): # During compilation other Dimensions may have been produced dimensions = FindSymbols('dimensions').visit(self) - ret.update(d for d in dimensions if d.is_PerfKnob) + + # NOTE: Should these dimensions be integrated into self._dimensions instead? + # In which case they would get picked up before this + struct_dims = derive_callback_dims(self._func_table) + + ret.update(d for d in dimensions if d.is_PerfKnob or d in struct_dims) ret = tuple(sorted(ret, key=attrgetter('name'))) @@ -526,7 +534,8 @@ def dimensions(self): @cached_property def input(self): - return tuple(i for i in self.parameters if i.is_Input) + struct_params = derive_struct_inputs(self.parameters) + return tuple(i for i in self.parameters+struct_params if i.is_Input) @cached_property def temporaries(self): diff --git a/devito/passes/clusters/__init__.py b/devito/passes/clusters/__init__.py index 1b14ce1aad..c41a628e06 100644 --- a/devito/passes/clusters/__init__.py +++ b/devito/passes/clusters/__init__.py @@ -8,4 +8,4 @@ from .implicit import * # noqa from .misc import * # noqa from .derivatives import * # noqa -from .unevaluate import * # noqa \ No newline at end of file +from .unevaluate import * # noqa diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py new file mode 100644 index 0000000000..d784d99c0d --- /dev/null +++ b/devito/passes/clusters/petsc.py @@ -0,0 +1,25 @@ +from devito.tools import timed_pass +from devito.petsc import LinearSolveExpr + +__all__ = ['petsc_lift'] + + +@timed_pass() +def petsc_lift(clusters): + """ + Lift the iteration space surrounding each PETSc equation to create + distinct iteration loops. This simplifys the movement of the loops + into specific callback functions generated at the IET level. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions, not the ones that stay inside the main kernel. + """ + processed = [] + for c in clusters: + + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index c09db00c9b..cf6a35de90 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,3 +8,4 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa +from .petsc import * # noqa diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py new file mode 100644 index 0000000000..8c8aeafa1c --- /dev/null +++ b/devito/passes/iet/petsc.py @@ -0,0 +1,298 @@ +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, MatVecAction, + Transformer, FindSymbols, LinearSolverExpression, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition) +from devito.types import (PetscMPIInt, PETScStruct, DMDALocalInfo, DM, Mat, + Vec, KSP, PC, SNES, PetscErrorCode) +from devito.symbolics import Byref, Macro, FieldFromPointer +import cgen as c + +__all__ = ['lower_petsc'] + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't + # appear in ccode + + # Check if PETScSolve was used and count occurrences. Each PETScSolve + # will have a unique MatVecAction. + is_petsc = FindNodes(MatVecAction).visit(iet) + + if is_petsc: + + # Collect all solution fields we're solving for + targets = [i.expr.rhs.target for i in is_petsc] + + # Initalize PETSc + init = init_petsc(**kwargs) + + # Create context data struct. + struct = build_struct(iet) + + objs = build_core_objects(targets[-1], struct, **kwargs) + + # Create core PETSc calls required by general linear solves, + # which only need to be generated once e.g create DMDA. + core = core_petsc(targets[-1], objs, **kwargs) + + matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) + + main_mapper = {} + + setup = [] + efuncs = [] + for target in unique_targets: + + solver_objs = build_solver_objs(target) + + matvec_callback_body_iters = [] + + solver_setup = False + + for iter, (matvec,) in matvec_mapper.items(): + + if matvec.expr.rhs.target == target: + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec) + setup.extend(solver) + solver_setup = True + + # Only need to generate solver setup once per target + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + solver_setup = True + + matvec_body = matvec_body_list._rebuild(body=[ + matvec_body_list.body, iter[0]]) + matvec_body_list = matvec_body_list._rebuild(body=matvec_body) + + main_mapper.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, matvec_callback_body_iters, solver_objs, objs, struct) + + setup.append(matvec_op) + setup.append(c.Line()) + + efuncs.append(matvec_callback) + + # Remove the LinSolveExpr from iet and efuncs that were used to carry + # metadata e.g solver_parameters + main_mapper.update(rebuild_expr_mapper(iet)) + efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} + + iet = Transformer(main_mapper).visit(iet) + efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] + + # Replace symbols appearing in each efunc with a pointer to the PETScStruct + efuncs = transform_efuncs(efuncs, struct) + + body = iet.body._rebuild(body=(tuple(init_setup) + iet.body.body)) + iet = iet._rebuild(body=body) + + return iet, {} + + +def init_petsc(**kwargs): + + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve. + # NOTE: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options. + initialize = Call('PetscCall', [Call('PetscInitialize', + arguments=['NULL', 'NULL', + 'NULL', 'NULL'])]) + + return tuple([initialize]) + + +def build_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct. + usr_ctx = [] + + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def core_petsc(target, objs, **kwargs): + # Assumption: all targets are generated from the same Grid, + # so we can use any target. + + # MPI + call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', + arguments=[objs['comm'], + Byref(objs['size'])])]) + # Create DMDA + dmda = create_dmda(target, objs) + dm_setup = Call('PetscCall', [Call('DMSetUp', arguments=[objs['da']])]) + dm_app_ctx = Call('PetscCall', [Call('DMSetApplicationContext', + arguments=[objs['da'], objs['struct']])]) + dm_mat_type = Call('PetscCall', [Call('DMSetMatType', + arguments=[objs['da'], 'MATSHELL'])]) + dm_local_info = Call('PetscCall', [Call('DMDAGetLocalInfo', + arguments=[objs['da'], Byref(objs['info'])])]) + + return tuple([call_mpi, dmda, dm_setup, dm_app_ctx, dm_mat_type, dm_local_info]) + + +def build_core_objects(target, struct, **kwargs): + + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return {'da': DM(name='da'), + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'struct': struct} + + +def create_dmda(target, objs): + + args = [objs['comm']] + + args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] + + # stencil type + args += ['DMDA_STENCIL_BOX'] + + # global dimensions + args += list(target.shape_global)[::-1] + + # no.of processors in each dimension + args += list(target.grid.distributor.topology)[::-1] + + args += [1, target.space_order] + + args += ['NULL' for _ in range(len(target.space_dimensions))] + + args += [Byref(objs['da'])] + + dmda = Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args) + + return dmda + + +def build_solver_objs(target): + + return {'Jac': Mat(name='J_'+str(target.name)), + 'x_global': Vec(name='x_global_'+str(target.name)), + 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), + 'b_global': Vec(name='b_global_'+str(target.name)), + 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), + 'ksp': KSP(name='ksp_'+str(target.name)), + 'pc': PC(name='pc_'+str(target.name)), + 'snes': SNES(name='snes_'+str(target.name)), + 'x': Vec(name='x_'+str(target.name)), + 'y': Vec(name='y_'+str(target.name))} + + +def generate_solver_calls(solver_objs, objs, matvec): + + snes_create = Call('PetscCall', [Call('SNESCreate', arguments=[ + objs['comm'], Byref(solver_objs['snes'])])]) + + snes_set_dm = Call('PetscCall', [Call('SNESSetDM', arguments=[ + solver_objs['snes'], objs['da']])]) + + create_matrix = Call('PetscCall', [Call('DMCreateMatrix', arguments=[ + objs['da'], Byref(solver_objs['Jac'])])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = Call('PetscCall', [Call('SNESSetType', arguments=[ + solver_objs['snes'], 'SNESKSPONLY'])]) + + global_x = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['x_global'])])]) + + local_x = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['x_local'])])]) + + global_b = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['b_global'])])]) + + local_b = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['b_local'])])]) + + snes_get_ksp = Call('PetscCall', [Call('SNESGetKSP', arguments=[ + solver_objs['snes'], Byref(solver_objs['ksp'])])]) + + return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, + global_x, local_x, global_b, local_b, snes_get_ksp]) + + +def create_matvec_callback(target, matvec_callback_body_iters, + solver_objs, objs, struct): + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + defn_struct = Definition(struct) + + get_context = Call('PetscCall', [Call('MatShellGetContext', + arguments=[solver_objs['Jac'], + Byref(struct)])]) + + body = List(body=[defn_struct, + get_context, + matvec_callback_body_iters]) + + matvec_callback = Callable('MyMatShellMult_'+str(target.name), + matvec_body, + retval=objs['err'], + parameters=(solver_objs['Jac'], + solver_objs['x'], + solver_objs['y'])) + + matvec_operation = Call('PetscCall', [ + Call('MatShellSetOperation', arguments=[solver_objs['Jac'], + 'MATOP_MULT', + Callback(matvec_callback.name, + Void, Void)])]) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(iet): + + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for + expr in FindNodes(LinearSolverExpression).visit(iet)} + + +def transform_efuncs(efuncs, struct): + + efuncs_new = [] + for efunc in efuncs: + new_body = efunc.body + for i in struct.usr_ctx: + new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) + efunc_with_new_body = efunc._rebuild(body=new_body) + efuncs_new.append(efunc_with_new_body) + + return efuncs_new + + +Null = Macro('NULL') +Void = Macro('void') + +petsc_call = String('PetscCall') +petsc_call_mpi = String('PetscCallMPI') +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') + +linear_solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 8aca92b035..5d77288f39 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -6,10 +6,9 @@ def petsc_lift(clusters): """ Lift the iteration space surrounding each PETSc equation to create - distinct iteration loops. This simplifys the movement of the loops - into specific callback functions generated at the IET level. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions, not the ones that stay inside the main kernel. + distinct iteration loops. + # TODO: Can probably remove this now due to recursive compilation, but + # leaving it for now. """ processed = [] for c in clusters: diff --git a/devito/petsc/iet.py b/devito/petsc/iet.py new file mode 100644 index 0000000000..659ac9f94d --- /dev/null +++ b/devito/petsc/iet.py @@ -0,0 +1,372 @@ +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, + Transformer, FindSymbols, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition, BlankLine, PointerCast) +<<<<<<< HEAD +from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, +======= +from devito.petsc import (PetscMPIInt, PETScStruct, DM, Mat, +>>>>>>> b22824875 (fix circular imports) + Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) +from devito.symbolics import Byref, Macro, FieldFromPointer +import cgen as c +from devito.petsc.nodes import MatVecAction, RHSLinearSystem, LinearSolverExpression + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # Check if PETScSolve was used. + petsc_nodes = FindNodes(MatVecAction).visit(iet) + + if not petsc_nodes: + return iet, {} + + else: + # Collect all petsc solution fields + unique_targets = list(set([i.expr.rhs.target for i in petsc_nodes])) + + # Initalize PETSc + init = init_petsc(**kwargs) + + # Create context data struct + struct = build_struct(iet) + + objs = build_core_objects(unique_targets[-1], **kwargs) + + # Create core PETSc calls (not specific to each PETScSolve) + core = core_petsc(unique_targets[-1], struct, objs, **kwargs) + + matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) + + main_mapper = {} + + setup = [] + efuncs = [] + + for target in unique_targets: + + solver_objs = build_solver_objs(target) + + matvec_body_list = List() + + solver_setup = False + + for iter, (matvec,) in matvec_mapper.items(): + + # Skip the MatVecAction if it is not associated with the target + # There will be more than one MatVecAction associated with the target + # e.g interior matvec + BC matvecs + if matvec.expr.rhs.target != target: + continue + + # Only need to generate solver setup once per target + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + solver_setup = True + + matvec_body = matvec_body_list._rebuild(body=[ + matvec_body_list.body, iter[0]]) + matvec_body_list = matvec_body_list._rebuild(body=matvec_body) + + main_mapper.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, matvec_body_list, solver_objs, objs, + struct) + + setup.append(matvec_op) + setup.append(BlankLine) + efuncs.append(matvec_callback) + + # Remove the LinSolveExpr from iet and efuncs that were used to carry + # metadata e.g solver_parameters + main_mapper.update(rebuild_expr_mapper(iet)) + efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} + + iet = Transformer(main_mapper).visit(iet) + efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] + + # Replace symbols appearing in each efunc with a pointer to the PETScStruct + efuncs = transform_efuncs(efuncs, struct) + + body = iet.body._rebuild(init=init, body=core + tuple(setup) + iet.body.body) + iet = iet._rebuild(body=body) + + return iet, {'efuncs': efuncs} + + +def init_petsc(**kwargs): + + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve + # TODO: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options + initialize = Call(petsc_call, [ + Call('PetscInitialize', arguments=[Null, Null, Null, Null])]) + + return tuple([petsc_func_begin_user, initialize]) + + +def build_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct + usr_ctx = [] + + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def core_petsc(target, struct, objs, **kwargs): + + # MPI + call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', + arguments=[objs['comm'], + Byref(objs['size'])])]) + # Create DMDA + dmda = create_dmda(target, objs) + dm_setup = Call(petsc_call, [ + Call('DMSetUp', arguments=[objs['da']])]) + dm_app_ctx = Call(petsc_call, [ + Call('DMSetApplicationContext', arguments=[objs['da'], struct])]) + dm_mat_type = Call(petsc_call, [ + Call('DMSetMatType', arguments=[objs['da'], 'MATSHELL'])]) + + return tuple([petsc_func_begin_user, call_mpi, dmda, dm_setup, + dm_app_ctx, dm_mat_type, BlankLine]) + + +def build_core_objects(target, **kwargs): + + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return {'da': DM(name='da', liveness='eager'), + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'err': PetscErrorCode(name='err')} + + +def create_dmda(target, objs): + + args = [objs['comm']] + + args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] + + # Stencil type + if len(target.space_dimensions) > 1: + args += ['DMDA_STENCIL_BOX'] + + # Global dimensions + args += list(target.shape_global)[::-1] + + # No.of processors in each dimension + if len(target.space_dimensions) > 1: + args += list(target.grid.distributor.topology)[::-1] + + args += [1, target.space_order] + + args += [Null for _ in range(len(target.space_dimensions))] + + args += [Byref(objs['da'])] + + dmda = Call(petsc_call, [ + Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args)]) + + return dmda + + +def build_solver_objs(target): + + return {'Jac': Mat(name='J_'+str(target.name)), + 'x_global': Vec(name='x_global_'+str(target.name)), + 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), + 'b_global': Vec(name='b_global_'+str(target.name)), + 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), + 'ksp': KSP(name='ksp_'+str(target.name)), + 'pc': PC(name='pc_'+str(target.name)), + 'snes': SNES(name='snes_'+str(target.name)), + 'X_global': Vec(name='X_global_'+str(target.name)), + 'Y_global': Vec(name='Y_global_'+str(target.name)), + 'X_local': Vec(name='X_local_'+str(target.name), liveness='eager'), + 'Y_local': Vec(name='Y_local_'+str(target.name), liveness='eager') + } + + +def generate_solver_calls(solver_objs, objs, matvec, target): + + solver_params = matvec.expr.rhs.solver_parameters + + snes_create = Call(petsc_call, [Call('SNESCreate', arguments=[ + objs['comm'], Byref(solver_objs['snes'])])]) + + snes_set_dm = Call(petsc_call, [Call('SNESSetDM', arguments=[ + solver_objs['snes'], objs['da']])]) + + create_matrix = Call(petsc_call, [Call('DMCreateMatrix', arguments=[ + objs['da'], Byref(solver_objs['Jac'])])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = Call(petsc_call, [Call('SNESSetType', arguments=[ + solver_objs['snes'], 'SNESKSPONLY'])]) + + global_x = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['x_global'])])]) + + local_x = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['x_local'])])]) + + global_b = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['b_global'])])]) + + local_b = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['b_local'])])]) + + snes_get_ksp = Call(petsc_call, [Call('SNESGetKSP', arguments=[ + solver_objs['snes'], Byref(solver_objs['ksp'])])]) + + vec_replace_array = Call(petsc_call, [ + Call('VecReplaceArray', arguments=[solver_objs[ + 'x_local'], FieldFromPointer(target._C_field_data, target._C_symbol)])]) + + ksp_set_tols = Call(petsc_call, [Call('KSPSetTolerances', arguments=[ + solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], + solver_params['ksp_divtol'], solver_params['ksp_max_it']])]) + + ksp_set_type = Call(petsc_call, [Call('KSPSetType', arguments=[ + solver_objs['ksp'], linear_solver_mapper[solver_params['ksp_type']]])]) + + ksp_get_pc = Call(petsc_call, [Call('KSPGetPC', arguments=[ + solver_objs['ksp'], Byref(solver_objs['pc'])])]) + + pc_set_type = Call(petsc_call, [Call('PCSetType', arguments=[ + solver_objs['pc'], linear_solver_mapper[solver_params['pc_type']]])]) + + ksp_set_from_ops = Call(petsc_call, [Call('KSPSetFromOptions', arguments=[ + solver_objs['ksp']])]) + + return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, + global_x, local_x, global_b, local_b, snes_get_ksp, + vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, + pc_set_type, ksp_set_from_ops]) + + +def create_matvec_callback(target, body, solver_objs, objs, struct): + + # There will be 2 PETScArrays within the body + petsc_arrays = [i.function for i in FindSymbols('indexedbases').visit(body) + if isinstance(i.function, PETScArray)] + + # There will only be one PETScArray that is written to within this body and + # one PETScArray which corresponds to the 'seed' vector + petsc_arr_write, = FindSymbols('writes').visit(body) + petsc_arr_seed, = [i for i in petsc_arrays if i.function != petsc_arr_write.function] + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + defn_struct = Definition(struct) + + mat_get_dm = Call(petsc_call, [Call('MatGetDM', arguments=[ + solver_objs['Jac'], Byref(objs['da'])])]) + + dm_get_app_context = Call(petsc_call, [Call('DMGetApplicationContext', arguments=[ + objs['da'], Byref(struct._C_symbol)])]) + + dm_get_local_xvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ + objs['da'], Byref(solver_objs['X_local'])])]) + + global_to_local_begin = Call(petsc_call, [Call('DMGlobalToLocalBegin', arguments=[ + objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) + + global_to_local_end = Call(petsc_call, [Call('DMGlobalToLocalEnd', arguments=[ + objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) + + dm_get_local_yvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ + objs['da'], Byref(solver_objs['Y_local'])])]) + + vec_get_array_y = Call(petsc_call, [Call('VecGetArray', arguments=[ + solver_objs['Y_local'], Byref(petsc_arr_write.function)])]) + + vec_get_array_x = Call(petsc_call, [Call('VecGetArray', arguments=[ + solver_objs['X_local'], Byref(petsc_arrays[0])])]) + + dm_get_local_info = Call(petsc_call, [Call('DMDAGetLocalInfo', arguments=[ + objs['da'], Byref(petsc_arrays[0].function.dmda_info)])]) + + casts = [PointerCast(i.function) for i in petsc_arrays] + + vec_restore_array_y = Call(petsc_call, [Call('VecRestoreArray', arguments=[ + solver_objs['Y_local'], Byref(petsc_arr_write[0])])]) + + vec_restore_array_x = Call(petsc_call, [Call('VecRestoreArray', arguments=[ + solver_objs['X_local'], Byref(petsc_arr_seed)])]) + + dm_local_to_global_begin = Call(petsc_call, [Call('DMLocalToGlobalBegin', arguments=[ + objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) + + dm_local_to_global_end = Call(petsc_call, [Call('DMLocalToGlobalEnd', arguments=[ + objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) + + func_return = Call('PetscFunctionReturn', arguments=[0]) + + matvec_body = List(body=[ + petsc_func_begin_user, defn_struct, mat_get_dm, dm_get_app_context, + dm_get_local_xvec, global_to_local_begin, global_to_local_end, + dm_get_local_yvec, vec_get_array_y, vec_get_array_x, dm_get_local_info, + casts, BlankLine, body, vec_restore_array_y, vec_restore_array_x, + dm_local_to_global_begin, dm_local_to_global_end, func_return]) + + matvec_callback = Callable( + 'MyMatShellMult_'+str(target.name), matvec_body, retval=objs['err'], + parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'])) + + matvec_operation = Call(petsc_call, [ + Call('MatShellSetOperation', arguments=[ + solver_objs['Jac'], 'MATOP_MULT', Callback(matvec_callback.name, + Void, Void)])]) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(callable): + + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for + expr in FindNodes(LinearSolverExpression).visit(callable)} + + +def transform_efuncs(efuncs, struct): + + efuncs_new = [] + for efunc in efuncs: + new_body = efunc.body + for i in struct.usr_ctx: + new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) + efunc_with_new_body = efunc._rebuild(body=new_body) + efuncs_new.append(efunc_with_new_body) + + return efuncs_new + + +Null = Macro('NULL') +Void = 'void' + +petsc_call = 'PetscCall' +petsc_call_mpi = 'PetscCallMPI' +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') + +linear_solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} diff --git a/devito/petsc/iet/__init__.py b/devito/petsc/iet/__init__.py index 7ce0b8c3f1..beb6d1f2d1 100644 --- a/devito/petsc/iet/__init__.py +++ b/devito/petsc/iet/__init__.py @@ -1 +1 @@ -from devito.petsc.iet import * # noqa \ No newline at end of file +from devito.petsc.iet import * # noqa diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index 4682ea91d2..5507f999e2 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -1,9 +1,9 @@ -from devito.ir.iet import Expression -from devito.ir.equations import OpMatVec, OpRHS +from devito.ir.iet import Expression, Callable, Callback +from devito.ir.equations import OpInjectSolve +from devito.tools import as_tuple class LinearSolverExpression(Expression): - """ Base class for general expressions required by a matrix-free linear solve of the form Ax=b. @@ -11,21 +11,33 @@ class LinearSolverExpression(Expression): pass -class MatVecAction(LinearSolverExpression): - +class InjectSolveDummy(LinearSolverExpression): """ - Expression representing matrix-vector multiplication. + Placeholder expression to run the iterative solver. """ - - def __init__(self, expr, pragmas=None, operation=OpMatVec): + def __init__(self, expr, pragmas=None, operation=OpInjectSolve): super().__init__(expr, pragmas=pragmas, operation=operation) -class RHSLinearSystem(LinearSolverExpression): +class PETScCallable(Callable): + def __init__(self, name, body, retval=None, parameters=None, + prefix=None, unused_parameters=None): + super().__init__(name, body, retval, parameters, prefix) + self._unused_parameters = as_tuple(unused_parameters) - """ - Expression to build the RHS of a linear system. - """ + @property + def unused_parameters(self): + return self._unused_parameters - def __init__(self, expr, pragmas=None, operation=OpRHS): - super().__init__(expr, pragmas=pragmas, operation=operation) + +class MatVecCallback(Callback): + @property + def callback_form(self): + param_types_str = ', '.join([str(t) for t in self.param_types]) + return "(%s (*)(%s))%s" % (self.retval, param_types_str, self.name) + + +class FormFunctionCallback(Callback): + @property + def callback_form(self): + return "%s" % self.name diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 3bc1e9a29b..fec645bfc8 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,49 +1,49 @@ -from collections import OrderedDict import cgen as c from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, - Transformer, FindSymbols, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition, BlankLine, PointerCast, filter_iterations, - retrieve_iteration_tree) +from devito.ir.iet import (FindNodes, Transformer, + MapNodes, Iteration, List, BlankLine, + filter_iterations, retrieve_iteration_tree, + Callable, CallableBody, DummyExpr, Call) from devito.symbolics import Byref, Macro, FieldFromPointer -from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, - Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) -from devito.petsc.iet.nodes import MatVecAction, LinearSolverExpression -from devito.petsc.utils import (solver_mapper, petsc_call, petsc_call_mpi, - core_metadata) +from devito.tools import filter_ordered +from devito.petsc.types import (PetscMPIInt, DM, Mat, LocalVec, + GlobalVec, KSP, PC, SNES, PetscErrorCode, DummyArg) +from devito.petsc.iet.nodes import InjectSolveDummy +from devito.petsc.utils import solver_mapper, core_metadata +from devito.petsc.iet.routines import PETScCallbackBuilder +from devito.petsc.iet.utils import petsc_call, petsc_call_mpi, petsc_struct @iet_pass def lower_petsc(iet, **kwargs): # Check if PETScSolve was used - petsc_nodes = FindNodes(MatVecAction).visit(iet) + petsc_nodes = FindNodes(InjectSolveDummy).visit(iet) if not petsc_nodes: return iet, {} - # Collect all petsc solution fields unique_targets = list({i.expr.rhs.target for i in petsc_nodes}) init = init_petsc(**kwargs) # Assumption is that all targets have the same grid so can use any target here objs = build_core_objects(unique_targets[-1], **kwargs) - objs['struct'] = build_petsc_struct(iet) # Create core PETSc calls (not specific to each PETScSolve) core = make_core_petsc_calls(objs, **kwargs) - # Create matvec mapper from the spatial iteration loops (exclude time loop if present) + # Create injectsolve mapper from the spatial iteration loops + # (exclude time loop if present) spatial_body = [] for tree in retrieve_iteration_tree(iet): root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] spatial_body.append(root) - matvec_mapper = MapNodes(Iteration, MatVecAction, - 'groupby').visit(List(body=spatial_body)) + injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy, + 'groupby').visit(List(body=spatial_body)) setup = [] + subs = {} # Create a different DMDA for each target with a unique space order unique_dmdas = create_dmda_objs(unique_targets) @@ -51,52 +51,48 @@ def lower_petsc(iet, **kwargs): for dmda in unique_dmdas.values(): setup.extend(create_dmda_calls(dmda, objs)) - subs = {} - efuncs = OrderedDict() + builder = PETScCallbackBuilder(**kwargs) # Create the PETSc calls which are specific to each target for target in unique_targets: solver_objs = build_solver_objs(target) - # Generate solver setup for target - for iter, (matvec,) in matvec_mapper.items(): - # Skip the MatVecAction if it is not associated with the target - # There will most likely be more than one MatVecAction - # associated with the target e.g interior matvec + BC matvecs - if matvec.expr.rhs.target != target: + # Generate the solver setup for target. This is required only + # once per target + for (injectsolve,) in injectsolve_mapper.values(): + # Skip if not associated with the target + if injectsolve.expr.rhs.target != target: continue - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) + solver_setup = generate_solver_setup(solver_objs, objs, injectsolve, target) + setup.extend(solver_setup) break - # Create the body of the matrix-vector callback for target - matvec_body_list = [] - for iter, (matvec,) in matvec_mapper.items(): - if matvec.expr.rhs.target != target: + # Generate all PETSc callback functions for the target via recusive compilation + for iter, (injectsolve,) in injectsolve_mapper.items(): + if injectsolve.expr.rhs.target != target: continue - matvec_body_list.append(iter[0]) - # Remove the iteration loop from the main kernel encapsulating - # the matvec equations since they are moved into the callback - subs.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, List(body=matvec_body_list), solver_objs, objs - ) + matvec_op, formfunc_op, runsolve = builder.make(injectsolve, + objs, solver_objs) + setup.extend([matvec_op, formfunc_op]) + subs.update({iter[0]: List(body=runsolve)}) + break - setup.extend([matvec_op, BlankLine]) - efuncs[matvec_callback.name] = matvec_callback + # Generate callback to populate main struct object + struct_main = petsc_struct('ctx', filter_ordered(builder.struct_params)) + struct_callback = generate_struct_callback(struct_main) + call_struct_callback = petsc_call(struct_callback.name, [Byref(struct_main)]) + calls_set_app_ctx = [petsc_call('DMSetApplicationContext', [i, Byref(struct_main)]) + for i in unique_dmdas] + setup.extend([BlankLine, call_struct_callback] + calls_set_app_ctx) - # Remove the LinSolveExpr's from iet and efuncs - subs.update(rebuild_expr_mapper(iet)) iet = Transformer(subs).visit(iet) - efuncs = transform_efuncs(efuncs, objs['struct']) - body = iet.body._rebuild(init=init, body=core+tuple(setup)+iet.body.body) + body = iet.body._rebuild( + init=init, body=core+tuple(setup)+(BlankLine,)+iet.body.body + ) iet = iet._rebuild(body=body) - metadata = core_metadata() - metadata.update({'efuncs': tuple(efuncs.values())}) + metadata.update({'efuncs': tuple(builder.efuncs.values())+(struct_callback,)}) return iet, metadata @@ -112,17 +108,6 @@ def init_petsc(**kwargs): return petsc_func_begin_user, initialize -def build_petsc_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct - usr_ctx = [] - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - def make_core_petsc_calls(objs, **kwargs): call_mpi = petsc_call_mpi('MPI_Comm_size', [objs['comm'], Byref(objs['size'])]) @@ -155,10 +140,9 @@ def create_dmda_objs(unique_targets): def create_dmda_calls(dmda, objs): dmda_create = create_dmda(dmda, objs) dm_setup = petsc_call('DMSetUp', [dmda]) - dm_app_ctx = petsc_call('DMSetApplicationContext', [dmda, objs['struct']]) dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) - - return dmda_create, dm_setup, dm_app_ctx, dm_mat_type, BlankLine + dm_get_local_info = petsc_call('DMDAGetLocalInfo', [dmda, Byref(dmda.info)]) + return dmda_create, dm_setup, dm_mat_type, dm_get_local_info, BlankLine def create_dmda(dmda, objs): @@ -199,24 +183,25 @@ def build_solver_objs(target): name = target.name return { 'Jac': Mat(name='J_%s' % name), - 'x_global': Vec(name='x_global_%s' % name), - 'x_local': Vec(name='x_local_%s' % name, liveness='eager'), - 'b_global': Vec(name='b_global_%s' % name), - 'b_local': Vec(name='b_local_%s' % name, liveness='eager'), + 'x_global': GlobalVec(name='x_global_%s' % name), + 'x_local': LocalVec(name='x_local_%s' % name, liveness='eager'), + 'b_global': GlobalVec(name='b_global_%s' % name), + 'b_local': LocalVec(name='b_local_%s' % name), 'ksp': KSP(name='ksp_%s' % name), 'pc': PC(name='pc_%s' % name), 'snes': SNES(name='snes_%s' % name), - 'X_global': Vec(name='X_global_%s' % name), - 'Y_global': Vec(name='Y_global_%s' % name), - 'X_local': Vec(name='X_local_%s' % name, liveness='eager'), - 'Y_local': Vec(name='Y_local_%s' % name, liveness='eager') + 'X_global': GlobalVec(name='X_global_%s' % name), + 'Y_global': GlobalVec(name='Y_global_%s' % name), + 'X_local': LocalVec(name='X_local_%s' % name, liveness='eager'), + 'Y_local': LocalVec(name='Y_local_%s' % name, liveness='eager'), + 'dummy': DummyArg(name='dummy_%s' % name) } -def generate_solver_calls(solver_objs, objs, matvec, target): +def generate_solver_setup(solver_objs, objs, injectsolve, target): dmda = objs['da_so_%s' % target.space_order] - solver_params = matvec.expr.rhs.solver_parameters + solver_params = injectsolve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) @@ -227,6 +212,11 @@ def generate_solver_calls(solver_objs, objs, matvec, target): # NOTE: Assumming all solves are linear for now. snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) + snes_set_jac = petsc_call( + 'SNESSetJacobian', [solver_objs['snes'], solver_objs['Jac'], + solver_objs['Jac'], 'MatMFFDComputeJacobian', Null] + ) + global_x = petsc_call('DMCreateGlobalVector', [dmda, Byref(solver_objs['x_global'])]) @@ -259,9 +249,8 @@ def generate_solver_calls(solver_objs, objs, matvec, target): ksp_get_pc = petsc_call('KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])]) - pc_set_type = petsc_call( - 'PCSetType', [solver_objs['pc'], solver_mapper[solver_params['pc_type']]] - ) + # Even though the default will be jacobi, set to PCNONE for now + pc_set_type = petsc_call('PCSetType', [solver_objs['pc'], 'PCNONE']) ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) @@ -269,6 +258,7 @@ def generate_solver_calls(solver_objs, objs, matvec, target): snes_create, snes_set_dm, create_matrix, + snes_set_jac, snes_set_type, global_x, local_x, @@ -284,131 +274,63 @@ def generate_solver_calls(solver_objs, objs, matvec, target): ) -def create_matvec_callback(target, body, solver_objs, objs): - dmda = objs['da_so_%s' % target.space_order] - - # There will be 2 PETScArrays within the body - petsc_arrays = [i for i in FindSymbols('indexedbases').visit(body) - if isinstance(i.function, PETScArray)] - - # There will only be one PETScArray that is written to within this body and - # one PETScArray which corresponds to the 'seed' vector - petsc_arr_write, = FindSymbols('writes').visit(body) - petsc_arr_seed, = [i.function for i in petsc_arrays - if i.function != petsc_arr_write.function] - - define_arrays = [Definition(i.function) for i in petsc_arrays] - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - define_struct = Definition(objs['struct']) - - mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) - - dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(objs['struct']._C_symbol)] +def generate_struct_callback(struct): + body = [DummyExpr(FieldFromPointer(i._C_symbol, struct), + i._C_symbol) for i in struct.fields] + struct_callback_body = CallableBody( + List(body=body), init=tuple([petsc_func_begin_user]), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) ) - - dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] - ) - - global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], - 'INSERT_VALUES', solver_objs['X_local']] - ) - - global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] - ]) - - dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + struct_callback = Callable( + 'PopulateMatContext', struct_callback_body, PetscErrorCode(name='err'), + parameters=[struct] ) + return struct_callback - vec_get_array_y = petsc_call( - 'VecGetArray', [solver_objs['Y_local'], Byref(petsc_arr_write._C_symbol)] - ) - vec_get_array_x = petsc_call( - 'VecGetArray', [solver_objs['X_local'], Byref(petsc_arr_seed._C_symbol)] - ) +@iet_pass +def sort_frees(iet): + frees = iet.body.frees - dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(petsc_arr_seed.function.dmda_info)] - ) + if not frees: + return iet, {} - casts = [PointerCast(i.function) for i in petsc_arrays] + destroys = ["VecDestroy", "MatDestroy", "SNESDestroy", "DMDestroy"] + priority = {k: i for i, k in enumerate(destroys, start=1)} - vec_restore_array_y = petsc_call( - 'VecRestoreArray', [solver_objs['Y_local'], Byref(petsc_arr_write._C_symbol)] - ) + def key(i): + for destroy, prio in priority.items(): + if destroy in str(i): + return prio + return float('inf') - vec_restore_array_x = petsc_call( - 'VecRestoreArray', [solver_objs['X_local'], Byref(petsc_arr_seed._C_symbol)] - ) + frees = sorted(frees, key=key) - dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] - ]) - - dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] - ]) - - func_return = Call('PetscFunctionReturn', arguments=[0]) - - matvec_body = List(body=[ - petsc_func_begin_user, - define_arrays, - define_struct, - mat_get_dm, - dm_get_app_context, - dm_get_local_xvec, - global_to_local_begin, - global_to_local_end, - dm_get_local_yvec, - vec_get_array_y, - vec_get_array_x, - dm_get_local_info, - casts, - BlankLine, - body, - vec_restore_array_y, - vec_restore_array_x, - dm_local_to_global_begin, - dm_local_to_global_end, - func_return]) - - matvec_callback = Callable( - 'MyMatShellMult_%s' % target.name, matvec_body, retval=objs['err'], - parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global']) - ) + body = iet.body._rebuild(frees=frees) + iet = iet._rebuild(body=body) + return iet, {} - matvec_operation = petsc_call( - 'MatShellSetOperation', [solver_objs['Jac'], 'MATOP_MULT', - Callback(matvec_callback.name, void, void)] - ) - return matvec_callback, matvec_operation +@iet_pass +def sort_frees(iet): + frees = iet.body.frees + if not frees: + return iet, {} -def rebuild_expr_mapper(callable): - # This mapper removes LinSolveExpr instances from the callable - # These expressions were previously used in lower_petc to carry metadata, - # such as solver_parameters - nodes = FindNodes(LinearSolverExpression).visit(callable) - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for expr in nodes} + destroys = ["VecDestroy", "MatDestroy", "SNESDestroy", "DMDestroy"] + priority = {k: i for i, k in enumerate(destroys, start=1)} + def key(i): + for destroy, prio in priority.items(): + if destroy in str(i): + return prio + return float('inf') -def transform_efuncs(efuncs, struct): - subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} - for efunc in efuncs.values(): - transformed_efunc = Transformer(rebuild_expr_mapper(efunc)).visit(efunc) - transformed_efunc = Uxreplace(subs).visit(transformed_efunc) - efuncs[efunc.name] = transformed_efunc - return efuncs + frees = sorted(frees, key=key) + body = iet.body._rebuild(frees=frees) + iet = iet._rebuild(body=body) + return iet, {} Null = Macro('NULL') diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py new file mode 100644 index 0000000000..c14e06269f --- /dev/null +++ b/devito/petsc/iet/routines.py @@ -0,0 +1,401 @@ +from collections import OrderedDict + +import cgen as c + +from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, + Callable) +from devito.symbolics import Byref, FieldFromPointer, Macro +from devito.petsc.types import PETScStruct +from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, + MatVecCallback) +from devito.petsc.utils import petsc_call + + +class PETScCallbackBuilder: + """ + Build IET routines to generate PETSc callback functions. + """ + def __new__(cls, rcompile=None, **kwargs): + obj = object.__new__(cls) + obj.rcompile = rcompile + obj._efuncs = OrderedDict() + obj._struct_params = [] + + return obj + + @property + def efuncs(self): + return self._efuncs + + @property + def struct_params(self): + return self._struct_params + + def make(self, injectsolve, objs, solver_objs): + matvec_callback, formfunc_callback, formrhs_callback = self.make_all( + injectsolve, objs, solver_objs + ) + + matvec_operation = petsc_call( + 'MatShellSetOperation', [solver_objs['Jac'], 'MATOP_MULT', + MatVecCallback(matvec_callback.name, void, void)] + ) + formfunc_operation = petsc_call( + 'SNESSetFunction', + [solver_objs['snes'], Null, + FormFunctionCallback(formfunc_callback.name, void, void), Null] + ) + runsolve = self.runsolve(solver_objs, objs, formrhs_callback, injectsolve) + + return matvec_operation, formfunc_operation, runsolve + + def make_all(self, injectsolve, objs, solver_objs): + matvec_callback = self.make_matvec(injectsolve, objs, solver_objs) + formfunc_callback = self.make_formfunc(injectsolve, objs, solver_objs) + formrhs_callback = self.make_formrhs(injectsolve, objs, solver_objs) + + self._efuncs[matvec_callback.name] = matvec_callback + self._efuncs[formfunc_callback.name] = formfunc_callback + self._efuncs[formrhs_callback.name] = formrhs_callback + + return matvec_callback, formfunc_callback, formrhs_callback + + def make_matvec(self, injectsolve, objs, solver_objs): + target = injectsolve.expr.rhs.target + # Compile matvec `eqns` into an IET via recursive compilation + irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, + options={'mpi': False}) + body_matvec = self.create_matvec_body(injectsolve, irs_matvec.uiet.body, + solver_objs, objs) + + matvec_callback = PETScCallable( + 'MyMatShellMult_%s' % target.name, body_matvec, retval=objs['err'], + parameters=( + solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'] + ) + ) + return matvec_callback + + def create_matvec_body(self, injectsolve, body, solver_objs, objs): + linsolveexpr = injectsolve.expr.rhs + + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + + struct = build_petsc_struct(body, 'matvec', liveness='eager') + + y_matvec = linsolveexpr.arrays['y_matvec'] + x_matvec = linsolveexpr.arrays['x_matvec'] + + mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] + ) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], + 'INSERT_VALUES', solver_objs['X_local']] + ) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + ) + + vec_get_array_y = petsc_call( + 'VecGetArray', [solver_objs['Y_local'], Byref(y_matvec._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [solver_objs['X_local'], Byref(x_matvec._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] + ) + + vec_restore_array_y = petsc_call( + 'VecRestoreArray', [solver_objs['Y_local'], Byref(y_matvec._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [solver_objs['X_local'], Byref(x_matvec._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + ]) + + # NOTE: Question: I have placed a chunk of the calls in the `stacks` argument + # of the `CallableBody` to ensure that these calls precede the `cast` statements. + # The 'casts' depend on the calls, so this order is necessary. By doing this, + # I avoid having to manually construct the 'casts' and can allow Devito to handle + # their construction. Are there any potential issues with this approach? + body = [body, + vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end] + + stacks = ( + mat_get_dm, + dm_get_app_context, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + vec_get_array_y, + vec_get_array_x, + dm_get_local_info + ) + + matvec_body = CallableBody( + List(body=body), + init=tuple([petsc_func_begin_user]), + stacks=stacks, + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])) + + # Replace data with pointer to data in struct + subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + matvec_body = Uxreplace(subs).visit(matvec_body) + + self._struct_params.extend(struct.usr_ctx) + + return matvec_body + + def make_formfunc(self, injectsolve, objs, solver_objs): + target = injectsolve.expr.rhs.target + # Compile formfunc `eqns` into an IET via recursive compilation + irs_formfunc, _ = self.rcompile(injectsolve.expr.rhs.formfuncs, + options={'mpi': False}) + body_formfunc = self.create_formfunc_body(injectsolve, irs_formfunc.uiet.body, + solver_objs, objs) + + formfunc_callback = PETScCallable( + 'FormFunction_%s' % target.name, body_formfunc, retval=objs['err'], + parameters=(solver_objs['snes'], solver_objs['X_global'], + solver_objs['Y_global']), unused_parameters=(solver_objs['dummy']) + ) + return formfunc_callback + + def create_formfunc_body(self, injectsolve, body, solver_objs, objs): + linsolveexpr = injectsolve.expr.rhs + + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + + struct = build_petsc_struct(body, 'formfunc', liveness='eager') + + y_formfunc = linsolveexpr.arrays['y_formfunc'] + x_formfunc = linsolveexpr.arrays['x_formfunc'] + + snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] + ) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], + 'INSERT_VALUES', solver_objs['X_local']] + ) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + ) + + vec_get_array_y = petsc_call( + 'VecGetArray', [solver_objs['Y_local'], Byref(y_formfunc._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [solver_objs['X_local'], Byref(x_formfunc._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] + ) + + vec_restore_array_y = petsc_call( + 'VecRestoreArray', [solver_objs['Y_local'], Byref(y_formfunc._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [solver_objs['X_local'], Byref(x_formfunc._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + ]) + + body = [body, + vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end] + + stacks = ( + snes_get_dm, + dm_get_app_context, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + vec_get_array_y, + vec_get_array_x, + dm_get_local_info + ) + + formfunc_body = CallableBody( + List(body=body), + init=tuple([petsc_func_begin_user]), + stacks=stacks, + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])) + + # Replace data with pointer to data in struct + subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + formfunc_body = Uxreplace(subs).visit(formfunc_body) + + self._struct_params.extend(struct.usr_ctx) + + return formfunc_body + + def make_formrhs(self, injectsolve, objs, solver_objs): + target = injectsolve.expr.rhs.target + # Compile formrhs `eqns` into an IET via recursive compilation + irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, + options={'mpi': False}) + body_formrhs = self.create_formrhs_body(injectsolve, irs_formrhs.uiet.body, + solver_objs, objs) + + formrhs_callback = Callable( + 'FormRHS_%s' % target.name, body_formrhs, retval=objs['err'], + parameters=( + solver_objs['snes'], solver_objs['b_local'] + ) + ) + + return formrhs_callback + + def create_formrhs_body(self, injectsolve, body, solver_objs, objs): + linsolveexpr = injectsolve.expr.rhs + + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + + snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + + b_arr = linsolveexpr.arrays['b_tmp'] + + vec_get_array = petsc_call( + 'VecGetArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] + ) + + struct = build_petsc_struct(body, 'formrhs', liveness='eager') + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] + ) + + vec_restore_array = petsc_call( + 'VecRestoreArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + ) + + body = [body, + vec_restore_array] + + stacks = ( + snes_get_dm, + dm_get_app_context, + vec_get_array, + dm_get_local_info, + ) + + formrhs_body = CallableBody( + List(body=[body]), + init=tuple([petsc_func_begin_user]), + stacks=stacks, + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])) + + # Replace data with pointer to data in struct + subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + formrhs_body = Uxreplace(subs).visit(formrhs_body) + + self._struct_params.extend(struct.usr_ctx) + + return formrhs_body + + def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): + target = injectsolve.expr.rhs.target + + dmda = objs['da_so_%s' % target.space_order] + + rhs_call = petsc_call(rhs_callback.name, list(rhs_callback.parameters)) + + dm_local_to_global_x = petsc_call( + 'DMLocalToGlobal', [dmda, solver_objs['x_local'], 'INSERT_VALUES', + solver_objs['x_global']] + ) + + dm_local_to_global_b = petsc_call( + 'DMLocalToGlobal', [dmda, solver_objs['b_local'], 'INSERT_VALUES', + solver_objs['b_global']] + ) + + snes_solve = petsc_call('SNESSolve', [ + solver_objs['snes'], solver_objs['b_global'], solver_objs['x_global']] + ) + + dm_global_to_local_x = petsc_call('DMGlobalToLocal', [ + dmda, solver_objs['x_global'], 'INSERT_VALUES', solver_objs['x_local']] + ) + + calls = (rhs_call, + dm_local_to_global_x, + dm_local_to_global_b, + snes_solve, + dm_global_to_local_x) + + return calls + + +def build_petsc_struct(iet, name, liveness): + # Place all context data required by the shell routines + # into a PETScStruct + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases|writes').visit(iet) + usr_ctx = [data for data in basics if data not in avoid] + return PETScStruct(name, usr_ctx, liveness=liveness) + + +Null = Macro('NULL') +void = 'void' + + +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 5bdad59b43..6f30b57f10 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,56 +1,61 @@ from functools import singledispatch -from sympy import simplify +import sympy -from devito.finite_differences.differentiable import Add, Mul, EvalDerivative, diffify +from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq +from devito.types.equation import InjectSolveEq from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions from devito.symbolics import uxreplace -from devito.petsc.types import PETScArray, LinearSolveExpr, MatVecEq, RHSEq +from devito.petsc.types import LinearSolveExpr, PETScArray + __all__ = ['PETScSolve'] def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - # TODO: Add check for time dimensions and utilise implicit dimensions. - - is_time_dep = any(dim.is_Time for dim in target.dimensions) - # TODO: Current assumption is rhs is part of pde that remains - # constant at each timestep. Need to insert function to extract this from eq. - y_matvec, x_matvec, b_tmp = [ - PETScArray(name=f'{prefix}_{target.name}', - dtype=target.dtype, - dimensions=target.space_dimensions, - shape=target.grid.shape, - liveness='eager', - halo=target.halo[1:] if is_time_dep else target.halo) - for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] - b, F_target = separate_eqn(eq, target) + prefixes = ['y_matvec', 'x_matvec', 'y_formfunc', 'x_formfunc', 'b_tmp'] - # Args were updated so need to update target to enable uxreplace on F_target - new_target = {f for f in retrieve_functions(F_target) if - f.function == target.function} - assert len(new_target) == 1 # Sanity check: only one target expected - new_target = new_target.pop() + arrays = { + p: PETScArray(name='%s_%s' % (p, target.name), + dtype=target.dtype, + dimensions=target.space_dimensions, + shape=target.grid.shape, + liveness='eager', + halo=[target.halo[d] for d in target.space_dimensions], + space_order=target.space_order) + for p in prefixes + } + + b, F_target = separate_eqn(eq, target) # TODO: Current assumption is that problem is linear and user has not provided # a jacobian. Hence, we can use F_target to form the jac-vec product - matvecaction = MatVecEq( - y_matvec, LinearSolveExpr(uxreplace(F_target, {new_target: x_matvec}), - target=target, solver_parameters=solver_parameters), - subdomain=eq.subdomain) + matvecaction = Eq(arrays['y_matvec'], + uxreplace(F_target, {target: arrays['x_matvec']}), + subdomain=eq.subdomain) - # Part of pde that remains constant at each timestep - rhs = RHSEq(b_tmp, LinearSolveExpr(b, target=target, - solver_parameters=solver_parameters), subdomain=eq.subdomain) + formfunction = Eq(arrays['y_formfunc'], + uxreplace(F_target, {target: arrays['x_formfunc']}), + subdomain=eq.subdomain) + + rhs = Eq(arrays['b_tmp'], b, subdomain=eq.subdomain) + + # Placeholder equation for inserting calls to the solver + inject_solve = InjectSolveEq(arrays['b_tmp'], LinearSolveExpr( + b, target=target, solver_parameters=solver_parameters, matvecs=[matvecaction], + formfuncs=[formfunction], formrhs=[rhs], arrays=arrays, + ), subdomain=eq.subdomain) if not bcs: - return [matvecaction, rhs] + return [inject_solve] + # NOTE: BELOW IS NOT FULLY TESTED/IMPLEMENTED YET bcs_for_matvec = [] + bcs_for_formfunc = [] + bcs_for_rhs = [] for bc in bcs: # TODO: Insert code to distiguish between essential and natural # boundary conditions since these are treated differently within @@ -59,14 +64,24 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): # (and move to rhs) but for now, they are included since this # is not trivial to implement when using DMDA # NOTE: Below is temporary -> Just using this as a palceholder for - # the actual BC implementation for the matvec callback - new_rhs = bc.rhs.subs(target, x_matvec) - bc_rhs = LinearSolveExpr( - new_rhs, target=target, solver_parameters=solver_parameters - ) - bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) - - return [matvecaction] + bcs_for_matvec + [rhs] + # the actual BC implementation + centre = centre_stencil(F_target, target) + bcs_for_matvec.append(Eq(arrays['y_matvec'], + centre.subs(target, arrays['x_matvec']), + subdomain=bc.subdomain)) + bcs_for_formfunc.append(Eq(arrays['y_formfunc'], + 0., subdomain=bc.subdomain)) + # NOTE: Temporary + bcs_for_rhs.append(Eq(arrays['b_tmp'], 0., subdomain=bc.subdomain)) + + inject_solve = InjectSolveEq(arrays['b_tmp'], LinearSolveExpr( + b, target=target, solver_parameters=solver_parameters, + matvecs=[matvecaction]+bcs_for_matvec, + formfuncs=[formfunction]+bcs_for_formfunc, formrhs=[rhs], + arrays=arrays, + ), subdomain=eq.subdomain) + + return [inject_solve] def separate_eqn(eqn, target): @@ -76,45 +91,37 @@ def separate_eqn(eqn, target): """ zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) tmp = eval_time_derivatives(zeroed_eqn.lhs) - b = remove_target(tmp, target) - F_target = diffify(simplify(tmp - b)) + b, F_target = remove_target(tmp, target) return -b, F_target @singledispatch def remove_target(expr, target): - return 0 if expr == target else expr + return (0, expr) if expr == target else (expr, 0) -@remove_target.register(Add) -@remove_target.register(EvalDerivative) +@remove_target.register(sympy.Add) def _(expr, target): if not expr.has(target): - return expr + return (expr, 0) - args = [remove_target(a, target) for a in expr.args] - return expr.func(*args, evaluate=False) + args_b, args_F = zip(*(remove_target(a, target) for a in expr.args)) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) @remove_target.register(Mul) def _(expr, target): if not expr.has(target): - return expr + return (expr, 0) - args = [] - for a in expr.args: - if not a.has(target): - args.append(a) - else: - a1 = remove_target(a, target) - args.append(a1) - - return expr.func(*args, evaluate=False) + args_b, args_F = zip(*[remove_target(a, target) if a.has(target) + else (a, a) for a in expr.args]) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) @remove_target.register(Derivative) def _(expr, target): - return 0 if expr.has(target) else expr + return (0, expr) if expr.has(target) else (expr, 0) @singledispatch @@ -127,8 +134,7 @@ def centre_stencil(expr, target): return expr if expr == target else 0 -@centre_stencil.register(Add) -@centre_stencil.register(EvalDerivative) +@centre_stencil.register(sympy.Add) def _(expr, target): if not expr.has(target): return 0 diff --git a/devito/petsc/types.py b/devito/petsc/types.py index 62c69f666a..72a0faa1e8 100644 --- a/devito/petsc/types.py +++ b/devito/petsc/types.py @@ -1,17 +1,19 @@ import sympy -from functools import cached_property -import numpy as np +<<<<<<< HEAD +from devito.tools import Reconstructable, sympy_mutex +======= from devito.tools import CustomDtype from devito.types import LocalObject, Eq, CompositeObject from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction, Symbol +from devito.types.basic import AbstractFunction, AbstractSymbol from devito.finite_differences.tools import fd_weights_registry -from devito.tools import Reconstructable, dtype_to_ctype -from devito.symbolics import FieldFromComposite -from devito.types.basic import IndexedBase +from devito.tools import dtype_to_ctype, Reconstructable, sympy_mutex +from devito.symbolics import FieldFromComposite, Byref + +from devito.petsc.utils import petsc_call class DM(LocalObject): @@ -28,6 +30,15 @@ def __init__(self, *args, stencil_width=None, **kwargs): def stencil_width(self): return self._stencil_width + @property + def info(self): + return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') + + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('DMDestroy', [Byref(self.function)]) + class Mat(LocalObject): """ @@ -35,13 +46,30 @@ class Mat(LocalObject): """ dtype = CustomDtype('Mat') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('MatDestroy', [Byref(self.function)]) + + +class LocalVec(LocalObject): + """ + PETSc Vector object (Vec). + """ + dtype = CustomDtype('Vec') + -class Vec(LocalObject): +class GlobalVec(LocalObject): """ PETSc Vector object (Vec). """ dtype = CustomDtype('Vec') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('VecDestroy', [Byref(self.function)]) + class PetscMPIInt(LocalObject): """ @@ -65,6 +93,11 @@ class SNES(LocalObject): """ dtype = CustomDtype('SNES') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('SNESDestroy', [Byref(self.function)]) + class PC(LocalObject): """ @@ -97,6 +130,10 @@ class PetscErrorCode(LocalObject): dtype = CustomDtype('PetscErrorCode') +class DummyArg(LocalObject): + dtype = CustomDtype('void', modifier='*') + + class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent @@ -114,7 +151,8 @@ class PETScArray(ArrayBasic, Differentiable): _default_fd = 'taylor' __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients')) + ('dimensions', 'shape', 'liveness', 'coefficients', + 'space_order')) def __init_finalize__(self, *args, **kwargs): @@ -126,6 +164,7 @@ def __init_finalize__(self, *args, **kwargs): raise ValueError("coefficients must be one of %s" " not %s" % (str(fd_weights_registry), self._coefficients)) self._shape = kwargs.get('shape') + self._space_order = kwargs.get('space_order', 1) @classmethod def __dtype_setup__(cls, **kwargs): @@ -140,6 +179,10 @@ def coefficients(self): def shape(self): return self._shape + @property + def space_order(self): + return self._space_order + @cached_property def _shape_with_inhalo(self): """ @@ -172,37 +215,28 @@ def shape_allocated(self): @cached_property def _C_ctype(self): - return CustomDtype(self.petsc_type, modifier='*') - - @property - def petsc_type(self): - return dtype_to_petsctype(self._dtype) - - @property - def dtype(self): - return CustomDtype(self.petsc_type) + # NOTE: Reverting to using float/double instead of PetscScalar for + # simplicity when opt='advanced'. Otherwise, Temp objects must also + # be converted to PetscScalar. Additional tests are needed to + # ensure this approach is fine. Previously, issues arose from + # mismatches between precision of Function objects in Devito and the + # precision of the PETSc configuration. + # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables + # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of + # the user's PETSc configuration. + return POINTER(dtype_to_ctype(self.dtype)) @property def symbolic_shape(self): field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] + FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] # Reverse it since DMDA is setup backwards to Devito dimensions. return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - @cached_property - def dmda_info(self): - # To access the local grid info via the DM in - # PETSc you use DMDAGetLocalInfo(da, &info) - # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz - return DMDALocalInfo(name='info', liveness='eager') - - @cached_property - def indexed(self): - return PETScIndexedData(self.name, shape=self._shape, function=self.function) - - -class PETScIndexedData(IndexedBase): - pass + @property + def dmda(self): + name = 'da_so_%s' % self.space_order + return DM(name=name, liveness='eager', stencil_width=self.space_order) def dtype_to_petsctype(dtype): @@ -216,85 +250,25 @@ def dtype_to_petsctype(dtype): }[dtype] -class MatVecEq(Eq): - """ - Represents the mathematical expression of applying a linear - operator to a vector. This is a key component - for running matrix-free solvers. - """ +class InjectSolveEq(Eq): pass - - -class RHSEq(Eq): - """ - Represents the mathematical expression of building the - rhs of a linear system. - """ - pass - - -class MockEq(Eq): - """ - Represents a mock/placeholder equation to ensure distinct iteration loops. - - For example, the mat-vec action iteration loop is to be isolated from the - expression loop used to build the RHS of the linear system. This separation - facilitates the utilisation of the mat-vec iteration loop in callback functions - created at the IET level. - """ - pass - - -def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - # TODO: Add check for time dimensions and utilise implicit dimensions. - - is_time_dep = any(dim.is_Time for dim in target.dimensions) - # TODO: Current assumption is rhs is part of pde that remains - # constant at each timestep. Need to insert function to extract this from eq. - y_matvec, x_matvec, b_tmp = [ - PETScArray(name=f'{prefix}_{target.name}', - dtype=target.dtype, - dimensions=target.space_dimensions, - shape=target.grid.shape, - liveness='eager', - halo=target.halo[1:] if is_time_dep else target.halo) - for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] - - # TODO: Extend to rearrange equation for implicit time stepping. - matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), - target=target, solver_parameters=solver_parameters), - subdomain=eq.subdomain) - - # Part of pde that remains constant at each timestep - rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, - solver_parameters=solver_parameters), subdomain=eq.subdomain) - - if not bcs: - return [matvecaction, rhs] - - bcs_for_matvec = [] - for bc in bcs: - # TODO: Insert code to distiguish between essential and natural - # boundary conditions since these are treated differently within - # the solver - # NOTE: May eventually remove the essential bcs from the solve - # (and move to rhs) but for now, they are included since this - # is not trivial to implement when using DMDA - # NOTE: Below is temporary -> Just using this as a palceholder for - # the actual BC implementation for the matvec callback - new_rhs = bc.rhs.subs(target, x_matvec) - bc_rhs = LinearSolveExpr( - new_rhs, target=target, solver_parameters=solver_parameters - ) - bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) - - return [matvecaction] + bcs_for_matvec + [rhs] +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) class LinearSolveExpr(sympy.Function, Reconstructable): __rargs__ = ('expr',) - __rkwargs__ = ('target', 'solver_parameters') +<<<<<<< HEAD +<<<<<<< HEAD + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays') +======= + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', 'formfuncs', 'formrhs') +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays') +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) defaults = { 'ksp_type': 'gmres', @@ -305,7 +279,16 @@ class LinearSolveExpr(sympy.Function, Reconstructable): 'ksp_max_it': 10000 # Maximum iterations } - def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): + def __new__(cls, expr, target=None, solver_parameters=None, +<<<<<<< HEAD +<<<<<<< HEAD + matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): +======= + matvecs=None, formfuncs=None, formrhs=None, **kwargs): +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) if solver_parameters is None: solver_parameters = cls.defaults @@ -313,10 +296,22 @@ def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): for key, val in cls.defaults.items(): solver_parameters[key] = solver_parameters.get(key, val) - obj = super().__new__(cls, expr) + with sympy_mutex: + obj = sympy.Basic.__new__(cls, expr) obj._expr = expr obj._target = target obj._solver_parameters = solver_parameters + obj._matvecs = matvecs + obj._formfuncs = formfuncs + obj._formrhs = formrhs +<<<<<<< HEAD +<<<<<<< HEAD + obj._arrays = arrays +======= +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + obj._arrays = arrays +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) return obj def __repr__(self): @@ -328,7 +323,12 @@ def _sympystr(self, printer): return str(self) def __hash__(self): - return hash(self.target) + return hash(self.expr) + + def __eq__(self, other): + return (isinstance(other, LinearSolveExpr) and + self.expr == other.expr and + self.target == other.target) @property def expr(self): @@ -342,43 +342,88 @@ def target(self): def solver_parameters(self): return self._solver_parameters - func = Reconstructable._rebuild - + @property + def matvecs(self): + return self._matvecs -def petsc_lift(clusters): - """ - Lift the iteration space associated with each PETSc equation. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions. - """ - processed = [] - for c in clusters: + @property + def formfuncs(self): + return self._formfuncs - ispace = c.ispace - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + @property + def formrhs(self): + return self._formrhs - processed.append(c.rebuild(ispace=ispace)) +<<<<<<< HEAD +<<<<<<< HEAD + @property + def formrhs(self): + return self._formrhs +======= +======= + @property + def arrays(self): + return self._arrays - return processed +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) + func = Reconstructable._rebuild class PETScStruct(CompositeObject): __rargs__ = ('name', 'usr_ctx') + __rkwargs__ = ('liveness',) - def __init__(self, name, usr_ctx): + def __init__(self, name, usr_ctx, liveness='lazy'): pfields = [(i._C_name, dtype_to_ctype(i.dtype)) - for i in usr_ctx if isinstance(i, Symbol)] + for i in usr_ctx if isinstance(i, AbstractSymbol)] self._usr_ctx = usr_ctx super().__init__(name, 'MatContext', pfields) +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) + + assert liveness in ['eager', 'lazy'] + self._liveness = liveness @property - def usr_ctx(self): - return self._usr_ctx + def arrays(self): + return self._arrays +<<<<<<< HEAD + func = Reconstructable._rebuild +======= def _arg_values(self, **kwargs): + # TODO: Handle passing `struct dataobj *restrict g1_vec` and similar + # elements within in the struct. It seemed necessary to convert the Numpy array + # into a pointer to ctypes.Struct within MatContext + # (i.e override _arg_finalize for PETScStruct here?), but this approach seemed + # to result in a NIL pointer in the generated code ... any suggestions? + # Alternative: Allow these pointers to appear in the main kernel arguments + # as usual and use ctx->g1_vec = g1_vec; to assign them to the struct? values = super()._arg_values(**kwargs) - for i in self.fields: - setattr(values[self.name]._obj, i, kwargs['args'][i]) + for i in self.usr_ctx: + setattr(values[self.name]._obj, i.name, kwargs['args'][i.name]) return values + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' +<<<<<<< HEAD +<<<<<<< HEAD +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + +======= + +>>>>>>> 9d611239f (compiler: Fix derive_callback_inputs) + @property + def fields(self): + return self._usr_ctx +>>>>>>> 41157b0bc (compiler: Place dataobj pointers inside ctx struct) diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 1d171aad2b..b24ad56a3d 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -1,14 +1,14 @@ import os -from devito.ir.equations import OpMatVec, OpRHS +from devito.ir.equations import OpInjectSolve from devito.tools import memoized_func -from devito.ir.iet import Call -from devito.petsc.iet.nodes import MatVecAction, RHSLinearSystem +from devito.ir.iet import Call, FindSymbols +from devito.petsc.iet.nodes import PETScCallable, InjectSolveDummy +from devito.petsc.types import PETScStruct # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. -petsc_iet_mapper = {OpMatVec: MatVecAction, - OpRHS: RHSLinearSystem} +petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} solver_mapper = { @@ -25,7 +25,6 @@ def get_petsc_dir(): petsc_dir = os.environ.get(i) if petsc_dir: return petsc_dir - # TODO: Raise error if PETSC_DIR is not set return None @@ -60,13 +59,3 @@ def core_metadata(): 'lib_dirs': lib_dir, 'ldflags': ('-Wl,-rpath,%s' % lib_dir) } - - -def petsc_call(specific_call, call_args): - general_call = 'PetscCall' - return Call(general_call, [Call(specific_call, arguments=call_args)]) - - -def petsc_call_mpi(specific_call, call_args): - general_call = 'PetscCallMPI' - return Call(general_call, [Call(specific_call, arguments=call_args)]) diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index ae70971c10..ba41309fd9 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -123,4 +123,3 @@ USER app EXPOSE 8888 ENTRYPOINT ["/docker-entrypoint.sh"] CMD ["/jupyter"] - diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc new file mode 100644 index 0000000000..782e0ac340 --- /dev/null +++ b/docker/Dockerfile.petsc @@ -0,0 +1,49 @@ +############################################################## +# This Dockerfile builds a base image to run Devito with PETSc +############################################################# + +# Base image +FROM ubuntu:22.04 as base + +ENV DEBIAN_FRONTEND noninteractive + +# Install python +RUN apt-get update && \ + apt-get install -y dh-autoreconf python3-venv python3-dev python3-pip libopenblas-serial-dev git pkgconf mpich + +# Install for basic base not containing it +RUN apt-get install -y vim wget git flex libnuma-dev tmux \ + numactl hwloc curl \ + autoconf libtool build-essential procps + +# Install tmpi +RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi + +# Install OpenGL library, necessary for the installation of GemPy +RUN apt-get install -y libgl1-mesa-glx + +# Install numpy from source +RUN python3 -m venv /venv && \ + /venv/bin/pip install --no-cache-dir --upgrade pip && \ + /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ + rm -rf ~/.cache/pip + +# Create a directory for PETSc +RUN mkdir -p /home/petsc + +# Clone and install PETSc +RUN cd /home/petsc \ + && git clone -b release https://gitlab.com/petsc/petsc.git petsc \ + && cd petsc \ + && git pull \ + && ./configure --with-fortran-bindings=0 --with-openblas-include=$(pkgconf --variable=includedir openblas) --with-openblas-lib=$(pkgconf --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build \ + && make all + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +EXPOSE 8888 +CMD ["/bin/bash"] + +ENV PETSC_ARCH="devito_build" +ENV PETSC_DIR="/home/petsc/petsc" diff --git a/tests/test_iet.py b/tests/test_iet.py index 1c36830084..eac5316087 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -10,8 +10,8 @@ from devito.ir.iet import (Call, Callable, Conditional, DeviceCall, DummyExpr, Iteration, List, KernelLaunch, Lambda, ElementalFunction, CGen, FindSymbols, filter_iterations, make_efunc, - retrieve_iteration_tree, Transformer, Callback, FindNodes, - Definition) + retrieve_iteration_tree, Transformer, Callback, + Definition, FindNodes) from devito.ir import SymbolRegistry from devito.passes.iet.engine import Graph from devito.passes.iet.languages.C import CDataManager @@ -131,19 +131,26 @@ def test_find_symbols_nested(mode, expected): def test_callback_cgen(): + class FunctionPtr(Callback): + @property + def callback_form(self): + param_types = ', '.join([str(t) for t in + self.param_types]) + return "(%s (*)(%s))%s" % (self.retval, param_types, self.name) + a = Symbol('a') b = Symbol('b') foo0 = Callable('foo0', Definition(a), 'void', parameters=[b]) - foo0_arg = Callback(foo0.name, foo0.retval, 'int') + foo0_arg = FunctionPtr(foo0.name, foo0.retval, 'int') code0 = CGen().visit(foo0_arg) assert str(code0) == '(void (*)(int))foo0' # Test nested calls with a Callback as an argument. - call = Call('foo1', [ - Call('foo2', [foo0_arg]) + call = Call('foo2', [ + Call('foo1', [foo0_arg]) ]) code1 = CGen().visit(call) - assert str(code1) == 'foo1(foo2((void (*)(int))foo0));' + assert str(code1) == 'foo2(foo1((void (*)(int))foo0));' callees = FindNodes(Call).visit(call) assert len(callees) == 3 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 9ce8b209e6..c639e26783 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -5,14 +5,14 @@ from conftest import skipif from devito import Grid, Function, TimeFunction, Eq, Operator, switchconfig from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, - FindNodes, - PointerCast, retrieve_iteration_tree) + FindNodes, PointerCast, retrieve_iteration_tree) +from devito.types import Constant, CCompositeObject from devito.passes.iet.languages.C import CDataManager -from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, +from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, - LinearSolveExpr, PETScStruct) + LinearSolveExpr) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil -from devito.petsc.iet.nodes import MatVecAction, RHSLinearSystem +from devito.petsc.iet.nodes import Expression @skipif('petsc') @@ -22,7 +22,7 @@ def test_petsc_local_object(): """ lo0 = DM('da', stencil_width=1) lo1 = Mat('A') - lo2 = Vec('x') + lo2 = LocalVec('x') lo3 = PetscMPIInt('size') lo4 = KSP('ksp') lo5 = PC('pc') @@ -68,11 +68,11 @@ def test_petsc_functions(): expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1) - assert str(defn0) == 'PetscScalar*restrict ptr0_vec;' - assert str(defn1) == 'const PetscScalar*restrict ptr1_vec;' - assert str(defn2) == 'const PetscScalar*restrict ptr2_vec;' - assert str(defn3) == 'PetscInt*restrict ptr3_vec;' - assert str(defn4) == 'const PetscInt*restrict ptr4_vec;' + assert str(defn0) == 'float *restrict ptr0_vec;' + assert str(defn1) == 'const float *restrict ptr1_vec;' + assert str(defn2) == 'const double *restrict ptr2_vec;' + assert str(defn3) == 'int *restrict ptr3_vec;' + assert str(defn4) == 'const long *restrict ptr4_vec;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' @@ -123,16 +123,18 @@ def test_petsc_solve(): matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_f'] - action_expr = FindNodes(MatVecAction).visit(matvec_callback[0]) - rhs_expr = FindNodes(RHSLinearSystem).visit(op) + formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS_f'] + + action_expr = FindNodes(Expression).visit(matvec_callback[0]) + rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) assert str(action_expr[-1].expr.rhs) == \ - 'ctx->h_x**(-2)*x_matvec_f[x + 1, y + 2]' + \ - ' - 2.0*ctx->h_x**(-2)*x_matvec_f[x + 2, y + 2]' + \ - ' + ctx->h_x**(-2)*x_matvec_f[x + 3, y + 2]' + \ - ' + ctx->h_y**(-2)*x_matvec_f[x + 2, y + 1]' + \ - ' - 2.0*ctx->h_y**(-2)*x_matvec_f[x + 2, y + 2]' + \ - ' + ctx->h_y**(-2)*x_matvec_f[x + 2, y + 3]' + 'matvec->h_x**(-2)*x_matvec_f[x + 1, y + 2]' + \ + ' - 2.0*matvec->h_x**(-2)*x_matvec_f[x + 2, y + 2]' + \ + ' + matvec->h_x**(-2)*x_matvec_f[x + 3, y + 2]' + \ + ' + matvec->h_y**(-2)*x_matvec_f[x + 2, y + 1]' + \ + ' - 2.0*matvec->h_y**(-2)*x_matvec_f[x + 2, y + 2]' + \ + ' + matvec->h_y**(-2)*x_matvec_f[x + 2, y + 3]' assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' @@ -142,7 +144,7 @@ def test_petsc_solve(): assert op.arguments().get('y_M') == 1 assert op.arguments().get('x_M') == 1 - assert len(retrieve_iteration_tree(op)) == 1 + assert len(retrieve_iteration_tree(op)) == 0 # TODO: Remove pragmas from PETSc callback functions assert len(matvec_callback[0].parameters) == 3 @@ -172,13 +174,9 @@ def test_multiple_petsc_solves(): callable_roots = [meta_call.root for meta_call in op._func_table.values()] - assert len(callable_roots) == 2 - - structs = [i for i in op.parameters if isinstance(i, PETScStruct)] - - # Only create 1 struct per Grid/DMDA - assert len(structs) == 1 - assert len(structs[0].fields) == 6 + # One FormRHS, one MatShellMult and one FormFunction per solve + # One PopulateMatContext for all solves + assert len(callable_roots) == 7 @skipif('petsc') @@ -194,40 +192,27 @@ def test_petsc_cast(): arr1 = PETScArray(name='arr1', dimensions=g1.dimensions, shape=g1.shape) arr2 = PETScArray(name='arr2', dimensions=g2.dimensions, shape=g2.shape) + arr3 = PETScArray(name='arr3', dimensions=g1.dimensions, + shape=g1.shape, space_order=4) + # Casts will be explictly generated and placed at specific locations in the C code, # specifically after various other PETSc calls have been executed. cast0 = PointerCast(arr0) cast1 = PointerCast(arr1) cast2 = PointerCast(arr2) + cast3 = PointerCast(arr3) assert str(cast0) == \ - 'PetscScalar (*restrict arr0) = (PetscScalar (*)) arr0_vec;' + 'float (*restrict arr0) = (float (*)) arr0_vec;' assert str(cast1) == \ - 'PetscScalar (*restrict arr1)[info.gxm] = (PetscScalar (*)[info.gxm]) arr1_vec;' + 'float (*restrict arr1)[da_so_1_info.gxm] = ' + \ + '(float (*)[da_so_1_info.gxm]) arr1_vec;' assert str(cast2) == \ - 'PetscScalar (*restrict arr2)[info.gym][info.gxm] = ' + \ - '(PetscScalar (*)[info.gym][info.gxm]) arr2_vec;' - - -@skipif('petsc') -def test_no_automatic_cast(): - """ - Verify that the compiler does not automatically generate casts for PETScArrays. - They will be generated at specific points within the C code, particularly after - other PETSc calls, rather than necessarily at the top of the Kernel. - """ - grid = Grid((2, 2)) - - f = Function(name='f', grid=grid, space_order=2) - - arr = PETScArray(name='arr', dimensions=f.dimensions, shape=f.shape) - - eqn = Eq(arr, f.laplace) - - with switchconfig(openmp=False): - op = Operator(eqn, opt='noop') - - assert len(op.body.casts) == 1 + 'float (*restrict arr2)[da_so_1_info.gym][da_so_1_info.gxm] = ' + \ + '(float (*)[da_so_1_info.gym][da_so_1_info.gxm]) arr2_vec;' + assert str(cast3) == \ + 'float (*restrict arr3)[da_so_4_info.gxm] = ' + \ + '(float (*)[da_so_4_info.gxm]) arr3_vec;' @skipif('petsc') @@ -319,13 +304,56 @@ def test_cinterface_petsc_struct(): assert 'include "%s.h"' % name in ccode # The public `struct MatContext` only appears in the header file - assert any(isinstance(i, PETScStruct) for i in op.parameters) assert 'struct MatContext\n{' not in ccode assert 'struct MatContext\n{' in hcode @skipif('petsc') -def test_separate_eqn(): +@pytest.mark.parametrize('eqn, target, expected', [ + ('Eq(f1.laplace, g1)', + 'f1', ('g1(x, y)', 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), + ('Eq(g1, f1.laplace)', + 'f1', ('-g1(x, y)', '-Derivative(f1(x, y), (x, 2)) - Derivative(f1(x, y), (y, 2))')), + ('Eq(g1, f1.laplace)', 'g1', + ('Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))', 'g1(x, y)')), + ('Eq(f1 + f1.laplace, g1)', 'f1', ('g1(x, y)', + 'f1(x, y) + Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), + ('Eq(g1.dx + f1.dx, g1)', 'f1', + ('g1(x, y) - Derivative(g1(x, y), x)', 'Derivative(f1(x, y), x)')), + ('Eq(g1.dx + f1.dx, g1)', 'g1', + ('-Derivative(f1(x, y), x)', '-g1(x, y) + Derivative(g1(x, y), x)')), + ('Eq(f1 * g1.dx, g1)', 'g1', ('0', 'f1(x, y)*Derivative(g1(x, y), x) - g1(x, y)')), + ('Eq(f1 * g1.dx, g1)', 'f1', ('g1(x, y)', 'f1(x, y)*Derivative(g1(x, y), x)')), + ('Eq((f1 * g1.dx).dy, f1)', 'f1', + ('0', '-f1(x, y) + Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')), + ('Eq((f1 * g1.dx).dy, f1)', 'g1', + ('f1(x, y)', 'Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')), + ('Eq(f2.laplace, g2)', 'g2', + ('-Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))', + '-g2(t, x, y)')), + ('Eq(f2.laplace, g2)', 'f2', ('g2(t, x, y)', + 'Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')), + ('Eq(f2.laplace, f2)', 'f2', ('0', + '-f2(t, x, y) + Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')), + ('Eq(f2*g2, f2)', 'f2', ('0', 'f2(t, x, y)*g2(t, x, y) - f2(t, x, y)')), + ('Eq(f2*g2, f2)', 'g2', ('f2(t, x, y)', 'f2(t, x, y)*g2(t, x, y)')), + ('Eq(g2*f2.laplace, f2)', 'g2', ('f2(t, x, y)', + '(Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2)))*g2(t, x, y)')), + ('Eq(f2.forward, f2)', 'f2.forward', ('f2(t, x, y)', 'f2(t + dt, x, y)')), + ('Eq(f2.forward, f2)', 'f2', ('-f2(t + dt, x, y)', '-f2(t, x, y)')), + ('Eq(f2.forward.laplace, f2)', 'f2.forward', ('f2(t, x, y)', + 'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')), + ('Eq(f2.forward.laplace, f2)', 'f2', + ('-Derivative(f2(t + dt, x, y), (x, 2)) - Derivative(f2(t + dt, x, y), (y, 2))', + '-f2(t, x, y)')), + ('Eq(f2.laplace + f2.forward.laplace, g2)', 'f2.forward', + ('g2(t, x, y) - Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))', + 'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')), + ('Eq(g2.laplace, f2 + g2.forward)', 'g2.forward', + ('f2(t, x, y) - Derivative(g2(t, x, y), (x, 2)) - Derivative(g2(t, x, y), (y, 2))', + '-g2(t + dt, x, y)')) +]) +def test_separate_eqn(eqn, target, expected): """ Test the separate_eqn function. @@ -335,35 +363,19 @@ def test_separate_eqn(): """ grid = Grid((2, 2)) - f1 = Function(name='f1', grid=grid, space_order=2) - g1 = Function(name='g1', grid=grid, space_order=2) - eq1 = Eq(f1.laplace, g1) - b1, F1 = separate_eqn(eq1, f1) - assert str(b1) == 'g1(x, y)' - assert str(F1) == 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))' - - f2 = TimeFunction(name='f2', grid=grid, space_order=2) - g2 = TimeFunction(name='g2', grid=grid, space_order=2) - eq2 = Eq(f2.dt, f2.laplace + g2) - b2, F2 = separate_eqn(eq2, f2.forward) - assert str(b2) == 'g2(t, x, y) + Derivative(f2(t, x, y), (x, 2))' + \ - ' + Derivative(f2(t, x, y), (y, 2)) + f2(t, x, y)/dt' - assert str(F2) == 'f2(t + dt, x, y)/dt' - - # Implicit Time Stepping - eqn3 = Eq(f2.dt, f2.forward.laplace + g2) - b3, F3 = separate_eqn(eqn3, f2.forward) - assert str(b3) == 'g2(t, x, y) + f2(t, x, y)/dt' - assert str(F3) == '-Derivative(f2(t + dt, x, y), (x, 2)) - ' + \ - 'Derivative(f2(t + dt, x, y), (y, 2)) + f2(t + dt, x, y)/dt' - - # Semi-implicit Time Stepping - eqn4 = Eq(f2.dt, f2.forward.laplace + f2.laplace + g2) - b4, F4 = separate_eqn(eqn4, f2.forward) - assert str(b4) == 'g2(t, x, y) + Derivative(f2(t, x, y), (x, 2))' + \ - ' + Derivative(f2(t, x, y), (y, 2)) + f2(t, x, y)/dt' - assert str(F4) == '-Derivative(f2(t + dt, x, y), (x, 2)) -' + \ - ' Derivative(f2(t + dt, x, y), (y, 2)) + f2(t + dt, x, y)/dt' + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F @skipif('petsc') @@ -403,12 +415,141 @@ def test_centre_stencil(expr, so, target, expected): f1 = Function(name='f1', grid=grid, space_order=so) # noqa g1 = Function(name='g1', grid=grid, space_order=so) # noqa - h1 = Function(name='h1', grid=grid, space_order=so) # noqa f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - h2 = TimeFunction(name='h2', grid=grid, space_order=so) # noqa centre = centre_stencil(eval(expr), eval(target)) assert str(centre) == expected + + +@skipif('petsc') +def test_callback_arguments(): + """ + Test the arguments of each callback function. + """ + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=2) + g1 = Function(name='g1', grid=grid, space_order=2) + + eqn1 = Eq(f1.laplace, g1) + + petsc1 = PETScSolve(eqn1, f1) + + with switchconfig(openmp=False): + op = Operator(petsc1) + + mv = op._func_table['MyMatShellMult_f1'].root + ff = op._func_table['FormFunction_f1'].root + + assert len(mv.parameters) == 3 + assert len(ff.parameters) == 4 + + assert str(mv.parameters) == '(J_f1, X_global_f1, Y_global_f1)' + assert str(ff.parameters) == '(snes_f1, X_global_f1, Y_global_f1, dummy_f1)' + + +@skipif('petsc') +def test_petsc_struct(): + + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=2) + g1 = Function(name='g1', grid=grid, space_order=2) + + mu1 = Constant(name='mu1', value=2.0) + mu2 = Constant(name='mu2', value=2.0) + + eqn1 = Eq(f1.laplace, g1*mu1) + petsc1 = PETScSolve(eqn1, f1) + + eqn2 = Eq(f1, g1*mu2) + + with switchconfig(openmp=False): + op = Operator([eqn2] + petsc1) + + arguments = op.arguments() + + # Check mu1 and mu2 in arguments + assert 'mu1' in arguments + assert 'mu2' in arguments + + # Check mu1 and mu2 in op.parameters + assert mu1 in op.parameters + assert mu2 in op.parameters + + # Check PETSc struct not in op.parameters + assert all(not isinstance(i, CCompositeObject) for i in op.parameters) + + +@skipif('petsc') +@pytest.mark.parallel(mode=1) +def test_apply(mode): + + grid = Grid(shape=(13, 13), dtype=np.float64) + + pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) + rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + mu = Constant(name='mu', value=2.0) + + eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) + + petsc = PETScSolve(eqn, pn) + + # Build the op + with switchconfig(openmp=False): + op = Operator(petsc) + + # Check the Operator runs without errors. Not verifying output for + # now. Need to consolidate BC implementation + op.apply() + + # Verify that users can override `mu` + mu_new = Constant(name='mu_new', value=4.0) + op.apply(mu=mu_new) + + +@skipif('petsc') +def test_petsc_frees(): + + grid = Grid((2, 2)) + + f = Function(name='f', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + eqn = Eq(f.laplace, g) + petsc = PETScSolve(eqn, f) + + with switchconfig(openmp=False): + op = Operator(petsc) + + frees = op.body.frees + + # Check the frees appear in the following order + assert str(frees[0]) == 'PetscCall(VecDestroy(&(b_global_f)));' + assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_f)));' + assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_f)));' + assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_f)));' + assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_so_2)));' + + +@skipif('petsc') +def test_calls_to_callbacks(): + + grid = Grid((2, 2)) + + f = Function(name='f', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + eqn = Eq(f.laplace, g) + petsc = PETScSolve(eqn, f) + + with switchconfig(openmp=False): + op = Operator(petsc) + + ccode = str(op.ccode) + + assert '(void (*)(void))MyMatShellMult_f' in ccode + assert 'PetscCall(SNESSetFunction(snes_f,NULL,FormFunction_f,NULL));' in ccode From 1fe94ea2907baeec7cc9ea8f440c67403eaa4e09 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 23 Jul 2024 21:18:45 +0100 Subject: [PATCH 097/151] compiler: Create mixin for ThreadCallable and PETScCallable --- .github/workflows/examples-mpi.yml | 2 +- .github/workflows/examples.yml | 2 +- .github/workflows/flake8.yml | 2 +- .github/workflows/pytest-core-mpi.yml | 2 +- .github/workflows/pytest-core-nompi.yml | 2 +- .github/workflows/pytest-petsc.yml | 2 +- .github/workflows/tutorials.yml | 2 +- devito/ir/equations/equation.py | 9 +- devito/ir/iet/algorithms.py | 3 +- devito/ir/iet/efunc.py | 4 +- devito/ir/iet/nodes.py | 23 +- devito/ir/iet/visitors.py | 2 +- devito/operator/operator.py | 20 +- devito/passes/clusters/petsc.py | 25 -- devito/passes/iet/__init__.py | 1 - devito/passes/iet/definitions.py | 8 +- devito/passes/iet/engine.py | 6 +- devito/passes/iet/mpi.py | 3 +- devito/passes/iet/petsc.py | 298 ---------------- devito/petsc/clusters.py | 22 +- devito/petsc/iet.py | 372 -------------------- devito/petsc/iet/nodes.py | 18 +- devito/petsc/iet/passes.py | 180 ++++------ devito/petsc/iet/routines.py | 252 ++++++++++---- devito/petsc/iet/utils.py | 83 +++++ devito/petsc/solve.py | 97 +++--- devito/petsc/types.py | 429 ------------------------ devito/petsc/types/__init__.py | 3 + devito/petsc/types/array.py | 117 +++++++ devito/petsc/types/object.py | 171 ++++++++++ devito/petsc/types/types.py | 97 ++++++ devito/petsc/utils.py | 8 - devito/symbolics/printer.py | 3 - devito/types/array.py | 14 - devito/types/basic.py | 31 ++ devito/types/equation.py | 4 + devito/types/object.py | 29 +- docker/Dockerfile.petsc | 49 --- tests/test_iet.py | 2 +- tests/test_petsc.py | 126 ++++++- 40 files changed, 1003 insertions(+), 1520 deletions(-) delete mode 100644 devito/passes/clusters/petsc.py delete mode 100644 devito/passes/iet/petsc.py delete mode 100644 devito/petsc/iet.py create mode 100644 devito/petsc/iet/utils.py delete mode 100644 devito/petsc/types.py create mode 100644 devito/petsc/types/__init__.py create mode 100644 devito/petsc/types/array.py create mode 100644 devito/petsc/types/object.py create mode 100644 devito/petsc/types/types.py delete mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 9df9251991..6a804a27bb 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,7 +17,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 4490f47547..ef14b8450c 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index b61e777454..fe1761f872 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index cc4b8ab683..ab89107047 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 802aa3afb3..a022ed08e1 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index ffb140819e..a556771c26 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index af7f17d255..9e5e1b6cb8 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 6585663c7a..67ed9269a4 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -9,8 +9,9 @@ Stencil, detect_io, detect_accesses) from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict -from devito.types import Eq, Inc, ReduceMax, ReduceMin -from devito.petsc.types import InjectSolveEq +from devito.types import (Eq, Inc, ReduceMax, ReduceMin, + relational_min) +from devito.types.equation import InjectSolveEq __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', 'identity_mapper', 'OpInjectSolve'] @@ -226,13 +227,13 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - # from IPython import embed; embed() + # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - # from IPython import embed; embed() + expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index af079f25d8..9d2db185db 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -4,7 +4,7 @@ Section, HaloSpot, ExpressionBundle) from devito.tools import timed_pass from devito.petsc.types import LinearSolveExpr -from devito.petsc.utils import petsc_iet_mapper +from devito.petsc.iet.utils import petsc_iet_mapper __all__ = ['iet_build'] @@ -24,7 +24,6 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: - # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): diff --git a/devito/ir/iet/efunc.py b/devito/ir/iet/efunc.py index 10aa8920e6..064f02c487 100644 --- a/devito/ir/iet/efunc.py +++ b/devito/ir/iet/efunc.py @@ -1,6 +1,6 @@ from functools import cached_property -from devito.ir.iet.nodes import Call, Callable +from devito.ir.iet.nodes import Call, Callable, FixedArgsCallable from devito.ir.iet.utils import derive_parameters from devito.symbolics import uxreplace from devito.tools import as_tuple @@ -131,7 +131,7 @@ class AsyncCall(Call): pass -class ThreadCallable(Callable): +class ThreadCallable(FixedArgsCallable): """ A Callable executed asynchronously by a thread. diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index c3569fa60b..fa7716aaf8 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1,6 +1,7 @@ """The Iteration/Expression Tree (IET) hierarchy.""" import abc +import ctypes import inspect from functools import cached_property from collections import OrderedDict, namedtuple @@ -28,7 +29,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'CallableBody', 'Transfer', 'Callback'] + 'CallableBody', 'Transfer', 'Callback', 'FixedArgsCallable'] # First-class IET nodes @@ -763,6 +764,15 @@ def defines(self): return self.all_parameters +class FixedArgsCallable(Callable): + + """ + A Callable class that enforces a fixed function signature. + """ + + pass + + class CallableBody(MultiTraversable): """ @@ -1041,8 +1051,8 @@ class Dereference(ExprStmt, Node): The following cases are supported: * `pointer` is a PointerArray or TempFunction, and `pointee` is an Array. - * `pointer` is an ArrayObject representing a pointer to a C struct, and - `pointee` is a field in `pointer`. + * `pointer` is an ArrayObject or CCompositeObject representing a pointer + to a C struct, and `pointee` is a field in `pointer`. """ is_Dereference = True @@ -1061,13 +1071,14 @@ def functions(self): @property def expr_symbols(self): - ret = [self.pointer.indexed] + ret = [] if self.pointer.is_PointerArray or self.pointer.is_TempFunction: - ret.append(self.pointee.indexed) + ret.extend([self.pointer.indexed, self.pointee.indexed]) ret.extend(flatten(i.free_symbols for i in self.pointee.symbolic_shape[1:])) ret.extend(self.pointer.free_symbols) else: - ret.append(self.pointee._C_symbol) + assert issubclass(self.pointer._C_ctype, ctypes._Pointer) + ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) return tuple(filter_ordered(ret)) @property diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 58d69430e9..48570fa89f 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -253,7 +253,7 @@ def _gen_value(self, obj, mode=1, masked=()): strtype = '%s%s' % (strtype, self._restrict_keyword) strtype = ' '.join(qualifiers + [strtype]) - if obj.is_LocalObject and obj._C_modifier is not None and mode == 2: + if obj.is_LocalType and obj._C_modifier is not None and mode == 2: strtype += obj._C_modifier strname = obj._C_name diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 39277fabbc..16c1191d6b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -355,7 +355,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - # from IPython import embed; embed() + return processed # Compilation -- Cluster level @@ -382,8 +382,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Build a sequence of Clusters from a sequence of Eqs clusters = clusterize(expressions, **kwargs) - # Lift iteration spaces surrounding PETSc equations to produce - # distinct iteration loops. clusters = petsc_lift(clusters) # Operation count before specialization @@ -391,6 +389,8 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): clusters = cls._specialize_clusters(clusters, **kwargs) + clusters = petsc_project(clusters) + # Operation count after specialization final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) try: @@ -430,7 +430,6 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters - # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) @@ -500,9 +499,6 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Target-independent optimizations minimize_symbols(graph) - # If necessary, sort frees into a specific order - sort_frees(graph) - return graph.root, graph # Read-only properties exposed to the outside world @@ -521,12 +517,7 @@ def dimensions(self): # During compilation other Dimensions may have been produced dimensions = FindSymbols('dimensions').visit(self) - - # NOTE: Should these dimensions be integrated into self._dimensions instead? - # In which case they would get picked up before this - struct_dims = derive_callback_dims(self._func_table) - - ret.update(d for d in dimensions if d.is_PerfKnob or d in struct_dims) + ret.update(dimensions) ret = tuple(sorted(ret, key=attrgetter('name'))) @@ -534,8 +525,7 @@ def dimensions(self): @cached_property def input(self): - struct_params = derive_struct_inputs(self.parameters) - return tuple(i for i in self.parameters+struct_params if i.is_Input) + return tuple(i for i in self.parameters if i.is_Input) @cached_property def temporaries(self): diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py deleted file mode 100644 index d784d99c0d..0000000000 --- a/devito/passes/clusters/petsc.py +++ /dev/null @@ -1,25 +0,0 @@ -from devito.tools import timed_pass -from devito.petsc import LinearSolveExpr - -__all__ = ['petsc_lift'] - - -@timed_pass() -def petsc_lift(clusters): - """ - Lift the iteration space surrounding each PETSc equation to create - distinct iteration loops. This simplifys the movement of the loops - into specific callback functions generated at the IET level. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions, not the ones that stay inside the main kernel. - """ - processed = [] - for c in clusters: - - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - - return processed diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index cf6a35de90..c09db00c9b 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,4 +8,3 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa -from .petsc import * # noqa diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 62c24b9c78..59de8f95c7 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -333,6 +333,10 @@ def _inject_definitions(self, iet, storage): init = self.lang['thread-num'](retobj=tid) frees.append(Block(header=header, body=[init] + body)) frees.extend(as_list(cbody.frees) + flatten(v.frees)) + frees = sorted(frees, key=lambda x: min( + (obj._C_free_priority for obj in FindSymbols().visit(x) + if obj.is_LocalObject), default=float('inf') + )) # maps/unmaps maps = as_list(cbody.maps) + flatten(v.maps) @@ -452,9 +456,7 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [ - i for i in bases - if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] + bases = [i for i in bases if i.name != i.function._C_name] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index cf71b7773c..a88cf58d45 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -3,7 +3,7 @@ from devito.ir.iet import (Call, ExprStmt, Iteration, SyncSpot, AsyncCallable, FindNodes, FindSymbols, MapNodes, MetaCall, Transformer, - EntryFunction, ThreadCallable, Uxreplace, + EntryFunction, FixedArgsCallable, Uxreplace, derive_parameters) from devito.ir.support import SymbolRegistry from devito.mpi.distributed import MPINeighborhood @@ -603,12 +603,12 @@ def update_args(root, efuncs, dag): foo(..., z) : root(x, z) """ - if isinstance(root, ThreadCallable): + if isinstance(root, FixedArgsCallable): return efuncs # The parameters/arguments lists may have changed since a pass may have: # 1) introduced a new symbol - new_params = derive_parameters(root) + new_params = derive_parameters(root, drop_locals=True) # 2) defined a symbol for which no definition was available yet (e.g. # via a malloc, or a Dereference) diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py index 930cc108b2..3cedd0fb23 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -72,7 +72,8 @@ def _hoist_invariant(iet): """ # Precompute scopes to save time - scopes = {i: Scope([e.expr for e in v]) for i, v in MapNodes().visit(iet).items()} + scopes = {i: Scope([e.expr for e in v if not isinstance(e, Call)]) + for i, v in MapNodes().visit(iet).items()} # Analysis hsmapper = {} diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py deleted file mode 100644 index 8c8aeafa1c..0000000000 --- a/devito/passes/iet/petsc.py +++ /dev/null @@ -1,298 +0,0 @@ -from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, MatVecAction, - Transformer, FindSymbols, LinearSolverExpression, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition) -from devito.types import (PetscMPIInt, PETScStruct, DMDALocalInfo, DM, Mat, - Vec, KSP, PC, SNES, PetscErrorCode) -from devito.symbolics import Byref, Macro, FieldFromPointer -import cgen as c - -__all__ = ['lower_petsc'] - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't - # appear in ccode - - # Check if PETScSolve was used and count occurrences. Each PETScSolve - # will have a unique MatVecAction. - is_petsc = FindNodes(MatVecAction).visit(iet) - - if is_petsc: - - # Collect all solution fields we're solving for - targets = [i.expr.rhs.target for i in is_petsc] - - # Initalize PETSc - init = init_petsc(**kwargs) - - # Create context data struct. - struct = build_struct(iet) - - objs = build_core_objects(targets[-1], struct, **kwargs) - - # Create core PETSc calls required by general linear solves, - # which only need to be generated once e.g create DMDA. - core = core_petsc(targets[-1], objs, **kwargs) - - matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) - - main_mapper = {} - - setup = [] - efuncs = [] - for target in unique_targets: - - solver_objs = build_solver_objs(target) - - matvec_callback_body_iters = [] - - solver_setup = False - - for iter, (matvec,) in matvec_mapper.items(): - - if matvec.expr.rhs.target == target: - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec) - setup.extend(solver) - solver_setup = True - - # Only need to generate solver setup once per target - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) - solver_setup = True - - matvec_body = matvec_body_list._rebuild(body=[ - matvec_body_list.body, iter[0]]) - matvec_body_list = matvec_body_list._rebuild(body=matvec_body) - - main_mapper.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, matvec_callback_body_iters, solver_objs, objs, struct) - - setup.append(matvec_op) - setup.append(c.Line()) - - efuncs.append(matvec_callback) - - # Remove the LinSolveExpr from iet and efuncs that were used to carry - # metadata e.g solver_parameters - main_mapper.update(rebuild_expr_mapper(iet)) - efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} - - iet = Transformer(main_mapper).visit(iet) - efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] - - # Replace symbols appearing in each efunc with a pointer to the PETScStruct - efuncs = transform_efuncs(efuncs, struct) - - body = iet.body._rebuild(body=(tuple(init_setup) + iet.body.body)) - iet = iet._rebuild(body=body) - - return iet, {} - - -def init_petsc(**kwargs): - - # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve. - # NOTE: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options. - initialize = Call('PetscCall', [Call('PetscInitialize', - arguments=['NULL', 'NULL', - 'NULL', 'NULL'])]) - - return tuple([initialize]) - - -def build_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct. - usr_ctx = [] - - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - -def core_petsc(target, objs, **kwargs): - # Assumption: all targets are generated from the same Grid, - # so we can use any target. - - # MPI - call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', - arguments=[objs['comm'], - Byref(objs['size'])])]) - # Create DMDA - dmda = create_dmda(target, objs) - dm_setup = Call('PetscCall', [Call('DMSetUp', arguments=[objs['da']])]) - dm_app_ctx = Call('PetscCall', [Call('DMSetApplicationContext', - arguments=[objs['da'], objs['struct']])]) - dm_mat_type = Call('PetscCall', [Call('DMSetMatType', - arguments=[objs['da'], 'MATSHELL'])]) - dm_local_info = Call('PetscCall', [Call('DMDAGetLocalInfo', - arguments=[objs['da'], Byref(objs['info'])])]) - - return tuple([call_mpi, dmda, dm_setup, dm_app_ctx, dm_mat_type, dm_local_info]) - - -def build_core_objects(target, struct, **kwargs): - - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' - - return {'da': DM(name='da'), - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'struct': struct} - - -def create_dmda(target, objs): - - args = [objs['comm']] - - args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] - - # stencil type - args += ['DMDA_STENCIL_BOX'] - - # global dimensions - args += list(target.shape_global)[::-1] - - # no.of processors in each dimension - args += list(target.grid.distributor.topology)[::-1] - - args += [1, target.space_order] - - args += ['NULL' for _ in range(len(target.space_dimensions))] - - args += [Byref(objs['da'])] - - dmda = Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args) - - return dmda - - -def build_solver_objs(target): - - return {'Jac': Mat(name='J_'+str(target.name)), - 'x_global': Vec(name='x_global_'+str(target.name)), - 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), - 'b_global': Vec(name='b_global_'+str(target.name)), - 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), - 'ksp': KSP(name='ksp_'+str(target.name)), - 'pc': PC(name='pc_'+str(target.name)), - 'snes': SNES(name='snes_'+str(target.name)), - 'x': Vec(name='x_'+str(target.name)), - 'y': Vec(name='y_'+str(target.name))} - - -def generate_solver_calls(solver_objs, objs, matvec): - - snes_create = Call('PetscCall', [Call('SNESCreate', arguments=[ - objs['comm'], Byref(solver_objs['snes'])])]) - - snes_set_dm = Call('PetscCall', [Call('SNESSetDM', arguments=[ - solver_objs['snes'], objs['da']])]) - - create_matrix = Call('PetscCall', [Call('DMCreateMatrix', arguments=[ - objs['da'], Byref(solver_objs['Jac'])])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = Call('PetscCall', [Call('SNESSetType', arguments=[ - solver_objs['snes'], 'SNESKSPONLY'])]) - - global_x = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['x_global'])])]) - - local_x = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['x_local'])])]) - - global_b = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['b_global'])])]) - - local_b = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['b_local'])])]) - - snes_get_ksp = Call('PetscCall', [Call('SNESGetKSP', arguments=[ - solver_objs['snes'], Byref(solver_objs['ksp'])])]) - - return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, - global_x, local_x, global_b, local_b, snes_get_ksp]) - - -def create_matvec_callback(target, matvec_callback_body_iters, - solver_objs, objs, struct): - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - defn_struct = Definition(struct) - - get_context = Call('PetscCall', [Call('MatShellGetContext', - arguments=[solver_objs['Jac'], - Byref(struct)])]) - - body = List(body=[defn_struct, - get_context, - matvec_callback_body_iters]) - - matvec_callback = Callable('MyMatShellMult_'+str(target.name), - matvec_body, - retval=objs['err'], - parameters=(solver_objs['Jac'], - solver_objs['x'], - solver_objs['y'])) - - matvec_operation = Call('PetscCall', [ - Call('MatShellSetOperation', arguments=[solver_objs['Jac'], - 'MATOP_MULT', - Callback(matvec_callback.name, - Void, Void)])]) - - return matvec_callback, matvec_operation - - -def rebuild_expr_mapper(iet): - - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for - expr in FindNodes(LinearSolverExpression).visit(iet)} - - -def transform_efuncs(efuncs, struct): - - efuncs_new = [] - for efunc in efuncs: - new_body = efunc.body - for i in struct.usr_ctx: - new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) - efunc_with_new_body = efunc._rebuild(body=new_body) - efuncs_new.append(efunc_with_new_body) - - return efuncs_new - - -Null = Macro('NULL') -Void = Macro('void') - -petsc_call = String('PetscCall') -petsc_call_mpi = String('PetscCallMPI') -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') - -linear_solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 5d77288f39..bfec3dc52a 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -1,5 +1,5 @@ from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr +from devito.petsc.types import LinearSolveExpr, CallbackExpr @timed_pass() @@ -7,12 +7,9 @@ def petsc_lift(clusters): """ Lift the iteration space surrounding each PETSc equation to create distinct iteration loops. - # TODO: Can probably remove this now due to recursive compilation, but - # leaving it for now. """ processed = [] for c in clusters: - if isinstance(c.exprs[0].rhs, LinearSolveExpr): ispace = c.ispace.lift(c.exprs[0].rhs.target.space_dimensions) processed.append(c.rebuild(ispace=ispace)) @@ -20,3 +17,20 @@ def petsc_lift(clusters): processed.append(c) return processed + + +@timed_pass() +def petsc_project(clusters): + """ + Drop time loop for clusters which appear in PETSc callback functions. + """ + processed = [] + for c in clusters: + if isinstance(c.exprs[0].rhs, CallbackExpr): + time_dims = [d for d in c.ispace.intervals.dimensions if d.is_Time] + ispace = c.ispace.project(lambda d: d not in time_dims) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed diff --git a/devito/petsc/iet.py b/devito/petsc/iet.py deleted file mode 100644 index 659ac9f94d..0000000000 --- a/devito/petsc/iet.py +++ /dev/null @@ -1,372 +0,0 @@ -from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, - Transformer, FindSymbols, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition, BlankLine, PointerCast) -<<<<<<< HEAD -from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, -======= -from devito.petsc import (PetscMPIInt, PETScStruct, DM, Mat, ->>>>>>> b22824875 (fix circular imports) - Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) -from devito.symbolics import Byref, Macro, FieldFromPointer -import cgen as c -from devito.petsc.nodes import MatVecAction, RHSLinearSystem, LinearSolverExpression - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # Check if PETScSolve was used. - petsc_nodes = FindNodes(MatVecAction).visit(iet) - - if not petsc_nodes: - return iet, {} - - else: - # Collect all petsc solution fields - unique_targets = list(set([i.expr.rhs.target for i in petsc_nodes])) - - # Initalize PETSc - init = init_petsc(**kwargs) - - # Create context data struct - struct = build_struct(iet) - - objs = build_core_objects(unique_targets[-1], **kwargs) - - # Create core PETSc calls (not specific to each PETScSolve) - core = core_petsc(unique_targets[-1], struct, objs, **kwargs) - - matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) - - main_mapper = {} - - setup = [] - efuncs = [] - - for target in unique_targets: - - solver_objs = build_solver_objs(target) - - matvec_body_list = List() - - solver_setup = False - - for iter, (matvec,) in matvec_mapper.items(): - - # Skip the MatVecAction if it is not associated with the target - # There will be more than one MatVecAction associated with the target - # e.g interior matvec + BC matvecs - if matvec.expr.rhs.target != target: - continue - - # Only need to generate solver setup once per target - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) - solver_setup = True - - matvec_body = matvec_body_list._rebuild(body=[ - matvec_body_list.body, iter[0]]) - matvec_body_list = matvec_body_list._rebuild(body=matvec_body) - - main_mapper.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, matvec_body_list, solver_objs, objs, - struct) - - setup.append(matvec_op) - setup.append(BlankLine) - efuncs.append(matvec_callback) - - # Remove the LinSolveExpr from iet and efuncs that were used to carry - # metadata e.g solver_parameters - main_mapper.update(rebuild_expr_mapper(iet)) - efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} - - iet = Transformer(main_mapper).visit(iet) - efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] - - # Replace symbols appearing in each efunc with a pointer to the PETScStruct - efuncs = transform_efuncs(efuncs, struct) - - body = iet.body._rebuild(init=init, body=core + tuple(setup) + iet.body.body) - iet = iet._rebuild(body=body) - - return iet, {'efuncs': efuncs} - - -def init_petsc(**kwargs): - - # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve - # TODO: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options - initialize = Call(petsc_call, [ - Call('PetscInitialize', arguments=[Null, Null, Null, Null])]) - - return tuple([petsc_func_begin_user, initialize]) - - -def build_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct - usr_ctx = [] - - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - -def core_petsc(target, struct, objs, **kwargs): - - # MPI - call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', - arguments=[objs['comm'], - Byref(objs['size'])])]) - # Create DMDA - dmda = create_dmda(target, objs) - dm_setup = Call(petsc_call, [ - Call('DMSetUp', arguments=[objs['da']])]) - dm_app_ctx = Call(petsc_call, [ - Call('DMSetApplicationContext', arguments=[objs['da'], struct])]) - dm_mat_type = Call(petsc_call, [ - Call('DMSetMatType', arguments=[objs['da'], 'MATSHELL'])]) - - return tuple([petsc_func_begin_user, call_mpi, dmda, dm_setup, - dm_app_ctx, dm_mat_type, BlankLine]) - - -def build_core_objects(target, **kwargs): - - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' - - return {'da': DM(name='da', liveness='eager'), - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'err': PetscErrorCode(name='err')} - - -def create_dmda(target, objs): - - args = [objs['comm']] - - args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] - - # Stencil type - if len(target.space_dimensions) > 1: - args += ['DMDA_STENCIL_BOX'] - - # Global dimensions - args += list(target.shape_global)[::-1] - - # No.of processors in each dimension - if len(target.space_dimensions) > 1: - args += list(target.grid.distributor.topology)[::-1] - - args += [1, target.space_order] - - args += [Null for _ in range(len(target.space_dimensions))] - - args += [Byref(objs['da'])] - - dmda = Call(petsc_call, [ - Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args)]) - - return dmda - - -def build_solver_objs(target): - - return {'Jac': Mat(name='J_'+str(target.name)), - 'x_global': Vec(name='x_global_'+str(target.name)), - 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), - 'b_global': Vec(name='b_global_'+str(target.name)), - 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), - 'ksp': KSP(name='ksp_'+str(target.name)), - 'pc': PC(name='pc_'+str(target.name)), - 'snes': SNES(name='snes_'+str(target.name)), - 'X_global': Vec(name='X_global_'+str(target.name)), - 'Y_global': Vec(name='Y_global_'+str(target.name)), - 'X_local': Vec(name='X_local_'+str(target.name), liveness='eager'), - 'Y_local': Vec(name='Y_local_'+str(target.name), liveness='eager') - } - - -def generate_solver_calls(solver_objs, objs, matvec, target): - - solver_params = matvec.expr.rhs.solver_parameters - - snes_create = Call(petsc_call, [Call('SNESCreate', arguments=[ - objs['comm'], Byref(solver_objs['snes'])])]) - - snes_set_dm = Call(petsc_call, [Call('SNESSetDM', arguments=[ - solver_objs['snes'], objs['da']])]) - - create_matrix = Call(petsc_call, [Call('DMCreateMatrix', arguments=[ - objs['da'], Byref(solver_objs['Jac'])])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = Call(petsc_call, [Call('SNESSetType', arguments=[ - solver_objs['snes'], 'SNESKSPONLY'])]) - - global_x = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['x_global'])])]) - - local_x = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['x_local'])])]) - - global_b = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['b_global'])])]) - - local_b = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['b_local'])])]) - - snes_get_ksp = Call(petsc_call, [Call('SNESGetKSP', arguments=[ - solver_objs['snes'], Byref(solver_objs['ksp'])])]) - - vec_replace_array = Call(petsc_call, [ - Call('VecReplaceArray', arguments=[solver_objs[ - 'x_local'], FieldFromPointer(target._C_field_data, target._C_symbol)])]) - - ksp_set_tols = Call(petsc_call, [Call('KSPSetTolerances', arguments=[ - solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], - solver_params['ksp_divtol'], solver_params['ksp_max_it']])]) - - ksp_set_type = Call(petsc_call, [Call('KSPSetType', arguments=[ - solver_objs['ksp'], linear_solver_mapper[solver_params['ksp_type']]])]) - - ksp_get_pc = Call(petsc_call, [Call('KSPGetPC', arguments=[ - solver_objs['ksp'], Byref(solver_objs['pc'])])]) - - pc_set_type = Call(petsc_call, [Call('PCSetType', arguments=[ - solver_objs['pc'], linear_solver_mapper[solver_params['pc_type']]])]) - - ksp_set_from_ops = Call(petsc_call, [Call('KSPSetFromOptions', arguments=[ - solver_objs['ksp']])]) - - return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, - global_x, local_x, global_b, local_b, snes_get_ksp, - vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, - pc_set_type, ksp_set_from_ops]) - - -def create_matvec_callback(target, body, solver_objs, objs, struct): - - # There will be 2 PETScArrays within the body - petsc_arrays = [i.function for i in FindSymbols('indexedbases').visit(body) - if isinstance(i.function, PETScArray)] - - # There will only be one PETScArray that is written to within this body and - # one PETScArray which corresponds to the 'seed' vector - petsc_arr_write, = FindSymbols('writes').visit(body) - petsc_arr_seed, = [i for i in petsc_arrays if i.function != petsc_arr_write.function] - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - defn_struct = Definition(struct) - - mat_get_dm = Call(petsc_call, [Call('MatGetDM', arguments=[ - solver_objs['Jac'], Byref(objs['da'])])]) - - dm_get_app_context = Call(petsc_call, [Call('DMGetApplicationContext', arguments=[ - objs['da'], Byref(struct._C_symbol)])]) - - dm_get_local_xvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ - objs['da'], Byref(solver_objs['X_local'])])]) - - global_to_local_begin = Call(petsc_call, [Call('DMGlobalToLocalBegin', arguments=[ - objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) - - global_to_local_end = Call(petsc_call, [Call('DMGlobalToLocalEnd', arguments=[ - objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) - - dm_get_local_yvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ - objs['da'], Byref(solver_objs['Y_local'])])]) - - vec_get_array_y = Call(petsc_call, [Call('VecGetArray', arguments=[ - solver_objs['Y_local'], Byref(petsc_arr_write.function)])]) - - vec_get_array_x = Call(petsc_call, [Call('VecGetArray', arguments=[ - solver_objs['X_local'], Byref(petsc_arrays[0])])]) - - dm_get_local_info = Call(petsc_call, [Call('DMDAGetLocalInfo', arguments=[ - objs['da'], Byref(petsc_arrays[0].function.dmda_info)])]) - - casts = [PointerCast(i.function) for i in petsc_arrays] - - vec_restore_array_y = Call(petsc_call, [Call('VecRestoreArray', arguments=[ - solver_objs['Y_local'], Byref(petsc_arr_write[0])])]) - - vec_restore_array_x = Call(petsc_call, [Call('VecRestoreArray', arguments=[ - solver_objs['X_local'], Byref(petsc_arr_seed)])]) - - dm_local_to_global_begin = Call(petsc_call, [Call('DMLocalToGlobalBegin', arguments=[ - objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) - - dm_local_to_global_end = Call(petsc_call, [Call('DMLocalToGlobalEnd', arguments=[ - objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) - - func_return = Call('PetscFunctionReturn', arguments=[0]) - - matvec_body = List(body=[ - petsc_func_begin_user, defn_struct, mat_get_dm, dm_get_app_context, - dm_get_local_xvec, global_to_local_begin, global_to_local_end, - dm_get_local_yvec, vec_get_array_y, vec_get_array_x, dm_get_local_info, - casts, BlankLine, body, vec_restore_array_y, vec_restore_array_x, - dm_local_to_global_begin, dm_local_to_global_end, func_return]) - - matvec_callback = Callable( - 'MyMatShellMult_'+str(target.name), matvec_body, retval=objs['err'], - parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'])) - - matvec_operation = Call(petsc_call, [ - Call('MatShellSetOperation', arguments=[ - solver_objs['Jac'], 'MATOP_MULT', Callback(matvec_callback.name, - Void, Void)])]) - - return matvec_callback, matvec_operation - - -def rebuild_expr_mapper(callable): - - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for - expr in FindNodes(LinearSolverExpression).visit(callable)} - - -def transform_efuncs(efuncs, struct): - - efuncs_new = [] - for efunc in efuncs: - new_body = efunc.body - for i in struct.usr_ctx: - new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) - efunc_with_new_body = efunc._rebuild(body=new_body) - efuncs_new.append(efunc_with_new_body) - - return efuncs_new - - -Null = Macro('NULL') -Void = 'void' - -petsc_call = 'PetscCall' -petsc_call_mpi = 'PetscCallMPI' -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') - -linear_solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index 5507f999e2..2137495487 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -1,6 +1,5 @@ -from devito.ir.iet import Expression, Callable, Callback +from devito.ir.iet import Expression, Callback, FixedArgsCallable, Call from devito.ir.equations import OpInjectSolve -from devito.tools import as_tuple class LinearSolverExpression(Expression): @@ -19,15 +18,8 @@ def __init__(self, expr, pragmas=None, operation=OpInjectSolve): super().__init__(expr, pragmas=pragmas, operation=operation) -class PETScCallable(Callable): - def __init__(self, name, body, retval=None, parameters=None, - prefix=None, unused_parameters=None): - super().__init__(name, body, retval, parameters, prefix) - self._unused_parameters = as_tuple(unused_parameters) - - @property - def unused_parameters(self): - return self._unused_parameters +class PETScCallable(FixedArgsCallable): + pass class MatVecCallback(Callback): @@ -41,3 +33,7 @@ class FormFunctionCallback(Callback): @property def callback_form(self): return "%s" % self.name + + +class PETScCall(Call): + pass diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index fec645bfc8..d9b875c477 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,81 +1,66 @@ import cgen as c from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Transformer, - MapNodes, Iteration, List, BlankLine, - filter_iterations, retrieve_iteration_tree, +from devito.ir.iet import (Transformer, MapNodes, Iteration, List, BlankLine, Callable, CallableBody, DummyExpr, Call) from devito.symbolics import Byref, Macro, FieldFromPointer from devito.tools import filter_ordered -from devito.petsc.types import (PetscMPIInt, DM, Mat, LocalVec, - GlobalVec, KSP, PC, SNES, PetscErrorCode, DummyArg) +from devito.petsc.types import (PetscMPIInt, DM, Mat, LocalVec, GlobalVec, + KSP, PC, SNES, PetscErrorCode, DummyArg, PetscInt) from devito.petsc.iet.nodes import InjectSolveDummy from devito.petsc.utils import solver_mapper, core_metadata from devito.petsc.iet.routines import PETScCallbackBuilder -from devito.petsc.iet.utils import petsc_call, petsc_call_mpi, petsc_struct +from devito.petsc.iet.utils import (petsc_call, petsc_call_mpi, petsc_struct, + spatial_injectsolve_iter, assign_time_iters, + retrieve_mod_dims) @iet_pass def lower_petsc(iet, **kwargs): - # Check if PETScSolve was used - petsc_nodes = FindNodes(InjectSolveDummy).visit(iet) + injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy, + 'groupby').visit(iet) - if not petsc_nodes: + if not injectsolve_mapper: return iet, {} - unique_targets = list({i.expr.rhs.target for i in petsc_nodes}) + targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()] init = init_petsc(**kwargs) # Assumption is that all targets have the same grid so can use any target here - objs = build_core_objects(unique_targets[-1], **kwargs) + objs = build_core_objects(targets[-1], **kwargs) # Create core PETSc calls (not specific to each PETScSolve) core = make_core_petsc_calls(objs, **kwargs) - # Create injectsolve mapper from the spatial iteration loops - # (exclude time loop if present) - spatial_body = [] - for tree in retrieve_iteration_tree(iet): - root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - spatial_body.append(root) - injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy, - 'groupby').visit(List(body=spatial_body)) - setup = [] subs = {} # Create a different DMDA for each target with a unique space order - unique_dmdas = create_dmda_objs(unique_targets) + unique_dmdas = create_dmda_objs(targets) objs.update(unique_dmdas) for dmda in unique_dmdas.values(): setup.extend(create_dmda_calls(dmda, objs)) builder = PETScCallbackBuilder(**kwargs) - # Create the PETSc calls which are specific to each target - for target in unique_targets: - solver_objs = build_solver_objs(target) - - # Generate the solver setup for target. This is required only - # once per target - for (injectsolve,) in injectsolve_mapper.values(): - # Skip if not associated with the target - if injectsolve.expr.rhs.target != target: - continue - solver_setup = generate_solver_setup(solver_objs, objs, injectsolve, target) - setup.extend(solver_setup) - break - - # Generate all PETSc callback functions for the target via recusive compilation - for iter, (injectsolve,) in injectsolve_mapper.items(): - if injectsolve.expr.rhs.target != target: - continue - matvec_op, formfunc_op, runsolve = builder.make(injectsolve, - objs, solver_objs) - setup.extend([matvec_op, formfunc_op]) - subs.update({iter[0]: List(body=runsolve)}) - break + for iters, (injectsolve,) in injectsolve_mapper.items(): + target = injectsolve.expr.rhs.target + solver_objs = build_solver_objs(target, **kwargs) + + # Generate the solver setup for each InjectSolveDummy + solver_setup = generate_solver_setup(solver_objs, objs, injectsolve) + setup.extend(solver_setup) + + # Retrieve `ModuloDimensions` for use in callback functions + solver_objs['mod_dims'] = retrieve_mod_dims(iters) + # Generate all PETSc callback functions for the target via recursive compilation + matvec_op, formfunc_op, runsolve = builder.make(injectsolve, + objs, solver_objs) + setup.extend([matvec_op, formfunc_op, BlankLine]) + # Only Transform the spatial iteration loop + space_iter, = spatial_injectsolve_iter(iters, injectsolve) + subs.update({space_iter: List(body=runsolve)}) # Generate callback to populate main struct object struct_main = petsc_struct('ctx', filter_ordered(builder.struct_params)) @@ -83,16 +68,21 @@ def lower_petsc(iet, **kwargs): call_struct_callback = petsc_call(struct_callback.name, [Byref(struct_main)]) calls_set_app_ctx = [petsc_call('DMSetApplicationContext', [i, Byref(struct_main)]) for i in unique_dmdas] - setup.extend([BlankLine, call_struct_callback] + calls_set_app_ctx) + setup.extend([call_struct_callback] + calls_set_app_ctx) iet = Transformer(subs).visit(iet) + iet = assign_time_iters(iet, struct_main) + + body = core + tuple(setup) + (BlankLine,) + iet.body.body body = iet.body._rebuild( - init=init, body=core+tuple(setup)+(BlankLine,)+iet.body.body + init=init, body=body, + frees=(c.Line("PetscCall(PetscFinalize());"),) ) iet = iet._rebuild(body=body) metadata = core_metadata() - metadata.update({'efuncs': tuple(builder.efuncs.values())+(struct_callback,)}) + efuncs = tuple(builder.efuncs.values())+(struct_callback,) + metadata.update({'efuncs': efuncs}) return iet, metadata @@ -179,26 +169,29 @@ def create_dmda(dmda, objs): return dmda -def build_solver_objs(target): - name = target.name +def build_solver_objs(target, **kwargs): + sreg = kwargs['sregistry'] return { - 'Jac': Mat(name='J_%s' % name), - 'x_global': GlobalVec(name='x_global_%s' % name), - 'x_local': LocalVec(name='x_local_%s' % name, liveness='eager'), - 'b_global': GlobalVec(name='b_global_%s' % name), - 'b_local': LocalVec(name='b_local_%s' % name), - 'ksp': KSP(name='ksp_%s' % name), - 'pc': PC(name='pc_%s' % name), - 'snes': SNES(name='snes_%s' % name), - 'X_global': GlobalVec(name='X_global_%s' % name), - 'Y_global': GlobalVec(name='Y_global_%s' % name), - 'X_local': LocalVec(name='X_local_%s' % name, liveness='eager'), - 'Y_local': LocalVec(name='Y_local_%s' % name, liveness='eager'), - 'dummy': DummyArg(name='dummy_%s' % name) + 'Jac': Mat(sreg.make_name(prefix='J_')), + 'x_global': GlobalVec(sreg.make_name(prefix='x_global_')), + 'x_local': LocalVec(sreg.make_name(prefix='x_local_'), liveness='eager'), + 'b_global': GlobalVec(sreg.make_name(prefix='b_global_')), + 'b_local': LocalVec(sreg.make_name(prefix='b_local_')), + 'ksp': KSP(sreg.make_name(prefix='ksp_')), + 'pc': PC(sreg.make_name(prefix='pc_')), + 'snes': SNES(sreg.make_name(prefix='snes_')), + 'X_global': GlobalVec(sreg.make_name(prefix='X_global_')), + 'Y_global': GlobalVec(sreg.make_name(prefix='Y_global_')), + 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), + 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), + 'dummy': DummyArg(sreg.make_name(prefix='dummy_')), + 'localsize': PetscInt(sreg.make_name(prefix='localsize_')) } -def generate_solver_setup(solver_objs, objs, injectsolve, target): +def generate_solver_setup(solver_objs, objs, injectsolve): + target = injectsolve.expr.rhs.target + dmda = objs['da_so_%s' % target.space_order] solver_params = injectsolve.expr.rhs.solver_parameters @@ -220,9 +213,6 @@ def generate_solver_setup(solver_objs, objs, injectsolve, target): global_x = petsc_call('DMCreateGlobalVector', [dmda, Byref(solver_objs['x_global'])]) - local_x = petsc_call('DMCreateLocalVector', - [dmda, Byref(solver_objs['x_local'])]) - global_b = petsc_call('DMCreateGlobalVector', [dmda, Byref(solver_objs['b_global'])]) @@ -232,11 +222,6 @@ def generate_solver_setup(solver_objs, objs, injectsolve, target): snes_get_ksp = petsc_call('SNESGetKSP', [solver_objs['snes'], Byref(solver_objs['ksp'])]) - vec_replace_array = petsc_call( - 'VecReplaceArray', [solver_objs['x_local'], - FieldFromPointer(target._C_field_data, target._C_symbol)] - ) - ksp_set_tols = petsc_call( 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], solver_params['ksp_divtol'], @@ -261,11 +246,9 @@ def generate_solver_setup(solver_objs, objs, injectsolve, target): snes_set_jac, snes_set_type, global_x, - local_x, global_b, local_b, snes_get_ksp, - vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, @@ -275,8 +258,10 @@ def generate_solver_setup(solver_objs, objs, injectsolve, target): def generate_struct_callback(struct): - body = [DummyExpr(FieldFromPointer(i._C_symbol, struct), - i._C_symbol) for i in struct.fields] + body = [ + DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol) + for i in struct.fields if i not in struct.time_dim_fields + ] struct_callback_body = CallableBody( List(body=body), init=tuple([petsc_func_begin_user]), retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) @@ -288,51 +273,6 @@ def generate_struct_callback(struct): return struct_callback -@iet_pass -def sort_frees(iet): - frees = iet.body.frees - - if not frees: - return iet, {} - - destroys = ["VecDestroy", "MatDestroy", "SNESDestroy", "DMDestroy"] - priority = {k: i for i, k in enumerate(destroys, start=1)} - - def key(i): - for destroy, prio in priority.items(): - if destroy in str(i): - return prio - return float('inf') - - frees = sorted(frees, key=key) - - body = iet.body._rebuild(frees=frees) - iet = iet._rebuild(body=body) - return iet, {} - - -@iet_pass -def sort_frees(iet): - frees = iet.body.frees - - if not frees: - return iet, {} - - destroys = ["VecDestroy", "MatDestroy", "SNESDestroy", "DMDestroy"] - priority = {k: i for i, k in enumerate(destroys, start=1)} - - def key(i): - for destroy, prio in priority.items(): - if destroy in str(i): - return prio - return float('inf') - - frees = sorted(frees, key=key) - body = iet.body._rebuild(frees=frees) - iet = iet._rebuild(body=body) - return iet, {} - - Null = Macro('NULL') void = 'void' diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index c14e06269f..d164d9663f 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -3,21 +3,27 @@ import cgen as c from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, - Callable) -from devito.symbolics import Byref, FieldFromPointer, Macro -from devito.petsc.types import PETScStruct + Dereference, DummyExpr, BlankLine) +from devito.symbolics import Byref, FieldFromPointer, Macro, Cast +from devito.symbolics.unevaluation import Mul +from devito.types.basic import AbstractFunction +from devito.types import LocalObject, ModuloDimension, TimeDimension +from devito.tools import ctypes_to_cstr, dtype_to_ctype, CustomDtype +from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatVecCallback) -from devito.petsc.utils import petsc_call +from devito.petsc.iet.utils import petsc_call, petsc_struct, drop_callbackexpr +from devito.ir.support import SymbolRegistry class PETScCallbackBuilder: """ Build IET routines to generate PETSc callback functions. """ - def __new__(cls, rcompile=None, **kwargs): + def __new__(cls, rcompile=None, sregistry=None, **kwargs): obj = object.__new__(cls) obj.rcompile = rcompile + obj.sregistry = sregistry obj._efuncs = OrderedDict() obj._struct_params = [] @@ -61,15 +67,16 @@ def make_all(self, injectsolve, objs, solver_objs): return matvec_callback, formfunc_callback, formrhs_callback def make_matvec(self, injectsolve, objs, solver_objs): - target = injectsolve.expr.rhs.target # Compile matvec `eqns` into an IET via recursive compilation irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, - options={'mpi': False}) - body_matvec = self.create_matvec_body(injectsolve, irs_matvec.uiet.body, + options={'mpi': False}, sregistry=SymbolRegistry()) + body_matvec = self.create_matvec_body(injectsolve, + List(body=irs_matvec.uiet.body), solver_objs, objs) matvec_callback = PETScCallable( - 'MyMatShellMult_%s' % target.name, body_matvec, retval=objs['err'], + self.sregistry.make_name(prefix='MyMatShellMult_'), body_matvec, + retval=objs['err'], parameters=( solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'] ) @@ -79,8 +86,12 @@ def make_matvec(self, injectsolve, objs, solver_objs): def create_matvec_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs + body = drop_callbackexpr(body) + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + body = uxreplace_mod_dims(body, solver_objs['mod_dims'], target=linsolveexpr.target) + struct = build_petsc_struct(body, 'matvec', liveness='eager') y_matvec = linsolveexpr.arrays['y_matvec'] @@ -137,16 +148,20 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] ]) - # NOTE: Question: I have placed a chunk of the calls in the `stacks` argument - # of the `CallableBody` to ensure that these calls precede the `cast` statements. - # The 'casts' depend on the calls, so this order is necessary. By doing this, - # I avoid having to manually construct the 'casts' and can allow Devito to handle - # their construction. Are there any potential issues with this approach? - body = [body, - vec_restore_array_y, - vec_restore_array_x, - dm_local_to_global_begin, - dm_local_to_global_end] + # TODO: Some of the calls are placed in the `stacks` argument of the + # `CallableBody` to ensure that they precede the `cast` statements. The + # 'casts' depend on the calls, so this order is necessary. By doing this, + # you avoid having to manually construct the `casts` and can allow + # Devito to handle their construction. This is a temporary solution and + # should be revisited + + body = body._rebuild( + body=body.body + + (vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end) + ) stacks = ( mat_get_dm, @@ -160,40 +175,52 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): dm_get_local_info ) + # Dereference function data in struct + dereference_funcs = [Dereference(i, struct) for i in + struct.fields if isinstance(i.function, AbstractFunction)] + matvec_body = CallableBody( List(body=body), - init=tuple([petsc_func_begin_user]), - stacks=stacks, - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])) + init=(petsc_func_begin_user,), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) - # Replace data with pointer to data in struct - subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for i in struct.fields} matvec_body = Uxreplace(subs).visit(matvec_body) - self._struct_params.extend(struct.usr_ctx) + self._struct_params.extend(struct.fields) return matvec_body def make_formfunc(self, injectsolve, objs, solver_objs): - target = injectsolve.expr.rhs.target # Compile formfunc `eqns` into an IET via recursive compilation - irs_formfunc, _ = self.rcompile(injectsolve.expr.rhs.formfuncs, - options={'mpi': False}) - body_formfunc = self.create_formfunc_body(injectsolve, irs_formfunc.uiet.body, + irs_formfunc, _ = self.rcompile( + injectsolve.expr.rhs.formfuncs, + options={'mpi': False}, sregistry=SymbolRegistry() + ) + body_formfunc = self.create_formfunc_body(injectsolve, + List(body=irs_formfunc.uiet.body), solver_objs, objs) formfunc_callback = PETScCallable( - 'FormFunction_%s' % target.name, body_formfunc, retval=objs['err'], + self.sregistry.make_name(prefix='FormFunction_'), body_formfunc, + retval=objs['err'], parameters=(solver_objs['snes'], solver_objs['X_global'], - solver_objs['Y_global']), unused_parameters=(solver_objs['dummy']) + solver_objs['Y_global'], solver_objs['dummy']) ) return formfunc_callback def create_formfunc_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs + body = drop_callbackexpr(body) + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + body = uxreplace_mod_dims(body, solver_objs['mod_dims'], linsolveexpr.target) + struct = build_petsc_struct(body, 'formfunc', liveness='eager') y_formfunc = linsolveexpr.arrays['y_formfunc'] @@ -250,11 +277,13 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] ]) - body = [body, - vec_restore_array_y, - vec_restore_array_x, - dm_local_to_global_begin, - dm_local_to_global_end] + body = body._rebuild( + body=body.body + + (vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end) + ) stacks = ( snes_get_dm, @@ -268,30 +297,34 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): dm_get_local_info ) + # Dereference function data in struct + dereference_funcs = [Dereference(i, struct) for i in + struct.fields if isinstance(i.function, AbstractFunction)] + formfunc_body = CallableBody( List(body=body), - init=tuple([petsc_func_begin_user]), - stacks=stacks, - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])) + init=(petsc_func_begin_user,), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) - # Replace data with pointer to data in struct - subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for i in struct.fields} formfunc_body = Uxreplace(subs).visit(formfunc_body) - self._struct_params.extend(struct.usr_ctx) + self._struct_params.extend(struct.fields) return formfunc_body def make_formrhs(self, injectsolve, objs, solver_objs): - target = injectsolve.expr.rhs.target # Compile formrhs `eqns` into an IET via recursive compilation irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, - options={'mpi': False}) - body_formrhs = self.create_formrhs_body(injectsolve, irs_formrhs.uiet.body, + options={'mpi': False}, sregistry=SymbolRegistry()) + body_formrhs = self.create_formrhs_body(injectsolve, + List(body=irs_formrhs.uiet.body), solver_objs, objs) - formrhs_callback = Callable( - 'FormRHS_%s' % target.name, body_formrhs, retval=objs['err'], + formrhs_callback = PETScCallable( + self.sregistry.make_name(prefix='FormRHS_'), body_formrhs, retval=objs['err'], parameters=( solver_objs['snes'], solver_objs['b_local'] ) @@ -302,6 +335,8 @@ def make_formrhs(self, injectsolve, objs, solver_objs): def create_formrhs_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs + body = drop_callbackexpr(body) + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) @@ -316,6 +351,8 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] ) + body = uxreplace_mod_dims(body, solver_objs['mod_dims'], linsolveexpr.target) + struct = build_petsc_struct(body, 'formrhs', liveness='eager') dm_get_app_context = petsc_call( @@ -326,27 +363,32 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): 'VecRestoreArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] ) - body = [body, - vec_restore_array] + body = body._rebuild(body=body.body + (vec_restore_array,)) stacks = ( snes_get_dm, dm_get_app_context, vec_get_array, - dm_get_local_info, + dm_get_local_info ) + # Dereference function data in struct + dereference_funcs = [Dereference(i, struct) for i in + struct.fields if isinstance(i.function, AbstractFunction)] + formrhs_body = CallableBody( List(body=[body]), - init=tuple([petsc_func_begin_user]), - stacks=stacks, - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])) + init=(petsc_func_begin_user,), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) - # Replace data with pointer to data in struct - subs = {i: FieldFromPointer(i, struct) for i in struct.usr_ctx} + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for + i in struct.fields if not isinstance(i.function, AbstractFunction)} formrhs_body = Uxreplace(subs).visit(formrhs_body) - self._struct_params.extend(struct.usr_ctx) + self._struct_params.extend(struct.fields) return formrhs_body @@ -357,6 +399,19 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): rhs_call = petsc_call(rhs_callback.name, list(rhs_callback.parameters)) + local_x = petsc_call('DMCreateLocalVector', + [dmda, Byref(solver_objs['x_local'])]) + + if any(i.is_Time for i in target.dimensions): + vec_replace_array = time_dep_replace( + injectsolve, target, solver_objs, objs, self.sregistry + ) + else: + field_from_ptr = FieldFromPointer(target._C_field_data, target._C_symbol) + vec_replace_array = (petsc_call( + 'VecReplaceArray', [solver_objs['x_local'], field_from_ptr] + ),) + dm_local_to_global_x = petsc_call( 'DMLocalToGlobal', [dmda, solver_objs['x_local'], 'INSERT_VALUES', solver_objs['x_global']] @@ -375,22 +430,83 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): dmda, solver_objs['x_global'], 'INSERT_VALUES', solver_objs['x_local']] ) - calls = (rhs_call, - dm_local_to_global_x, - dm_local_to_global_b, - snes_solve, - dm_global_to_local_x) - - return calls + return ( + rhs_call, + local_x + ) + vec_replace_array + ( + dm_local_to_global_x, + dm_local_to_global_b, + snes_solve, + dm_global_to_local_x, + BlankLine, + ) def build_petsc_struct(iet, name, liveness): - # Place all context data required by the shell routines - # into a PETScStruct + # Place all context data required by the shell routines into a struct + # TODO: Clean this search up basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases|writes').visit(iet) - usr_ctx = [data for data in basics if data not in avoid] - return PETScStruct(name, usr_ctx, liveness=liveness) + time_dims = [i for i in FindSymbols('dimensions').visit(iet) + if isinstance(i, (TimeDimension, ModuloDimension))] + avoid0 = [i for i in FindSymbols('indexedbases').visit(iet) + if isinstance(i.function, PETScArray)] + avoid1 = [i for i in FindSymbols('dimensions|writes').visit(iet) + if i not in time_dims] + fields = [data.function for data in basics if data not in avoid0+avoid1] + return petsc_struct(name, fields, liveness) + + +def time_dep_replace(injectsolve, target, solver_objs, objs, sregistry): + target_time = injectsolve.expr.lhs + target_time = [i for i, d in zip(target_time.indices, + target_time.dimensions) if d.is_Time] + assert len(target_time) == 1 + target_time = target_time.pop() + + ctype_str = ctypes_to_cstr(dtype_to_ctype(target.dtype)) + + class BarCast(Cast): + _base_typ = ctype_str + + class StartPtr(LocalObject): + dtype = CustomDtype(ctype_str, modifier=' *') + + start_ptr = StartPtr(sregistry.make_name(prefix='start_ptr_')) + + vec_get_size = petsc_call( + 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] + ) + + # TODO: What is the correct way to use Mul here? Devito Mul? Sympy Mul? + field_from_ptr = FieldFromPointer(target._C_field_data, target._C_symbol) + expr = DummyExpr( + start_ptr, BarCast(field_from_ptr, ' *') + + Mul(target_time, solver_objs['localsize']), init=True + ) + + vec_replace_array = petsc_call('VecReplaceArray', [solver_objs['x_local'], start_ptr]) + return (vec_get_size, expr, vec_replace_array) + + +def uxreplace_mod_dims(body, mod_dims, target): + """ + Replace ModuloDimensions in callback functions with the corresponding + ModuloDimensions generated by the initial lowering. They must match because + they are assigned and updated in the struct at each time step. This is a valid + uxreplace because all functions appearing in the callback functions are + passed through the initial lowering. + """ + # old_mod_dims = [ + # i for i in FindSymbols('dimensions').visit(body) if isinstance(i, ModuloDimension) + # ] + # from IPython import embed; embed() + # if not old_mod_dims: + # return body + # t_tmp = + # from IPython import embed; embed() + # body = Uxreplace({i: mod_dims[i] for i in old_mod_dims}).visit(body) + # return Uxreplace({i: mod_dims[i.origin] for i in old_mod_dims}).visit(body) + return body Null = Macro('NULL') diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py new file mode 100644 index 0000000000..8f1b0a34aa --- /dev/null +++ b/devito/petsc/iet/utils.py @@ -0,0 +1,83 @@ +from devito.ir.iet.nodes import Expression +from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall +from devito.ir.equations import OpInjectSolve +from devito.ir.iet import (FindNodes, retrieve_iteration_tree, + filter_iterations, Transformer, Iteration, + DummyExpr, List) +from devito.symbolics import FieldFromComposite + + +def petsc_call(specific_call, call_args): + return PETScCall('PetscCall', [PETScCall(specific_call, arguments=call_args)]) + + +def petsc_call_mpi(specific_call, call_args): + return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)]) + + +def petsc_struct(name, fields, liveness='lazy'): + # TODO: Fix this circular import + from devito.petsc.types.object import PETScStruct + return PETScStruct(name=name, pname='MatContext', + fields=fields, liveness=liveness) + + +def spatial_injectsolve_iter(iter, injectsolve): + spatial_body = [] + for tree in retrieve_iteration_tree(iter[0]): + root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] + if injectsolve in FindNodes(InjectSolveDummy).visit(root): + spatial_body.append(root) + return spatial_body + + +# Mapping special Eq operations to their corresponding IET Expression subclass types. +# These operations correspond to subclasses of Eq utilised within PETScSolve. +petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} + + +def drop_callbackexpr(body): + # TODO: fix this circular import + from devito.petsc.types import CallbackExpr + nodes = FindNodes(Expression).visit(body) + mapper = { + expr: expr._rebuild(expr=expr.expr._rebuild(rhs=expr.expr.rhs.args[0])) + for expr in nodes + if isinstance(expr.expr.rhs, CallbackExpr) + } + return Transformer(mapper).visit(body) + + +def assign_time_iters(iet, struct): + """ + Assign time iterators to the struct within loops containing PETScCalls. + Ensure that assignment occurs only once per time loop, if necessary. + Assign only the iterators that are common between the struct fields + and the actual Iteration. + """ + time_iters = [ + i for i in FindNodes(Iteration).visit(iet) + if i.dim.is_Time and FindNodes(PETScCall).visit(i) + ] + + if not time_iters: + return iet + + mapper = {} + for iter in time_iters: + common_dims = [dim for dim in iter.dimensions if dim in struct.fields] + common_dims = [ + DummyExpr(FieldFromComposite(dim, struct), dim) for dim in common_dims + ] + iter_new = iter._rebuild(nodes=List(body=tuple(common_dims)+iter.nodes)) + mapper.update({iter: iter_new}) + + return Transformer(mapper).visit(iet) + + +def retrieve_mod_dims(iters): + outer_iter_dims = iters[0].dimensions + if any(dim.is_Time for dim in outer_iter_dims): + mod_dims = [dim for dim in outer_iter_dims if dim.is_Modulo] + return {dim.origin: dim for dim in mod_dims} + return {} diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 6f30b57f10..fd5403e319 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -4,18 +4,17 @@ from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative -from devito.types import Eq +from devito.types import Eq, Symbol from devito.types.equation import InjectSolveEq from devito.operations.solve import eval_time_derivatives -from devito.symbolics import uxreplace -from devito.petsc.types import LinearSolveExpr, PETScArray +from devito.symbolics import retrieve_functions +from devito.petsc.types import LinearSolveExpr, PETScArray, CallbackExpr __all__ = ['PETScSolve'] -def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - +def PETScSolve(eqns, target, solver_parameters=None, **kwargs): prefixes = ['y_matvec', 'x_matvec', 'y_formfunc', 'x_formfunc', 'b_tmp'] arrays = { @@ -29,56 +28,46 @@ def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): for p in prefixes } - b, F_target = separate_eqn(eq, target) - - # TODO: Current assumption is that problem is linear and user has not provided - # a jacobian. Hence, we can use F_target to form the jac-vec product - matvecaction = Eq(arrays['y_matvec'], - uxreplace(F_target, {target: arrays['x_matvec']}), - subdomain=eq.subdomain) - - formfunction = Eq(arrays['y_formfunc'], - uxreplace(F_target, {target: arrays['x_formfunc']}), - subdomain=eq.subdomain) - - rhs = Eq(arrays['b_tmp'], b, subdomain=eq.subdomain) - - # Placeholder equation for inserting calls to the solver - inject_solve = InjectSolveEq(arrays['b_tmp'], LinearSolveExpr( - b, target=target, solver_parameters=solver_parameters, matvecs=[matvecaction], - formfuncs=[formfunction], formrhs=[rhs], arrays=arrays, - ), subdomain=eq.subdomain) - - if not bcs: - return [inject_solve] - - # NOTE: BELOW IS NOT FULLY TESTED/IMPLEMENTED YET - bcs_for_matvec = [] - bcs_for_formfunc = [] - bcs_for_rhs = [] - for bc in bcs: - # TODO: Insert code to distiguish between essential and natural - # boundary conditions since these are treated differently within - # the solver - # NOTE: May eventually remove the essential bcs from the solve - # (and move to rhs) but for now, they are included since this - # is not trivial to implement when using DMDA - # NOTE: Below is temporary -> Just using this as a palceholder for - # the actual BC implementation - centre = centre_stencil(F_target, target) - bcs_for_matvec.append(Eq(arrays['y_matvec'], - centre.subs(target, arrays['x_matvec']), - subdomain=bc.subdomain)) - bcs_for_formfunc.append(Eq(arrays['y_formfunc'], - 0., subdomain=bc.subdomain)) - # NOTE: Temporary - bcs_for_rhs.append(Eq(arrays['b_tmp'], 0., subdomain=bc.subdomain)) - - inject_solve = InjectSolveEq(arrays['b_tmp'], LinearSolveExpr( - b, target=target, solver_parameters=solver_parameters, - matvecs=[matvecaction]+bcs_for_matvec, - formfuncs=[formfunction]+bcs_for_formfunc, formrhs=[rhs], + matvecs = [] + formfuncs = [] + formrhs = [] + + eqns = eqns if isinstance(eqns, (list, tuple)) else [eqns] + for eq in eqns: + b, F_target = separate_eqn(eq, target) + + # TODO: Current assumption is that problem is linear and user has not provided + # a jacobian. Hence, we can use F_target to form the jac-vec product + matvecs.append(Eq( + arrays['y_matvec'], + CallbackExpr(F_target.subs({target: arrays['x_matvec']})), + subdomain=eq.subdomain + )) + + formfuncs.append(Eq( + arrays['y_formfunc'], + CallbackExpr(F_target.subs({target: arrays['x_formfunc']})), + subdomain=eq.subdomain + )) + time_dim = Symbol('t_tmp') + time_mapper = {target.grid.stepping_dim: time_dim} + formrhs.append(Eq( + arrays['b_tmp'], + CallbackExpr(b).subs(time_mapper), + subdomain=eq.subdomain + )) + + # Placeholder equation for inserting calls to the solver and generating + # correct time loop etc + inject_solve = InjectSolveEq(target, LinearSolveExpr( + expr=tuple(retrieve_functions(eqns)), + target=target, + solver_parameters=solver_parameters, + matvecs=matvecs, + formfuncs=formfuncs, + formrhs=formrhs, arrays=arrays, + time_dim=time_dim ), subdomain=eq.subdomain) return [inject_solve] diff --git a/devito/petsc/types.py b/devito/petsc/types.py deleted file mode 100644 index 72a0faa1e8..0000000000 --- a/devito/petsc/types.py +++ /dev/null @@ -1,429 +0,0 @@ -import sympy - -<<<<<<< HEAD -from devito.tools import Reconstructable, sympy_mutex -======= -from devito.tools import CustomDtype -from devito.types import LocalObject, Eq, CompositeObject -from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic -from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction, AbstractSymbol -from devito.finite_differences.tools import fd_weights_registry -from devito.tools import dtype_to_ctype, Reconstructable, sympy_mutex -from devito.symbolics import FieldFromComposite, Byref - -from devito.petsc.utils import petsc_call - - -class DM(LocalObject): - """ - PETSc Data Management object (DM). - """ - dtype = CustomDtype('DM') - - def __init__(self, *args, stencil_width=None, **kwargs): - super().__init__(*args, **kwargs) - self._stencil_width = stencil_width - - @property - def stencil_width(self): - return self._stencil_width - - @property - def info(self): - return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('DMDestroy', [Byref(self.function)]) - - -class Mat(LocalObject): - """ - PETSc Matrix object (Mat). - """ - dtype = CustomDtype('Mat') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('MatDestroy', [Byref(self.function)]) - - -class LocalVec(LocalObject): - """ - PETSc Vector object (Vec). - """ - dtype = CustomDtype('Vec') - - -class GlobalVec(LocalObject): - """ - PETSc Vector object (Vec). - """ - dtype = CustomDtype('Vec') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('VecDestroy', [Byref(self.function)]) - - -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 SNES(LocalObject): - """ - PETSc SNES : Non-Linear Systems Solvers. - """ - dtype = CustomDtype('SNES') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('SNESDestroy', [Byref(self.function)]) - - -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 DMDALocalInfo(LocalObject): - """ - PETSc object - C struct containing information - about the local grid. - """ - dtype = CustomDtype('DMDALocalInfo') - - -class PetscErrorCode(LocalObject): - """ - PETSc datatype used to return PETSc error codes. - https://petsc.org/release/manualpages/Sys/PetscErrorCode/ - """ - dtype = CustomDtype('PetscErrorCode') - - -class DummyArg(LocalObject): - dtype = CustomDtype('void', modifier='*') - - -class PETScArray(ArrayBasic, Differentiable): - """ - 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. - Differentiable enables compatability with standard Function objects, - allowing for the use of the `subs` method. - TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) - and then PETScArray(ArrayBasic). - """ - - _data_alignment = False - - # Default method for the finite difference approximation weights computation. - _default_fd = 'taylor' - - __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients', - 'space_order')) - - def __init_finalize__(self, *args, **kwargs): - - super().__init_finalize__(*args, **kwargs) - - # Symbolic (finite difference) coefficients - self._coefficients = kwargs.get('coefficients', self._default_fd) - if self._coefficients not in fd_weights_registry: - raise ValueError("coefficients must be one of %s" - " not %s" % (str(fd_weights_registry), self._coefficients)) - self._shape = kwargs.get('shape') - self._space_order = kwargs.get('space_order', 1) - - @classmethod - def __dtype_setup__(cls, **kwargs): - return kwargs.get('dtype', np.float32) - - @property - def coefficients(self): - """Form of the coefficients of the function.""" - return self._coefficients - - @property - def shape(self): - return self._shape - - @property - def space_order(self): - return self._space_order - - @cached_property - def _shape_with_inhalo(self): - """ - Shape of the domain+inhalo region. The inhalo region comprises the - outhalo as well as any additional "ghost" layers for MPI halo - exchanges. Data in the inhalo region are exchanged when running - Operators to maintain consistent values as in sequential runs. - - Notes - ----- - Typically, this property won't be used in user code, but it may come - in handy for testing or debugging - """ - return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) - - @cached_property - def shape_allocated(self): - """ - Shape of the allocated data of the Function type object from which - this PETScArray was derived. It includes the domain and inhalo regions, - as well as any additional padding surrounding the halo. - - Notes - ----- - In an MPI context, this is the *local* with_halo region shape. - """ - return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, - self._padding)], - getters=self.dimensions) - - @cached_property - def _C_ctype(self): - # NOTE: Reverting to using float/double instead of PetscScalar for - # simplicity when opt='advanced'. Otherwise, Temp objects must also - # be converted to PetscScalar. Additional tests are needed to - # ensure this approach is fine. Previously, issues arose from - # mismatches between precision of Function objects in Devito and the - # precision of the PETSc configuration. - # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables - # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of - # the user's PETSc configuration. - return POINTER(dtype_to_ctype(self.dtype)) - - @property - def symbolic_shape(self): - field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] - # Reverse it since DMDA is setup backwards to Devito dimensions. - return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - - @property - def dmda(self): - name = 'da_so_%s' % self.space_order - return DM(name=name, liveness='eager', stencil_width=self.space_order) - - -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] - - -class InjectSolveEq(Eq): - pass ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) - - -class LinearSolveExpr(sympy.Function, Reconstructable): - - __rargs__ = ('expr',) -<<<<<<< HEAD -<<<<<<< HEAD - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays') -======= - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', 'formfuncs', 'formrhs') ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays') ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - - defaults = { - 'ksp_type': 'gmres', - 'pc_type': 'jacobi', - 'ksp_rtol': 1e-7, # Relative tolerance - 'ksp_atol': 1e-50, # Absolute tolerance - 'ksp_divtol': 1e4, # Divergence tolerance - 'ksp_max_it': 10000 # Maximum iterations - } - - def __new__(cls, expr, target=None, solver_parameters=None, -<<<<<<< HEAD -<<<<<<< HEAD - matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): -======= - matvecs=None, formfuncs=None, formrhs=None, **kwargs): ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - - if solver_parameters is None: - solver_parameters = cls.defaults - else: - for key, val in cls.defaults.items(): - solver_parameters[key] = solver_parameters.get(key, val) - - with sympy_mutex: - obj = sympy.Basic.__new__(cls, expr) - obj._expr = expr - obj._target = target - obj._solver_parameters = solver_parameters - obj._matvecs = matvecs - obj._formfuncs = formfuncs - obj._formrhs = formrhs -<<<<<<< HEAD -<<<<<<< HEAD - obj._arrays = arrays -======= ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - obj._arrays = arrays ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - return obj - - def __repr__(self): - return "%s(%s)" % (self.__class__.__name__, self.expr) - - __str__ = __repr__ - - def _sympystr(self, printer): - return str(self) - - def __hash__(self): - return hash(self.expr) - - def __eq__(self, other): - return (isinstance(other, LinearSolveExpr) and - self.expr == other.expr and - self.target == other.target) - - @property - def expr(self): - return self._expr - - @property - def target(self): - return self._target - - @property - def solver_parameters(self): - return self._solver_parameters - - @property - def matvecs(self): - return self._matvecs - - @property - def formfuncs(self): - return self._formfuncs - - @property - def formrhs(self): - return self._formrhs - -<<<<<<< HEAD -<<<<<<< HEAD - @property - def formrhs(self): - return self._formrhs -======= -======= - @property - def arrays(self): - return self._arrays - ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - func = Reconstructable._rebuild - - -class PETScStruct(CompositeObject): - - __rargs__ = ('name', 'usr_ctx') - __rkwargs__ = ('liveness',) - - def __init__(self, name, usr_ctx, liveness='lazy'): - pfields = [(i._C_name, dtype_to_ctype(i.dtype)) - for i in usr_ctx if isinstance(i, AbstractSymbol)] - self._usr_ctx = usr_ctx - super().__init__(name, 'MatContext', pfields) ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) - - assert liveness in ['eager', 'lazy'] - self._liveness = liveness - - @property - def arrays(self): - return self._arrays - -<<<<<<< HEAD - func = Reconstructable._rebuild -======= - def _arg_values(self, **kwargs): - # TODO: Handle passing `struct dataobj *restrict g1_vec` and similar - # elements within in the struct. It seemed necessary to convert the Numpy array - # into a pointer to ctypes.Struct within MatContext - # (i.e override _arg_finalize for PETScStruct here?), but this approach seemed - # to result in a NIL pointer in the generated code ... any suggestions? - # Alternative: Allow these pointers to appear in the main kernel arguments - # as usual and use ctx->g1_vec = g1_vec; to assign them to the struct? - values = super()._arg_values(**kwargs) - for i in self.usr_ctx: - setattr(values[self.name]._obj, i.name, kwargs['args'][i.name]) - return values - - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' -<<<<<<< HEAD -<<<<<<< HEAD ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - -======= - ->>>>>>> 9d611239f (compiler: Fix derive_callback_inputs) - @property - def fields(self): - return self._usr_ctx ->>>>>>> 41157b0bc (compiler: Place dataobj pointers inside ctx struct) diff --git a/devito/petsc/types/__init__.py b/devito/petsc/types/__init__.py new file mode 100644 index 0000000000..ebcceb8d45 --- /dev/null +++ b/devito/petsc/types/__init__.py @@ -0,0 +1,3 @@ +from .array import * # noqa +from .types import * # noqa +from .object import * # noqa diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py new file mode 100644 index 0000000000..a150ea3247 --- /dev/null +++ b/devito/petsc/types/array.py @@ -0,0 +1,117 @@ +from functools import cached_property +import numpy as np +from ctypes import POINTER + +from devito.types.utils import DimensionTuple +from devito.types.array import ArrayBasic +from devito.finite_differences import Differentiable +from devito.types.basic import AbstractFunction +from devito.finite_differences.tools import fd_weights_registry +from devito.tools import dtype_to_ctype +from devito.symbolics import FieldFromComposite + +from .object import DM + + +class PETScArray(ArrayBasic, Differentiable): + """ + PETScArrays are generated by the compiler only and represent + a customised variant of ArrayBasic. + Differentiable enables compatability with standard Function objects, + allowing for the use of the `subs` method. + TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) + and then PETScArray(ArrayBasic). + """ + + _data_alignment = False + + # Default method for the finite difference approximation weights computation. + _default_fd = 'taylor' + + __rkwargs__ = (AbstractFunction.__rkwargs__ + + ('dimensions', 'shape', 'liveness', 'coefficients', + 'space_order')) + + def __init_finalize__(self, *args, **kwargs): + + super().__init_finalize__(*args, **kwargs) + + # Symbolic (finite difference) coefficients + self._coefficients = kwargs.get('coefficients', self._default_fd) + if self._coefficients not in fd_weights_registry: + raise ValueError("coefficients must be one of %s" + " not %s" % (str(fd_weights_registry), self._coefficients)) + self._shape = kwargs.get('shape') + self._space_order = kwargs.get('space_order', 1) + + @classmethod + def __dtype_setup__(cls, **kwargs): + return kwargs.get('dtype', np.float32) + + @property + def coefficients(self): + """Form of the coefficients of the function.""" + return self._coefficients + + @property + def shape(self): + return self._shape + + @property + def space_order(self): + return self._space_order + + @cached_property + def _shape_with_inhalo(self): + """ + Shape of the domain+inhalo region. The inhalo region comprises the + outhalo as well as any additional "ghost" layers for MPI halo + exchanges. Data in the inhalo region are exchanged when running + Operators to maintain consistent values as in sequential runs. + + Notes + ----- + Typically, this property won't be used in user code, but it may come + in handy for testing or debugging + """ + return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) + + @cached_property + def shape_allocated(self): + """ + Shape of the allocated data of the Function type object from which + this PETScArray was derived. It includes the domain and inhalo regions, + as well as any additional padding surrounding the halo. + + Notes + ----- + In an MPI context, this is the *local* with_halo region shape. + """ + return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, + self._padding)], + getters=self.dimensions) + + @cached_property + def _C_ctype(self): + # NOTE: Reverting to using float/double instead of PetscScalar for + # simplicity when opt='advanced'. Otherwise, Temp objects must also + # be converted to PetscScalar. Additional tests are needed to + # ensure this approach is fine. Previously, issues arose from + # mismatches between precision of Function objects in Devito and the + # precision of the PETSc configuration. + # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables + # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of + # the user's PETSc configuration. + return POINTER(dtype_to_ctype(self.dtype)) + + @property + def symbolic_shape(self): + field_from_composites = [ + FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] + # Reverse it since DMDA is setup backwards to Devito dimensions. + return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + + @cached_property + def dmda(self): + name = 'da_so_%s' % self.space_order + return DM(name=name, liveness='eager', stencil_width=self.space_order) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py new file mode 100644 index 0000000000..c36664298d --- /dev/null +++ b/devito/petsc/types/object.py @@ -0,0 +1,171 @@ +from ctypes import POINTER + +from devito.tools import CustomDtype +from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension +from devito.symbolics import Byref + +from devito.petsc.iet.utils import petsc_call + + +class DM(LocalObject): + """ + PETSc Data Management object (DM). + """ + dtype = CustomDtype('DM') + + def __init__(self, *args, stencil_width=None, **kwargs): + super().__init__(*args, **kwargs) + self._stencil_width = stencil_width + + @property + def stencil_width(self): + return self._stencil_width + + @property + def info(self): + return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') + + @property + def _C_free(self): + return petsc_call('DMDestroy', [Byref(self.function)]) + + @property + def _C_free_priority(self): + return 3 + + +class Mat(LocalObject): + """ + PETSc Matrix object (Mat). + """ + dtype = CustomDtype('Mat') + + @property + def _C_free(self): + return petsc_call('MatDestroy', [Byref(self.function)]) + + @property + def _C_free_priority(self): + return 1 + + +class LocalVec(LocalObject): + """ + PETSc Vector object (Vec). + """ + dtype = CustomDtype('Vec') + + +class GlobalVec(LocalObject): + """ + PETSc Vector object (Vec). + """ + dtype = CustomDtype('Vec') + + @property + def _C_free(self): + return petsc_call('VecDestroy', [Byref(self.function)]) + + @property + def _C_free_priority(self): + return 0 + + +class PetscMPIInt(LocalObject): + """ + PETSc datatype used to represent `int` parameters + to MPI functions. + """ + dtype = CustomDtype('PetscMPIInt') + + +class PetscInt(LocalObject): + """ + PETSc datatype used to represent `int` parameters + to PETSc functions. + """ + dtype = CustomDtype('PetscInt') + + +class KSP(LocalObject): + """ + PETSc KSP : Linear Systems Solvers. + Manages Krylov Methods. + """ + dtype = CustomDtype('KSP') + + +class SNES(LocalObject): + """ + PETSc SNES : Non-Linear Systems Solvers. + """ + dtype = CustomDtype('SNES') + + @property + def _C_free(self): + return petsc_call('SNESDestroy', [Byref(self.function)]) + + @property + def _C_free_priority(self): + return 2 + + +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 DMDALocalInfo(LocalObject): + """ + PETSc object - C struct containing information + about the local grid. + """ + dtype = CustomDtype('DMDALocalInfo') + + +class PetscErrorCode(LocalObject): + """ + PETSc datatype used to return PETSc error codes. + https://petsc.org/release/manualpages/Sys/PetscErrorCode/ + """ + dtype = CustomDtype('PetscErrorCode') + + +class DummyArg(LocalObject): + dtype = CustomDtype('void', modifier='*') + + +class PETScStruct(CCompositeObject): + + __rargs__ = ('name', 'pname', 'fields') + + def __init__(self, name, pname, fields, liveness='lazy'): + pfields = [(i._C_name, i._C_ctype) for i in fields] + super().__init__(name, pname, pfields, liveness) + self._fields = fields + + @property + def fields(self): + return self._fields + + @property + def time_dim_fields(self): + return [f for f in self.fields + if isinstance(f, (ModuloDimension, TimeDimension))] + + @property + def _C_ctype(self): + return POINTER(self.dtype) if self.liveness == \ + 'eager' else self.dtype + + _C_modifier = ' *' diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py new file mode 100644 index 0000000000..dca9140077 --- /dev/null +++ b/devito/petsc/types/types.py @@ -0,0 +1,97 @@ +import sympy + +from devito.tools import Reconstructable, sympy_mutex + + +class LinearSolveExpr(sympy.Function, Reconstructable): + + __rargs__ = ('expr',) + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays', 'time_dim') + + defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'jacobi', + 'ksp_rtol': 1e-7, # Relative tolerance + 'ksp_atol': 1e-50, # Absolute tolerance + 'ksp_divtol': 1e4, # Divergence tolerance + 'ksp_max_it': 10000 # Maximum iterations + } + + def __new__(cls, expr, target=None, solver_parameters=None, + matvecs=None, formfuncs=None, formrhs=None, arrays=None, time_dim=None, **kwargs): + + if solver_parameters is None: + solver_parameters = cls.defaults + else: + for key, val in cls.defaults.items(): + solver_parameters[key] = solver_parameters.get(key, val) + + with sympy_mutex: + obj = sympy.Function.__new__(cls, expr) + + obj._expr = expr + obj._target = target + obj._solver_parameters = solver_parameters + obj._matvecs = matvecs + obj._formfuncs = formfuncs + obj._formrhs = formrhs + obj._arrays = arrays + obj._time_dim = time_dim + return obj + + def __repr__(self): + return "%s(%s)" % (self.__class__.__name__, self.expr) + + __str__ = __repr__ + + def _sympystr(self, printer): + return str(self) + + def __hash__(self): + return hash(self.expr) + + def __eq__(self, other): + return (isinstance(other, LinearSolveExpr) and + self.expr == other.expr and + self.target == other.target) + + @property + def expr(self): + return self._expr + + @property + def target(self): + return self._target + + @property + def solver_parameters(self): + return self._solver_parameters + + @property + def matvecs(self): + return self._matvecs + + @property + def formfuncs(self): + return self._formfuncs + + @property + def formrhs(self): + return self._formrhs + + @property + def arrays(self): + return self._arrays + + @classmethod + def eval(cls, *args): + return None + + func = Reconstructable._rebuild + + +class CallbackExpr(sympy.Function): + @classmethod + def eval(cls, *args): + return None diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index b24ad56a3d..d898db23cb 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -1,14 +1,6 @@ import os -from devito.ir.equations import OpInjectSolve from devito.tools import memoized_func -from devito.ir.iet import Call, FindSymbols -from devito.petsc.iet.nodes import PETScCallable, InjectSolveDummy -from devito.petsc.types import PETScStruct - -# Mapping special Eq operations to their corresponding IET Expression subclass types. -# These operations correspond to subclasses of Eq utilised within PETScSolve. -petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} solver_mapper = { diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index e7b395d436..2c366a389e 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,9 +242,6 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) - def _print_PETScRHS(self, expr): - return "%s" % self._print(expr._expr) - _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/devito/types/array.py b/devito/types/array.py index f478add0e3..44d7fabf9d 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -20,8 +20,6 @@ class ArrayBasic(AbstractFunction, LocalType): __rkwargs__ = AbstractFunction.__rkwargs__ + ('is_const', 'liveness') def __init_finalize__(self, *args, **kwargs): - - self._liveness = kwargs.setdefault('liveness', 'lazy') super().__init_finalize__(*args, **kwargs) self._liveness = kwargs.get('liveness', 'lazy') @@ -62,18 +60,6 @@ def shape_allocated(self): def is_const(self): return self._is_const - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' - class Array(ArrayBasic): diff --git a/devito/types/basic.py b/devito/types/basic.py index 14e3949fb8..660ceff437 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1513,6 +1513,37 @@ def __getnewargs_ex__(self): return args, kwargs +class LocalType(Basic): + """ + This is the abstract base class for local types, which are + generated by the compiler in C rather than in Python. + Notes + ----- + Subclasses should setup `_liveness`. + """ + + is_LocalType = True + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' + + """ + A modifier added to the subclass C declaration when it appears + in a function signature. For example, a subclass might define `_C_modifier = '&'` + to impose pass-by-reference semantics. + """ + _C_modifier = None + + # Extended SymPy hierarchy follows, for essentially two reasons: # - To keep track of `function` # - To override SymPy caching behaviour diff --git a/devito/types/equation.py b/devito/types/equation.py index 3b49213f92..947d5072d0 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -219,3 +219,7 @@ class ReduceMax(Reduction): class ReduceMin(Reduction): pass + + +class InjectSolveEq(Eq): + pass diff --git a/devito/types/object.py b/devito/types/object.py index cba54b0add..aa738bd19b 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -1,5 +1,4 @@ from ctypes import byref - import sympy from devito.tools import Pickable, as_tuple, sympy_mutex @@ -8,7 +7,8 @@ from devito.types.basic import Basic, LocalType from devito.types.utils import CtypesFactory -__all__ = ['Object', 'LocalObject', 'CompositeObject'] + +__all__ = ['Object', 'LocalObject', 'CompositeObject', 'CCompositeObject'] class AbstractObject(Basic, sympy.Basic, Pickable): @@ -138,6 +138,7 @@ def __init__(self, name, pname, pfields, value=None): dtype = CtypesFactory.generate(pname, pfields) value = self.__value_setup__(dtype, value) super().__init__(name, dtype, value) + self._pname = pname def __value_setup__(self, dtype, value): return value or byref(dtype._type_()) @@ -148,7 +149,7 @@ def pfields(self): @property def pname(self): - return self.dtype._type_.__name__ + return self._pname @property def fields(self): @@ -231,6 +232,28 @@ def _C_free(self): """ return None + @property + def _C_free_priority(self): + return float('inf') + @property def _mem_global(self): return self._is_global + + +class CCompositeObject(CompositeObject, LocalType): + + """ + Object with composite type (e.g., a C struct) defined in C. + """ + + __rargs__ = ('name', 'pname', 'pfields') + + def __init__(self, name, pname, pfields, liveness='lazy'): + super().__init__(name, pname, pfields) + assert liveness in ['eager', 'lazy'] + self._liveness = liveness + + @property + def dtype(self): + return self._dtype._type_ diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc deleted file mode 100644 index 782e0ac340..0000000000 --- a/docker/Dockerfile.petsc +++ /dev/null @@ -1,49 +0,0 @@ -############################################################## -# This Dockerfile builds a base image to run Devito with PETSc -############################################################# - -# Base image -FROM ubuntu:22.04 as base - -ENV DEBIAN_FRONTEND noninteractive - -# Install python -RUN apt-get update && \ - apt-get install -y dh-autoreconf python3-venv python3-dev python3-pip libopenblas-serial-dev git pkgconf mpich - -# Install for basic base not containing it -RUN apt-get install -y vim wget git flex libnuma-dev tmux \ - numactl hwloc curl \ - autoconf libtool build-essential procps - -# Install tmpi -RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi - -# Install OpenGL library, necessary for the installation of GemPy -RUN apt-get install -y libgl1-mesa-glx - -# Install numpy from source -RUN python3 -m venv /venv && \ - /venv/bin/pip install --no-cache-dir --upgrade pip && \ - /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ - rm -rf ~/.cache/pip - -# Create a directory for PETSc -RUN mkdir -p /home/petsc - -# Clone and install PETSc -RUN cd /home/petsc \ - && git clone -b release https://gitlab.com/petsc/petsc.git petsc \ - && cd petsc \ - && git pull \ - && ./configure --with-fortran-bindings=0 --with-openblas-include=$(pkgconf --variable=includedir openblas) --with-openblas-lib=$(pkgconf --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build \ - && make all - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -EXPOSE 8888 -CMD ["/bin/bash"] - -ENV PETSC_ARCH="devito_build" -ENV PETSC_DIR="/home/petsc/petsc" diff --git a/tests/test_iet.py b/tests/test_iet.py index eac5316087..c1c78eb7ea 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -134,7 +134,7 @@ def test_callback_cgen(): class FunctionPtr(Callback): @property def callback_form(self): - param_types = ', '.join([str(t) for t in + param_types = ', '.join([str(t) for t in self.param_types]) return "(%s (*)(%s))%s" % (self.retval, param_types, self.name) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c639e26783..16e268b0b3 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -121,9 +121,9 @@ def test_petsc_solve(): callable_roots = [meta_call.root for meta_call in op._func_table.values()] - matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_f'] + matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_0'] - formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS_f'] + formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS_0'] action_expr = FindNodes(Expression).visit(matvec_callback[0]) rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) @@ -195,8 +195,6 @@ def test_petsc_cast(): arr3 = PETScArray(name='arr3', dimensions=g1.dimensions, shape=g1.shape, space_order=4) - # Casts will be explictly generated and placed at specific locations in the C code, - # specifically after various other PETSc calls have been executed. cast0 = PointerCast(arr0) cast1 = PointerCast(arr1) cast2 = PointerCast(arr2) @@ -441,14 +439,14 @@ def test_callback_arguments(): with switchconfig(openmp=False): op = Operator(petsc1) - mv = op._func_table['MyMatShellMult_f1'].root - ff = op._func_table['FormFunction_f1'].root + mv = op._func_table['MyMatShellMult_0'].root + ff = op._func_table['FormFunction_0'].root assert len(mv.parameters) == 3 assert len(ff.parameters) == 4 - assert str(mv.parameters) == '(J_f1, X_global_f1, Y_global_f1)' - assert str(ff.parameters) == '(snes_f1, X_global_f1, Y_global_f1, dummy_f1)' + assert str(mv.parameters) == '(J_0, X_global_0, Y_global_0)' + assert str(ff.parameters) == '(snes_0, X_global_0, Y_global_0, dummy_0)' @skipif('petsc') @@ -485,7 +483,7 @@ def test_petsc_struct(): @skipif('petsc') -@pytest.mark.parallel(mode=1) +@pytest.mark.parallel(mode=[2, 4, 8]) def test_apply(mode): grid = Grid(shape=(13, 13), dtype=np.float64) @@ -499,7 +497,7 @@ def test_apply(mode): petsc = PETScSolve(eqn, pn) # Build the op - with switchconfig(openmp=False): + with switchconfig(openmp=False, mpi=True): op = Operator(petsc) # Check the Operator runs without errors. Not verifying output for @@ -528,10 +526,10 @@ def test_petsc_frees(): frees = op.body.frees # Check the frees appear in the following order - assert str(frees[0]) == 'PetscCall(VecDestroy(&(b_global_f)));' - assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_f)));' - assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_f)));' - assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_f)));' + assert str(frees[0]) == 'PetscCall(VecDestroy(&(b_global_0)));' + assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_0)));' + assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_0)));' + assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_0)));' assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_so_2)));' @@ -551,5 +549,101 @@ def test_calls_to_callbacks(): ccode = str(op.ccode) - assert '(void (*)(void))MyMatShellMult_f' in ccode - assert 'PetscCall(SNESSetFunction(snes_f,NULL,FormFunction_f,NULL));' in ccode + assert '(void (*)(void))MyMatShellMult_0' in ccode + assert 'PetscCall(SNESSetFunction(snes_0,NULL,FormFunction_0,NULL));' in ccode + + +@skipif('petsc') +def test_start_ptr(): + """ + Verify that a pointer to the start of the memory address is correctly + generated for TimeFunction objects. This pointer should indicate the + beginning of the multidimensional array that will be overwritten at + the current time step. + This functionality is crucial for VecReplaceArray operations, as it ensures + that the correct memory location is accessed and modified during each time step. + """ + grid = Grid((11, 11)) + u1 = TimeFunction(name='u1', grid=grid, space_order=2, dtype=np.float32) + eq1 = Eq(u1.dt, u1.laplace, subdomain=grid.interior) + petsc1 = PETScSolve(eq1, u1.forward) + + with switchconfig(openmp=False): + op1 = Operator(petsc1) + + # Verify the case with modulo time stepping + assert 'float * start_ptr_0 = t1*localsize_0 + (float *)(u1_vec->data);' in str(op1) + + # Verify the case with no modulo time stepping + u2 = TimeFunction(name='u2', grid=grid, space_order=2, dtype=np.float32, save=5) + eq2 = Eq(u2.dt, u2.laplace, subdomain=grid.interior) + petsc2 = PETScSolve(eq2, u2.forward) + + with switchconfig(openmp=False): + op2 = Operator(petsc2) + + assert 'float * start_ptr_0 = (time + 1)*localsize_0 + ' + \ + '(float *)(u2_vec->data);' in str(op2) + + +@skipif('petsc') +def test_time_loop(): + """ + Verify the following: + - Modulo dimensions are correctly assigned and updated in the PETSc struct + at each time step. + - Only assign/update the modulo dimensions required by any of the + PETSc callback functions. + """ + grid = Grid((11, 11)) + + # Modulo time stepping + u1 = TimeFunction(name='u1', grid=grid, space_order=2) + v1 = Function(name='v1', grid=grid, space_order=2) + eq1 = Eq(v1.laplace, u1) + petsc1 = PETScSolve(eq1, v1) + op1 = Operator(petsc1) + body1 = str(op1.body) + rhs1 = str(op1._func_table['FormRHS_0'].root.ccode) + + assert 'ctx.t0 = t0' in body1 + assert 'ctx.t1 = t1' not in body1 + assert 'formrhs->t0' in rhs1 + assert 'formrhs->t1' not in rhs1 + + # Non-modulo time stepping + u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) + v2 = Function(name='v2', grid=grid, space_order=2, save=5) + eq2 = Eq(v2.laplace, u2) + petsc2 = PETScSolve(eq2, v2) + op2 = Operator(petsc2) + body2 = str(op2.body) + rhs2 = str(op2._func_table['FormRHS_0'].root.ccode) + + assert 'ctx.time = time' in body2 + assert 'formrhs->time' in rhs2 + + # Modulo time stepping with more than one time step + # used in one of the callback functions + eq3 = Eq(v1.laplace, u1 + u1.forward) + petsc3 = PETScSolve(eq3, v1) + op3 = Operator(petsc3) + body3 = str(op3.body) + rhs3 = str(op3._func_table['FormRHS_0'].root.ccode) + + assert 'ctx.t0 = t0' in body3 + assert 'ctx.t1 = t1' in body3 + assert 'formrhs->t0' in rhs3 + assert 'formrhs->t1' in rhs3 + + # Multiple petsc solves within the same time loop + v2 = Function(name='v2', grid=grid, space_order=2) + eq4 = Eq(v1.laplace, u1) + petsc4 = PETScSolve(eq4, v1) + eq5 = Eq(v2.laplace, u1) + petsc5 = PETScSolve(eq5, v2) + op4 = Operator(petsc4 + petsc5) + body4 = str(op4.body) + + assert 'ctx.t0 = t0' in body4 + assert body4.count('ctx.t0 = t0') == 1 From feba0a9d9c2d368188c632c490feb06656d924c7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 2 Oct 2024 20:22:24 +0100 Subject: [PATCH 098/151] compiler: Don't ever generate time loop for petsc callbacks --- devito/operator/operator.py | 4 +- devito/passes/iet/mpi.py | 5 +- devito/petsc/clusters.py | 22 +-------- devito/petsc/iet/passes.py | 21 ++++---- devito/petsc/iet/routines.py | 92 +++++++++++++----------------------- devito/petsc/iet/utils.py | 30 +++++------- devito/petsc/solve.py | 46 ++++++++++++++---- devito/petsc/types/object.py | 8 +++- devito/petsc/types/types.py | 17 ++++--- tests/test_petsc.py | 16 ++++--- 10 files changed, 125 insertions(+), 136 deletions(-) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 16c1191d6b..22dcad3dc5 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -382,6 +382,8 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Build a sequence of Clusters from a sequence of Eqs clusters = clusterize(expressions, **kwargs) + # Lift iteration space surrounding each PETSc solve to create + # distinct iteration loops clusters = petsc_lift(clusters) # Operation count before specialization @@ -389,8 +391,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): clusters = cls._specialize_clusters(clusters, **kwargs) - clusters = petsc_project(clusters) - # Operation count after specialization final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) try: diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py index 3cedd0fb23..180d2231a6 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -72,8 +72,9 @@ def _hoist_invariant(iet): """ # Precompute scopes to save time - scopes = {i: Scope([e.expr for e in v if not isinstance(e, Call)]) - for i, v in MapNodes().visit(iet).items()} + scopes = {} + for i, v in MapNodes(child_types=Expression).visit(iet).items(): + scopes[i] = Scope([e.expr for e in v]) # Analysis hsmapper = {} diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index bfec3dc52a..8c83aa098c 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -1,11 +1,11 @@ from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr, CallbackExpr +from devito.petsc.types import LinearSolveExpr @timed_pass() def petsc_lift(clusters): """ - Lift the iteration space surrounding each PETSc equation to create + Lift the iteration space surrounding each PETSc solve to create distinct iteration loops. """ processed = [] @@ -15,22 +15,4 @@ def petsc_lift(clusters): processed.append(c.rebuild(ispace=ispace)) else: processed.append(c) - - return processed - - -@timed_pass() -def petsc_project(clusters): - """ - Drop time loop for clusters which appear in PETSc callback functions. - """ - processed = [] - for c in clusters: - if isinstance(c.exprs[0].rhs, CallbackExpr): - time_dims = [d for d in c.ispace.intervals.dimensions if d.is_Time] - ispace = c.ispace.project(lambda d: d not in time_dims) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - return processed diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index d9b875c477..ddd6cd2060 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -6,13 +6,14 @@ from devito.symbolics import Byref, Macro, FieldFromPointer from devito.tools import filter_ordered from devito.petsc.types import (PetscMPIInt, DM, Mat, LocalVec, GlobalVec, - KSP, PC, SNES, PetscErrorCode, DummyArg, PetscInt) + KSP, PC, SNES, PetscErrorCode, DummyArg, PetscInt, + StartPtr) from devito.petsc.iet.nodes import InjectSolveDummy from devito.petsc.utils import solver_mapper, core_metadata from devito.petsc.iet.routines import PETScCallbackBuilder from devito.petsc.iet.utils import (petsc_call, petsc_call_mpi, petsc_struct, spatial_injectsolve_iter, assign_time_iters, - retrieve_mod_dims) + retrieve_time_dims) @iet_pass @@ -45,15 +46,12 @@ def lower_petsc(iet, **kwargs): builder = PETScCallbackBuilder(**kwargs) for iters, (injectsolve,) in injectsolve_mapper.items(): - target = injectsolve.expr.rhs.target - solver_objs = build_solver_objs(target, **kwargs) + solver_objs = build_solver_objs(injectsolve, iters, **kwargs) # Generate the solver setup for each InjectSolveDummy solver_setup = generate_solver_setup(solver_objs, objs, injectsolve) setup.extend(solver_setup) - # Retrieve `ModuloDimensions` for use in callback functions - solver_objs['mod_dims'] = retrieve_mod_dims(iters) # Generate all PETSc callback functions for the target via recursive compilation matvec_op, formfunc_op, runsolve = builder.make(injectsolve, objs, solver_objs) @@ -169,7 +167,8 @@ def create_dmda(dmda, objs): return dmda -def build_solver_objs(target, **kwargs): +def build_solver_objs(injectsolve, iters, **kwargs): + target = injectsolve.expr.rhs.target sreg = kwargs['sregistry'] return { 'Jac': Mat(sreg.make_name(prefix='J_')), @@ -185,12 +184,16 @@ def build_solver_objs(target, **kwargs): 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), 'dummy': DummyArg(sreg.make_name(prefix='dummy_')), - 'localsize': PetscInt(sreg.make_name(prefix='localsize_')) + 'localsize': PetscInt(sreg.make_name(prefix='localsize_')), + 'start_ptr': StartPtr(sreg.make_name(prefix='start_ptr_'), target.dtype), + 'true_dims': retrieve_time_dims(iters), + 'target': target, + 'time_mapper': injectsolve.expr.rhs.time_mapper, } def generate_solver_setup(solver_objs, objs, injectsolve): - target = injectsolve.expr.rhs.target + target = solver_objs['target'] dmda = objs['da_so_%s' % target.space_order] diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index d164d9663f..ffe8132e9f 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -4,15 +4,14 @@ from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine) -from devito.symbolics import Byref, FieldFromPointer, Macro, Cast +from devito.symbolics import Byref, FieldFromPointer, Macro, cast_mapper from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction -from devito.types import LocalObject, ModuloDimension, TimeDimension -from devito.tools import ctypes_to_cstr, dtype_to_ctype, CustomDtype +from devito.types import ModuloDimension, TimeDimension, Temp from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatVecCallback) -from devito.petsc.iet.utils import petsc_call, petsc_struct, drop_callbackexpr +from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.ir.support import SymbolRegistry @@ -86,11 +85,9 @@ def make_matvec(self, injectsolve, objs, solver_objs): def create_matvec_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs - body = drop_callbackexpr(body) - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] - body = uxreplace_mod_dims(body, solver_objs['mod_dims'], target=linsolveexpr.target) + body = uxreplace_time(body, solver_objs) struct = build_petsc_struct(body, 'matvec', liveness='eager') @@ -215,11 +212,9 @@ def make_formfunc(self, injectsolve, objs, solver_objs): def create_formfunc_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs - body = drop_callbackexpr(body) - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] - body = uxreplace_mod_dims(body, solver_objs['mod_dims'], linsolveexpr.target) + body = uxreplace_time(body, solver_objs) struct = build_petsc_struct(body, 'formfunc', liveness='eager') @@ -335,8 +330,6 @@ def make_formrhs(self, injectsolve, objs, solver_objs): def create_formrhs_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs - body = drop_callbackexpr(body) - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) @@ -351,7 +344,7 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] ) - body = uxreplace_mod_dims(body, solver_objs['mod_dims'], linsolveexpr.target) + body = uxreplace_time(body, solver_objs) struct = build_petsc_struct(body, 'formrhs', liveness='eager') @@ -386,6 +379,7 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for i in struct.fields if not isinstance(i.function, AbstractFunction)} + formrhs_body = Uxreplace(subs).visit(formrhs_body) self._struct_params.extend(struct.fields) @@ -404,7 +398,7 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): if any(i.is_Time for i in target.dimensions): vec_replace_array = time_dep_replace( - injectsolve, target, solver_objs, objs, self.sregistry + injectsolve, solver_objs, objs, self.sregistry ) else: field_from_ptr = FieldFromPointer(target._C_field_data, target._C_symbol) @@ -444,43 +438,34 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): def build_petsc_struct(iet, name, liveness): # Place all context data required by the shell routines into a struct - # TODO: Clean this search up - basics = FindSymbols('basics').visit(iet) - time_dims = [i for i in FindSymbols('dimensions').visit(iet) - if isinstance(i, (TimeDimension, ModuloDimension))] - avoid0 = [i for i in FindSymbols('indexedbases').visit(iet) - if isinstance(i.function, PETScArray)] - avoid1 = [i for i in FindSymbols('dimensions|writes').visit(iet) - if i not in time_dims] - fields = [data.function for data in basics if data not in avoid0+avoid1] + fields = [ + i.function for i in FindSymbols('basics').visit(iet) + if not isinstance(i.function, (PETScArray, Temp)) + and not (i.is_Dimension and not isinstance(i, (TimeDimension, ModuloDimension))) + ] return petsc_struct(name, fields, liveness) -def time_dep_replace(injectsolve, target, solver_objs, objs, sregistry): - target_time = injectsolve.expr.lhs - target_time = [i for i, d in zip(target_time.indices, - target_time.dimensions) if d.is_Time] +def time_dep_replace(injectsolve, solver_objs, objs, sregistry): + target = injectsolve.expr.lhs + target_time = [ + i for i, d in zip(target.indices, target.dimensions) if d.is_Time + ] assert len(target_time) == 1 target_time = target_time.pop() - ctype_str = ctypes_to_cstr(dtype_to_ctype(target.dtype)) - - class BarCast(Cast): - _base_typ = ctype_str - - class StartPtr(LocalObject): - dtype = CustomDtype(ctype_str, modifier=' *') - - start_ptr = StartPtr(sregistry.make_name(prefix='start_ptr_')) + start_ptr = solver_objs['start_ptr'] vec_get_size = petsc_call( 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] ) - # TODO: What is the correct way to use Mul here? Devito Mul? Sympy Mul? - field_from_ptr = FieldFromPointer(target._C_field_data, target._C_symbol) + field_from_ptr = FieldFromPointer( + target.function._C_field_data, target.function._C_symbol + ) + expr = DummyExpr( - start_ptr, BarCast(field_from_ptr, ' *') + + start_ptr, cast_mapper[(target.dtype, '*')](field_from_ptr) + Mul(target_time, solver_objs['localsize']), init=True ) @@ -488,25 +473,16 @@ class StartPtr(LocalObject): return (vec_get_size, expr, vec_replace_array) -def uxreplace_mod_dims(body, mod_dims, target): - """ - Replace ModuloDimensions in callback functions with the corresponding - ModuloDimensions generated by the initial lowering. They must match because - they are assigned and updated in the struct at each time step. This is a valid - uxreplace because all functions appearing in the callback functions are - passed through the initial lowering. - """ - # old_mod_dims = [ - # i for i in FindSymbols('dimensions').visit(body) if isinstance(i, ModuloDimension) - # ] - # from IPython import embed; embed() - # if not old_mod_dims: - # return body - # t_tmp = - # from IPython import embed; embed() - # body = Uxreplace({i: mod_dims[i] for i in old_mod_dims}).visit(body) - # return Uxreplace({i: mod_dims[i.origin] for i in old_mod_dims}).visit(body) - return body +def uxreplace_time(body, solver_objs): + time_spacing = solver_objs['target'].grid.stepping_dim.spacing + true_dims = solver_objs['true_dims'] + + time_mapper = { + v: k.xreplace({time_spacing: 1, -time_spacing: -1}) + for k, v in solver_objs['time_mapper'].items() + } + subs = {symb: true_dims[time_mapper[symb]] for symb in time_mapper} + return Uxreplace(subs).visit(body) Null = Macro('NULL') diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 8f1b0a34aa..16897f1501 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,4 +1,3 @@ -from devito.ir.iet.nodes import Expression from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall from devito.ir.equations import OpInjectSolve from devito.ir.iet import (FindNodes, retrieve_iteration_tree, @@ -36,18 +35,6 @@ def spatial_injectsolve_iter(iter, injectsolve): petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} -def drop_callbackexpr(body): - # TODO: fix this circular import - from devito.petsc.types import CallbackExpr - nodes = FindNodes(Expression).visit(body) - mapper = { - expr: expr._rebuild(expr=expr.expr._rebuild(rhs=expr.expr.rhs.args[0])) - for expr in nodes - if isinstance(expr.expr.rhs, CallbackExpr) - } - return Transformer(mapper).visit(body) - - def assign_time_iters(iet, struct): """ Assign time iterators to the struct within loops containing PETScCalls. @@ -75,9 +62,14 @@ def assign_time_iters(iet, struct): return Transformer(mapper).visit(iet) -def retrieve_mod_dims(iters): - outer_iter_dims = iters[0].dimensions - if any(dim.is_Time for dim in outer_iter_dims): - mod_dims = [dim for dim in outer_iter_dims if dim.is_Modulo] - return {dim.origin: dim for dim in mod_dims} - return {} +def retrieve_time_dims(iters): + time_iter = [i for i in iters if any(dim.is_Time for dim in i.dimensions)] + mapper = {} + if not time_iter: + return mapper + for dim in time_iter[0].dimensions: + if dim.is_Modulo: + mapper[dim.origin] = dim + elif dim.is_Time: + mapper[dim] = dim + return mapper diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index fd5403e319..17c03ae435 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -4,11 +4,12 @@ from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative -from devito.types import Eq, Symbol +from devito.types import Eq, Symbol, SteppingDimension from devito.types.equation import InjectSolveEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions -from devito.petsc.types import LinearSolveExpr, PETScArray, CallbackExpr +from devito.tools import as_tuple +from devito.petsc.types import LinearSolveExpr, PETScArray __all__ = ['PETScSolve'] @@ -32,42 +33,45 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): formfuncs = [] formrhs = [] - eqns = eqns if isinstance(eqns, (list, tuple)) else [eqns] + eqns = as_tuple(eqns) + funcs = retrieve_functions(eqns) + time_mapper = generate_time_mapper(funcs) + for eq in eqns: b, F_target = separate_eqn(eq, target) + b, F_target = b.subs(time_mapper), F_target.subs(time_mapper) # TODO: Current assumption is that problem is linear and user has not provided # a jacobian. Hence, we can use F_target to form the jac-vec product matvecs.append(Eq( arrays['y_matvec'], - CallbackExpr(F_target.subs({target: arrays['x_matvec']})), + F_target.subs({target: arrays['x_matvec']}), subdomain=eq.subdomain )) formfuncs.append(Eq( arrays['y_formfunc'], - CallbackExpr(F_target.subs({target: arrays['x_formfunc']})), + F_target.subs({target: arrays['x_formfunc']}), subdomain=eq.subdomain )) - time_dim = Symbol('t_tmp') - time_mapper = {target.grid.stepping_dim: time_dim} + formrhs.append(Eq( arrays['b_tmp'], - CallbackExpr(b).subs(time_mapper), + b, subdomain=eq.subdomain )) # Placeholder equation for inserting calls to the solver and generating # correct time loop etc inject_solve = InjectSolveEq(target, LinearSolveExpr( - expr=tuple(retrieve_functions(eqns)), + expr=tuple(funcs), target=target, solver_parameters=solver_parameters, matvecs=matvecs, formfuncs=formfuncs, formrhs=formrhs, arrays=arrays, - time_dim=time_dim + time_mapper=time_mapper, ), subdomain=eq.subdomain) return [inject_solve] @@ -153,3 +157,25 @@ def _(expr, target): return 0 args = [centre_stencil(a, target) for a in expr.evaluate.args] return expr.evaluate.func(*args) + + +def generate_time_mapper(funcs): + """ + Replace time indices with `Symbols` in equations used within + PETSc callback functions. These symbols are Uxreplaced at the IET + level to align with the `TimeDimension` and `ModuloDimension` objects + present in the inital lowering. + NOTE: All functions used in PETSc callback functions are attached to + the `LinearSolveExpr` object, which is passed through the initial lowering + (and subsequently dropped and replaced with calls to run the solver). + Therefore, the appropriate time loop will always be correctly generated inside + the main kernel. + """ + time_indices = list({ + i if isinstance(d, SteppingDimension) else d + for f in funcs + for i, d in zip(f.indices, f.dimensions) + if d.is_Time + }) + tau_symbs = [Symbol('tau%d' % i) for i in range(len(time_indices))] + return {time: tau for time, tau in zip(time_indices, tau_symbs)} diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index c36664298d..9f8cbe4cbb 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,6 +1,6 @@ from ctypes import POINTER -from devito.tools import CustomDtype +from devito.tools import CustomDtype, dtype_to_cstr from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension from devito.symbolics import Byref @@ -169,3 +169,9 @@ def _C_ctype(self): 'eager' else self.dtype _C_modifier = ' *' + + +class StartPtr(LocalObject): + def __init__(self, name, dtype): + super().__init__(name=name) + self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index dca9140077..eda2fa40d4 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -7,7 +7,7 @@ class LinearSolveExpr(sympy.Function, Reconstructable): __rargs__ = ('expr',) __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays', 'time_dim') + 'formfuncs', 'formrhs', 'arrays', 'time_mapper') defaults = { 'ksp_type': 'gmres', @@ -19,7 +19,8 @@ class LinearSolveExpr(sympy.Function, Reconstructable): } def __new__(cls, expr, target=None, solver_parameters=None, - matvecs=None, formfuncs=None, formrhs=None, arrays=None, time_dim=None, **kwargs): + matvecs=None, formfuncs=None, formrhs=None, + arrays=None, time_mapper=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -37,7 +38,7 @@ def __new__(cls, expr, target=None, solver_parameters=None, obj._formfuncs = formfuncs obj._formrhs = formrhs obj._arrays = arrays - obj._time_dim = time_dim + obj._time_mapper = time_mapper return obj def __repr__(self): @@ -84,14 +85,12 @@ def formrhs(self): def arrays(self): return self._arrays + @property + def time_mapper(self): + return self._time_mapper + @classmethod def eval(cls, *args): return None func = Reconstructable._rebuild - - -class CallbackExpr(sympy.Function): - @classmethod - def eval(cls, *args): - return None diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 16e268b0b3..f5616d4f36 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -572,7 +572,7 @@ def test_start_ptr(): op1 = Operator(petsc1) # Verify the case with modulo time stepping - assert 'float * start_ptr_0 = t1*localsize_0 + (float *)(u1_vec->data);' in str(op1) + assert 'float * start_ptr_0 = t1*localsize_0 + (float*)(u1_vec->data);' in str(op1) # Verify the case with no modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, dtype=np.float32, save=5) @@ -583,7 +583,7 @@ def test_start_ptr(): op2 = Operator(petsc2) assert 'float * start_ptr_0 = (time + 1)*localsize_0 + ' + \ - '(float *)(u2_vec->data);' in str(op2) + '(float*)(u2_vec->data);' in str(op2) @skipif('petsc') @@ -602,7 +602,8 @@ def test_time_loop(): v1 = Function(name='v1', grid=grid, space_order=2) eq1 = Eq(v1.laplace, u1) petsc1 = PETScSolve(eq1, v1) - op1 = Operator(petsc1) + with switchconfig(openmp=False): + op1 = Operator(petsc1) body1 = str(op1.body) rhs1 = str(op1._func_table['FormRHS_0'].root.ccode) @@ -616,7 +617,8 @@ def test_time_loop(): v2 = Function(name='v2', grid=grid, space_order=2, save=5) eq2 = Eq(v2.laplace, u2) petsc2 = PETScSolve(eq2, v2) - op2 = Operator(petsc2) + with switchconfig(openmp=False): + op2 = Operator(petsc2) body2 = str(op2.body) rhs2 = str(op2._func_table['FormRHS_0'].root.ccode) @@ -627,7 +629,8 @@ def test_time_loop(): # used in one of the callback functions eq3 = Eq(v1.laplace, u1 + u1.forward) petsc3 = PETScSolve(eq3, v1) - op3 = Operator(petsc3) + with switchconfig(openmp=False): + op3 = Operator(petsc3) body3 = str(op3.body) rhs3 = str(op3._func_table['FormRHS_0'].root.ccode) @@ -642,7 +645,8 @@ def test_time_loop(): petsc4 = PETScSolve(eq4, v1) eq5 = Eq(v2.laplace, u1) petsc5 = PETScSolve(eq5, v2) - op4 = Operator(petsc4 + petsc5) + with switchconfig(openmp=False): + op4 = Operator(petsc4 + petsc5) body4 = str(op4.body) assert 'ctx.t0 = t0' in body4 From 10602d081d809559388744128732a3d6ad0d5bf3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 4 Oct 2024 12:03:46 +0100 Subject: [PATCH 099/151] compiler: Move passes from utils into passes and create petsc_preprocess --- devito/operator/operator.py | 6 +-- devito/petsc/clusters.py | 9 ++++ devito/petsc/iet/passes.py | 83 ++++++++++++++++++++++++------------ devito/petsc/iet/routines.py | 42 +++++++++++++++--- devito/petsc/iet/utils.py | 53 ----------------------- 5 files changed, 103 insertions(+), 90 deletions(-) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 22dcad3dc5..fe4b5ab48c 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -37,7 +37,6 @@ from devito.petsc.iet.passes import lower_petsc from devito.petsc.clusters import petsc_preprocess - __all__ = ['Operator'] @@ -382,9 +381,8 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Build a sequence of Clusters from a sequence of Eqs clusters = clusterize(expressions, **kwargs) - # Lift iteration space surrounding each PETSc solve to create - # distinct iteration loops - clusters = petsc_lift(clusters) + # Preprocess clusters for PETSc lowering + clusters = petsc_preprocess(clusters) # Operation count before specialization init_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 8c83aa098c..0a6679b2dd 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -3,6 +3,15 @@ @timed_pass() +def petsc_preprocess(clusters): + """ + Preprocess the clusters to make them suitable for PETSc + code generation. + """ + clusters = petsc_lift(clusters) + return clusters + + def petsc_lift(clusters): """ Lift the iteration space surrounding each PETSc solve to create diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index ddd6cd2060..835ad102fb 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -2,18 +2,16 @@ from devito.passes.iet.engine import iet_pass from devito.ir.iet import (Transformer, MapNodes, Iteration, List, BlankLine, - Callable, CallableBody, DummyExpr, Call) -from devito.symbolics import Byref, Macro, FieldFromPointer -from devito.tools import filter_ordered + DummyExpr, FindNodes, retrieve_iteration_tree, + filter_iterations) +from devito.symbolics import Byref, Macro, FieldFromComposite from devito.petsc.types import (PetscMPIInt, DM, Mat, LocalVec, GlobalVec, KSP, PC, SNES, PetscErrorCode, DummyArg, PetscInt, StartPtr) -from devito.petsc.iet.nodes import InjectSolveDummy +from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall from devito.petsc.utils import solver_mapper, core_metadata from devito.petsc.iet.routines import PETScCallbackBuilder -from devito.petsc.iet.utils import (petsc_call, petsc_call_mpi, petsc_struct, - spatial_injectsolve_iter, assign_time_iters, - retrieve_time_dims) +from devito.petsc.iet.utils import petsc_call, petsc_call_mpi @iet_pass @@ -61,16 +59,12 @@ def lower_petsc(iet, **kwargs): subs.update({space_iter: List(body=runsolve)}) # Generate callback to populate main struct object - struct_main = petsc_struct('ctx', filter_ordered(builder.struct_params)) - struct_callback = generate_struct_callback(struct_main) - call_struct_callback = petsc_call(struct_callback.name, [Byref(struct_main)]) - calls_set_app_ctx = [petsc_call('DMSetApplicationContext', [i, Byref(struct_main)]) - for i in unique_dmdas] - setup.extend([call_struct_callback] + calls_set_app_ctx) + struct, struct_calls = builder.make_main_struct(unique_dmdas, objs) + setup.extend(struct_calls) iet = Transformer(subs).visit(iet) - iet = assign_time_iters(iet, struct_main) + iet = assign_time_iters(iet, struct) body = core + tuple(setup) + (BlankLine,) + iet.body.body body = iet.body._rebuild( @@ -79,7 +73,7 @@ def lower_petsc(iet, **kwargs): ) iet = iet._rebuild(body=body) metadata = core_metadata() - efuncs = tuple(builder.efuncs.values())+(struct_callback,) + efuncs = tuple(builder.efuncs.values()) metadata.update({'efuncs': efuncs}) return iet, metadata @@ -260,20 +254,53 @@ def generate_solver_setup(solver_objs, objs, injectsolve): ) -def generate_struct_callback(struct): - body = [ - DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol) - for i in struct.fields if i not in struct.time_dim_fields +def assign_time_iters(iet, struct): + """ + Assign time iterators to the struct within loops containing PETScCalls. + Ensure that assignment occurs only once per time loop, if necessary. + Assign only the iterators that are common between the struct fields + and the actual Iteration. + """ + time_iters = [ + i for i in FindNodes(Iteration).visit(iet) + if i.dim.is_Time and FindNodes(PETScCall).visit(i) ] - struct_callback_body = CallableBody( - List(body=body), init=tuple([petsc_func_begin_user]), - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) - ) - struct_callback = Callable( - 'PopulateMatContext', struct_callback_body, PetscErrorCode(name='err'), - parameters=[struct] - ) - return struct_callback + + if not time_iters: + return iet + + mapper = {} + for iter in time_iters: + common_dims = [dim for dim in iter.dimensions if dim in struct.fields] + common_dims = [ + DummyExpr(FieldFromComposite(dim, struct), dim) for dim in common_dims + ] + iter_new = iter._rebuild(nodes=List(body=tuple(common_dims)+iter.nodes)) + mapper.update({iter: iter_new}) + + return Transformer(mapper).visit(iet) + + +def retrieve_time_dims(iters): + time_iter = [i for i in iters if any(dim.is_Time for dim in i.dimensions)] + mapper = {} + if not time_iter: + return mapper + for dim in time_iter[0].dimensions: + if dim.is_Modulo: + mapper[dim.origin] = dim + elif dim.is_Time: + mapper[dim] = dim + return mapper + + +def spatial_injectsolve_iter(iter, injectsolve): + spatial_body = [] + for tree in retrieve_iteration_tree(iter[0]): + root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] + if injectsolve in FindNodes(InjectSolveDummy).visit(root): + spatial_body.append(root) + return spatial_body Null = Macro('NULL') diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index ffe8132e9f..a516b1bc81 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -3,11 +3,12 @@ import cgen as c from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, - Dereference, DummyExpr, BlankLine) + Dereference, DummyExpr, BlankLine, Callable) from devito.symbolics import Byref, FieldFromPointer, Macro, cast_mapper from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import ModuloDimension, TimeDimension, Temp +from devito.tools import filter_ordered from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatVecCallback) @@ -89,7 +90,7 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): body = uxreplace_time(body, solver_objs) - struct = build_petsc_struct(body, 'matvec', liveness='eager') + struct = build_local_struct(body, 'matvec', liveness='eager') y_matvec = linsolveexpr.arrays['y_matvec'] x_matvec = linsolveexpr.arrays['x_matvec'] @@ -216,7 +217,7 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): body = uxreplace_time(body, solver_objs) - struct = build_petsc_struct(body, 'formfunc', liveness='eager') + struct = build_local_struct(body, 'formfunc', liveness='eager') y_formfunc = linsolveexpr.arrays['y_formfunc'] x_formfunc = linsolveexpr.arrays['x_formfunc'] @@ -346,7 +347,7 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): body = uxreplace_time(body, solver_objs) - struct = build_petsc_struct(body, 'formrhs', liveness='eager') + struct = build_local_struct(body, 'formrhs', liveness='eager') dm_get_app_context = petsc_call( 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] @@ -435,8 +436,36 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): BlankLine, ) + def make_main_struct(self, unique_dmdas, objs): + struct_main = petsc_struct('ctx', filter_ordered(self.struct_params)) + struct_callback = self.generate_struct_callback(struct_main, objs) + call_struct_callback = petsc_call(struct_callback.name, [Byref(struct_main)]) + calls_set_app_ctx = [ + petsc_call('DMSetApplicationContext', [i, Byref(struct_main)]) + for i in unique_dmdas + ] + calls = [call_struct_callback] + calls_set_app_ctx -def build_petsc_struct(iet, name, liveness): + self._efuncs[struct_callback.name] = struct_callback + return struct_main, calls + + def generate_struct_callback(self, struct, objs): + body = [ + DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol) + for i in struct.fields if i not in struct.time_dim_fields + ] + struct_callback_body = CallableBody( + List(body=body), init=tuple([petsc_func_begin_user]), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + ) + struct_callback = Callable( + 'PopulateMatContext', struct_callback_body, objs['err'], + parameters=[struct] + ) + return struct_callback + + +def build_local_struct(iet, name, liveness): # Place all context data required by the shell routines into a struct fields = [ i.function for i in FindSymbols('basics').visit(iet) @@ -474,6 +503,9 @@ def time_dep_replace(injectsolve, solver_objs, objs, sregistry): def uxreplace_time(body, solver_objs): + # TODO: Potentially introduce a TimeIteration abstraction to simplify + # all the time processing that is done (searches, replacements, ...) + # "manually" via free functions time_spacing = solver_objs['target'].grid.stepping_dim.spacing true_dims = solver_objs['true_dims'] diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 16897f1501..a7855fbb36 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,9 +1,5 @@ from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall from devito.ir.equations import OpInjectSolve -from devito.ir.iet import (FindNodes, retrieve_iteration_tree, - filter_iterations, Transformer, Iteration, - DummyExpr, List) -from devito.symbolics import FieldFromComposite def petsc_call(specific_call, call_args): @@ -21,55 +17,6 @@ def petsc_struct(name, fields, liveness='lazy'): fields=fields, liveness=liveness) -def spatial_injectsolve_iter(iter, injectsolve): - spatial_body = [] - for tree in retrieve_iteration_tree(iter[0]): - root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if injectsolve in FindNodes(InjectSolveDummy).visit(root): - spatial_body.append(root) - return spatial_body - - # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} - - -def assign_time_iters(iet, struct): - """ - Assign time iterators to the struct within loops containing PETScCalls. - Ensure that assignment occurs only once per time loop, if necessary. - Assign only the iterators that are common between the struct fields - and the actual Iteration. - """ - time_iters = [ - i for i in FindNodes(Iteration).visit(iet) - if i.dim.is_Time and FindNodes(PETScCall).visit(i) - ] - - if not time_iters: - return iet - - mapper = {} - for iter in time_iters: - common_dims = [dim for dim in iter.dimensions if dim in struct.fields] - common_dims = [ - DummyExpr(FieldFromComposite(dim, struct), dim) for dim in common_dims - ] - iter_new = iter._rebuild(nodes=List(body=tuple(common_dims)+iter.nodes)) - mapper.update({iter: iter_new}) - - return Transformer(mapper).visit(iet) - - -def retrieve_time_dims(iters): - time_iter = [i for i in iters if any(dim.is_Time for dim in i.dimensions)] - mapper = {} - if not time_iter: - return mapper - for dim in time_iter[0].dimensions: - if dim.is_Modulo: - mapper[dim.origin] = dim - elif dim.is_Time: - mapper[dim] = dim - return mapper From ad1fb03df85ebf0b5fd3e80de00eb5a934a2f120 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 4 Oct 2024 14:24:08 +0100 Subject: [PATCH 100/151] misc: Remove workflow triggers --- .github/workflows/examples-mpi.yml | 1 - .github/workflows/examples.yml | 1 - .github/workflows/flake8.yml | 1 - .github/workflows/pytest-core-mpi.yml | 1 - .github/workflows/pytest-core-nompi.yml | 1 - .github/workflows/pytest-petsc.yml | 1 - .github/workflows/tutorials.yml | 1 - 7 files changed, 7 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 6a804a27bb..3bbebabb1c 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,7 +17,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index ef14b8450c..026331c881 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index fe1761f872..95a94f6e79 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index ab89107047..4f2853e46e 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index a022ed08e1..54b5b05b1f 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index a556771c26..6f0a729e95 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index 9e5e1b6cb8..d27af78fcf 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master From 107d86d4f17fa154d1cb3e11713284dd5879197d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 14 Oct 2024 11:01:50 +0100 Subject: [PATCH 101/151] rebase leftover --- devito/types/basic.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/devito/types/basic.py b/devito/types/basic.py index 660ceff437..14e3949fb8 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1513,37 +1513,6 @@ def __getnewargs_ex__(self): return args, kwargs -class LocalType(Basic): - """ - This is the abstract base class for local types, which are - generated by the compiler in C rather than in Python. - Notes - ----- - Subclasses should setup `_liveness`. - """ - - is_LocalType = True - - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' - - """ - A modifier added to the subclass C declaration when it appears - in a function signature. For example, a subclass might define `_C_modifier = '&'` - to impose pass-by-reference semantics. - """ - _C_modifier = None - - # Extended SymPy hierarchy follows, for essentially two reasons: # - To keep track of `function` # - To override SymPy caching behaviour From f3e1701651f2749486770c5c8c1348ab7369dff3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 16 Oct 2024 11:14:35 +0100 Subject: [PATCH 102/151] dsl: Edit PETScSolve to support modified stencils --- devito/petsc/solve.py | 71 +++++++++++++++++----- tests/test_petsc.py | 134 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 188 insertions(+), 17 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 17c03ae435..ef792c8a64 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -38,20 +38,20 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): time_mapper = generate_time_mapper(funcs) for eq in eqns: - b, F_target = separate_eqn(eq, target) + b, F_target, targets = separate_eqn(eq, target) b, F_target = b.subs(time_mapper), F_target.subs(time_mapper) # TODO: Current assumption is that problem is linear and user has not provided # a jacobian. Hence, we can use F_target to form the jac-vec product matvecs.append(Eq( arrays['y_matvec'], - F_target.subs({target: arrays['x_matvec']}), + F_target.subs(targets_to_arrays(arrays['x_matvec'], targets)), subdomain=eq.subdomain )) formfuncs.append(Eq( arrays['y_formfunc'], - F_target.subs({target: arrays['x_formfunc']}), + F_target.subs(targets_to_arrays(arrays['x_formfunc'], targets)), subdomain=eq.subdomain )) @@ -83,38 +83,77 @@ def separate_eqn(eqn, target): where F(target) = b. """ zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) - tmp = eval_time_derivatives(zeroed_eqn.lhs) - b, F_target = remove_target(tmp, target) - return -b, F_target + zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + target_funcs = set(generate_targets(zeroed_eqn, target)) + b, F_target = remove_target(zeroed_eqn, target_funcs) + return -b, F_target, target_funcs + + +def generate_targets(eq, target): + """ + Extract all the functions that share the same time index as the target, + but may have different spatial indices. + """ + funcs = retrieve_functions(eq) + if any(dim.is_Time for dim in target.dimensions): + time_idx = [ + i for i, d in zip(target.indices, target.dimensions) if d.is_Time + ] + targets = [ + func for func in funcs + if func.function is target.function and time_idx[0] + in func.indices + ] + else: + targets = [ + func for func in funcs + if func.function is target.function + ] + return targets + + +def targets_to_arrays(array, targets): + space_indices = [ + tuple( + i for i, d in zip(func.indices, func.dimensions) if d.is_Space + ) for func in targets + ] + array_targets = [ + array.subs( + {arr_idx: target_idx for arr_idx, target_idx in zip(array.indices, indices)} + ) + for indices in space_indices + ] + return {target: array for target, array in zip(targets, array_targets)} @singledispatch -def remove_target(expr, target): - return (0, expr) if expr == target else (expr, 0) +def remove_target(expr, targets): + return (0, expr) if expr in targets else (expr, 0) @remove_target.register(sympy.Add) -def _(expr, target): - if not expr.has(target): +def _(expr, targets): + if not any(expr.has(t) for t in targets): return (expr, 0) - args_b, args_F = zip(*(remove_target(a, target) for a in expr.args)) + args_b, args_F = zip(*(remove_target(a, targets) for a in expr.args)) return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) @remove_target.register(Mul) -def _(expr, target): - if not expr.has(target): +def _(expr, targets): + if not any(expr.has(t) for t in targets): return (expr, 0) - args_b, args_F = zip(*[remove_target(a, target) if a.has(target) + args_b, args_F = zip(*[remove_target(a, targets) if any(a.has(t) for t in targets) else (a, a) for a in expr.args]) return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) @remove_target.register(Derivative) -def _(expr, target): - return (0, expr) if expr.has(target) else (expr, 0) +def _(expr, targets): + return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) @singledispatch diff --git a/tests/test_petsc.py b/tests/test_petsc.py index f5616d4f36..8261f69e43 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -369,7 +369,139 @@ def test_separate_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F = separate_eqn(eval(eqn), eval(target)) + b, F, _ = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F + + +@skipif('petsc') +@pytest.mark.parametrize('eqn, target, expected', [ + ('Eq(f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2' + )), + ('Eq(g1, f1.laplace).evaluate', 'f1', + ( + '-g1(x, y)', + '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) ' + '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)' + )), + ('Eq(g1, f1.laplace).evaluate', 'g1', + ( + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2', + 'g1(x, y)' + )), + ('Eq(f1 + f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0' + '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1', + ( + '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)', + '-f1(x, y)/h_x + f1(x + h_x, y)/h_x' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1', + ( + '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)', + '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'g1', + ( + '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'f1', + ( + 'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1', + ( + '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) ' + '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y ' + '- f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1', + ( + 'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + ' + '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y' + )), + ('Eq(f2.laplace, g2).evaluate', 'g2', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)' + '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)' + )), + ('Eq(f2.laplace, g2).evaluate', 'f2', + ( + 'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)' + '/h_y**2 + f2(t, x, y + h_y)/h_y**2' + )), + ('Eq(f2.laplace, f2).evaluate', 'f2', + ( + '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)' + )), + ('Eq(g2*f2.laplace, f2).evaluate', 'g2', + ( + 'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward', + ( + 'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2' + ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2', + ( + '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + ' + 'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)', + '-f2(t, x, y)' + )), + ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/' + 'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + ' + 'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*' + 'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + ' + 'f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward', + ( + '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + ' + 'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)' + '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)' + )) +]) +def test_separate_eval_eqn(eqn, target, expected): + """ + Test the separate_eqn function on pre-evaluated equations. + This ensures that evaluated equations can be passed to PETScSolve, + allowing users to modify stencils for specific boundary conditions, + such as implementing free surface boundary conditions. + """ + grid = Grid((2, 2)) + + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F, _ = separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b From e825ab83fcc81b0b42ec11fb60b13e9ad99b9fd6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Oct 2024 11:07:42 +0100 Subject: [PATCH 103/151] dsl: Map time indices after forming eqn sets --- devito/petsc/solve.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index ef792c8a64..e42bb1c0f8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -34,12 +34,9 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): formrhs = [] eqns = as_tuple(eqns) - funcs = retrieve_functions(eqns) - time_mapper = generate_time_mapper(funcs) for eq in eqns: b, F_target, targets = separate_eqn(eq, target) - b, F_target = b.subs(time_mapper), F_target.subs(time_mapper) # TODO: Current assumption is that problem is linear and user has not provided # a jacobian. Hence, we can use F_target to form the jac-vec product @@ -61,6 +58,11 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): subdomain=eq.subdomain )) + funcs = retrieve_functions(eqns) + time_mapper = generate_time_mapper(funcs) + matvecs, formfuncs, formrhs = ( + [eq.subs(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs) + ) # Placeholder equation for inserting calls to the solver and generating # correct time loop etc inject_solve = InjectSolveEq(target, LinearSolveExpr( From ccb1901d04a54977f6571a6563953b8dbd59a80d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 20 Oct 2024 15:54:43 +0100 Subject: [PATCH 104/151] dsl: Simplify index extraction --- devito/petsc/iet/passes.py | 16 +++++----- devito/petsc/solve.py | 64 +++++++++++++++++++++----------------- 2 files changed, 43 insertions(+), 37 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 835ad102fb..7d7c156dcb 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -271,9 +271,9 @@ def assign_time_iters(iet, struct): mapper = {} for iter in time_iters: - common_dims = [dim for dim in iter.dimensions if dim in struct.fields] + common_dims = [d for d in iter.dimensions if d in struct.fields] common_dims = [ - DummyExpr(FieldFromComposite(dim, struct), dim) for dim in common_dims + DummyExpr(FieldFromComposite(d, struct), d) for d in common_dims ] iter_new = iter._rebuild(nodes=List(body=tuple(common_dims)+iter.nodes)) mapper.update({iter: iter_new}) @@ -282,15 +282,15 @@ def assign_time_iters(iet, struct): def retrieve_time_dims(iters): - time_iter = [i for i in iters if any(dim.is_Time for dim in i.dimensions)] + time_iter = [i for i in iters if any(d.is_Time for d in i.dimensions)] mapper = {} if not time_iter: return mapper - for dim in time_iter[0].dimensions: - if dim.is_Modulo: - mapper[dim.origin] = dim - elif dim.is_Time: - mapper[dim] = dim + for d in time_iter[0].dimensions: + if d.is_Modulo: + mapper[d.origin] = d + elif d.is_Time: + mapper[d] = d return mapper diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index e42bb1c0f8..c5df6b859f 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -4,7 +4,7 @@ from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative -from devito.types import Eq, Symbol, SteppingDimension +from devito.types import Eq, Symbol, SteppingDimension, TimeFunction from devito.types.equation import InjectSolveEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions @@ -87,73 +87,79 @@ def separate_eqn(eqn, target): zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) target_funcs = set(generate_targets(zeroed_eqn, target)) - b, F_target = remove_target(zeroed_eqn, target_funcs) + b, F_target = remove_targets(zeroed_eqn, target_funcs) return -b, F_target, target_funcs def generate_targets(eq, target): """ - Extract all the functions that share the same time index as the target, + Extract all the functions that share the same time index as the target but may have different spatial indices. """ funcs = retrieve_functions(eq) - if any(dim.is_Time for dim in target.dimensions): - time_idx = [ - i for i, d in zip(target.indices, target.dimensions) if d.is_Time - ] + if isinstance(target, TimeFunction): + time_idx = target.indices[target.time_dim] targets = [ - func for func in funcs - if func.function is target.function and time_idx[0] - in func.indices + f for f in funcs if f.function is target.function and time_idx + in f.indices ] else: - targets = [ - func for func in funcs - if func.function is target.function - ] + targets = [f for f in funcs if f.function is target.function] return targets def targets_to_arrays(array, targets): + """ + Map each target in `targets` to a corresponding array generated from `array`, + matching the spatial indices of the target. + + Example: + -------- + >>> array + vec_u(x, y) + + >>> targets + {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} + + >>> targets_to_arrays(array, targets) + {u(t + dt, x - h_x, y): vec_u(x - h_x, y), + u(t + dt, x + h_x, y): vec_u(x + h_x, y), + u(t + dt, x, y): vec_u(x, y)} + """ space_indices = [ - tuple( - i for i, d in zip(func.indices, func.dimensions) if d.is_Space - ) for func in targets + tuple(f.indices[d] for d in f.space_dimensions) for f in targets ] array_targets = [ - array.subs( - {arr_idx: target_idx for arr_idx, target_idx in zip(array.indices, indices)} - ) - for indices in space_indices + array.subs(dict(zip(array.indices, i))) for i in space_indices ] - return {target: array for target, array in zip(targets, array_targets)} + return dict(zip(targets, array_targets)) @singledispatch -def remove_target(expr, targets): +def remove_targets(expr, targets): return (0, expr) if expr in targets else (expr, 0) -@remove_target.register(sympy.Add) +@remove_targets.register(sympy.Add) def _(expr, targets): if not any(expr.has(t) for t in targets): return (expr, 0) - args_b, args_F = zip(*(remove_target(a, targets) for a in expr.args)) + args_b, args_F = zip(*(remove_targets(a, targets) for a in expr.args)) return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) -@remove_target.register(Mul) +@remove_targets.register(Mul) def _(expr, targets): if not any(expr.has(t) for t in targets): return (expr, 0) - args_b, args_F = zip(*[remove_target(a, targets) if any(a.has(t) for t in targets) + args_b, args_F = zip(*[remove_targets(a, targets) if any(a.has(t) for t in targets) else (a, a) for a in expr.args]) return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) -@remove_target.register(Derivative) +@remove_targets.register(Derivative) def _(expr, targets): return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) @@ -219,4 +225,4 @@ def generate_time_mapper(funcs): if d.is_Time }) tau_symbs = [Symbol('tau%d' % i) for i in range(len(time_indices))] - return {time: tau for time, tau in zip(time_indices, tau_symbs)} + return dict(zip(time_indices, tau_symbs)) From f0483f400921743f7ab4a26ed241aaa627f1d5eb Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 105/151] misc: Split Dockerfile.cpu into .cpu and .intel --- docker/Dockerfile.intel | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index e013fccc03..1fab7147d6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,6 +35,7 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu +<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -50,6 +51,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +======= +# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal +RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg +RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list + +RUN apt-get update -y && apt-get dist-upgrade -y && \ + apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ + intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ + libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -76,7 +87,11 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD # ICX OpenMP image +======= +# ICX image +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -93,6 +108,35 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD +======= +# ICX hpc image +############################################################## +FROM oneapi as icx-hpc + +# Install both icc and icx to avoid missing dependencies +RUN apt-get update -y && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic + +# Missig components +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager +RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ + chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ + rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# Devito config +ENV DEVITO_ARCH="icx" +ENV DEVITO_LANGUAGE="openmp" +# MPICC compiler for mpi4py +ENV MPICC=mpiicc +ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' + +############################################################## +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -107,6 +151,12 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl +<<<<<<< HEAD +======= +# NOTE: the name of this file ends with ".cpu" but this is a GPU image. +# It then feels a bit akward, so some restructuring might be needed + +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 57b99e69b427ffe8cec0292a8024d149c15de24d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 106/151] misc: Update file due to incorrect rebase --- docker/Dockerfile.intel | 50 ----------------------------------------- 1 file changed, 50 deletions(-) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index 1fab7147d6..e013fccc03 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,7 +35,6 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu -<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -51,16 +50,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev -======= -# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal -RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg -RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list - -RUN apt-get update -y && apt-get dist-upgrade -y && \ - apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ - intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ - libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -87,11 +76,7 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD # ICX OpenMP image -======= -# ICX image ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -108,35 +93,6 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD -======= -# ICX hpc image -############################################################## -FROM oneapi as icx-hpc - -# Install both icc and icx to avoid missing dependencies -RUN apt-get update -y && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic - -# Missig components -# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager -RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ - chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ - rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -# Devito config -ENV DEVITO_ARCH="icx" -ENV DEVITO_LANGUAGE="openmp" -# MPICC compiler for mpi4py -ENV MPICC=mpiicc -ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' - -############################################################## ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -151,12 +107,6 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl -<<<<<<< HEAD -======= -# NOTE: the name of this file ends with ".cpu" but this is a GPU image. -# It then feels a bit akward, so some restructuring might be needed - ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From ca20534760148b3a4ccf64509df05893c0252ac9 Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 107/151] misc: Split Dockerfile.cpu into .cpu and .intel --- docker/Dockerfile.intel | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index e013fccc03..1fab7147d6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,6 +35,7 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu +<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -50,6 +51,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +======= +# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal +RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg +RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list + +RUN apt-get update -y && apt-get dist-upgrade -y && \ + apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ + intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ + libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -76,7 +87,11 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD # ICX OpenMP image +======= +# ICX image +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -93,6 +108,35 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD +======= +# ICX hpc image +############################################################## +FROM oneapi as icx-hpc + +# Install both icc and icx to avoid missing dependencies +RUN apt-get update -y && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic + +# Missig components +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager +RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ + chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ + rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# Devito config +ENV DEVITO_ARCH="icx" +ENV DEVITO_LANGUAGE="openmp" +# MPICC compiler for mpi4py +ENV MPICC=mpiicc +ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' + +############################################################## +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -107,6 +151,12 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl +<<<<<<< HEAD +======= +# NOTE: the name of this file ends with ".cpu" but this is a GPU image. +# It then feels a bit akward, so some restructuring might be needed + +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 15ed63e3ac71020fe6641e3dd2665ebd2da532ff Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 108/151] misc: Update file due to incorrect rebase --- devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 2 +- devito/ir/iet/algorithms.py | 1 + devito/ir/stree/algorithms.py | 1 + devito/operator/operator.py | 1 + devito/passes/iet/definitions.py | 4 +- devito/petsc/types.py | 384 +++++++++++++++++++++++++++++++ devito/symbolics/printer.py | 3 + docker/Dockerfile.intel | 50 ---- 9 files changed, 395 insertions(+), 53 deletions(-) create mode 100644 devito/petsc/types.py diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index f0b0d718ea..5aa71a0529 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -376,7 +376,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 67ed9269a4..651aa7d1d3 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -239,7 +239,7 @@ def __new__(cls, *args, **kwargs): expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + return expr @property diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 9d2db185db..af48dc73c7 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -24,6 +24,7 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: + # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index a85b93460d..c738db645a 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,6 +163,7 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: + # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index fe4b5ab48c..a0c0981d68 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -428,6 +428,7 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters + # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 59de8f95c7..98fb8cabfe 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -456,7 +456,9 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [i for i in bases if i.name != i.function._C_name] + bases = [ + i for i in bases + if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/petsc/types.py b/devito/petsc/types.py new file mode 100644 index 0000000000..62c69f666a --- /dev/null +++ b/devito/petsc/types.py @@ -0,0 +1,384 @@ +import sympy +from functools import cached_property +import numpy as np + +from devito.tools import CustomDtype +from devito.types import LocalObject, Eq, CompositeObject +from devito.types.utils import DimensionTuple +from devito.types.array import ArrayBasic +from devito.finite_differences import Differentiable +from devito.types.basic import AbstractFunction, Symbol +from devito.finite_differences.tools import fd_weights_registry +from devito.tools import Reconstructable, dtype_to_ctype +from devito.symbolics import FieldFromComposite +from devito.types.basic import IndexedBase + + +class DM(LocalObject): + """ + PETSc Data Management object (DM). + """ + dtype = CustomDtype('DM') + + def __init__(self, *args, stencil_width=None, **kwargs): + super().__init__(*args, **kwargs) + self._stencil_width = stencil_width + + @property + def stencil_width(self): + return self._stencil_width + + +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 SNES(LocalObject): + """ + PETSc SNES : Non-Linear Systems Solvers. + """ + dtype = CustomDtype('SNES') + + +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 DMDALocalInfo(LocalObject): + """ + PETSc object - C struct containing information + about the local grid. + """ + dtype = CustomDtype('DMDALocalInfo') + + +class PetscErrorCode(LocalObject): + """ + PETSc datatype used to return PETSc error codes. + https://petsc.org/release/manualpages/Sys/PetscErrorCode/ + """ + dtype = CustomDtype('PetscErrorCode') + + +class PETScArray(ArrayBasic, Differentiable): + """ + 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. + Differentiable enables compatability with standard Function objects, + allowing for the use of the `subs` method. + TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) + and then PETScArray(ArrayBasic). + """ + + _data_alignment = False + + # Default method for the finite difference approximation weights computation. + _default_fd = 'taylor' + + __rkwargs__ = (AbstractFunction.__rkwargs__ + + ('dimensions', 'shape', 'liveness', 'coefficients')) + + def __init_finalize__(self, *args, **kwargs): + + super().__init_finalize__(*args, **kwargs) + + # Symbolic (finite difference) coefficients + self._coefficients = kwargs.get('coefficients', self._default_fd) + if self._coefficients not in fd_weights_registry: + raise ValueError("coefficients must be one of %s" + " not %s" % (str(fd_weights_registry), self._coefficients)) + self._shape = kwargs.get('shape') + + @classmethod + def __dtype_setup__(cls, **kwargs): + return kwargs.get('dtype', np.float32) + + @property + def coefficients(self): + """Form of the coefficients of the function.""" + return self._coefficients + + @property + def shape(self): + return self._shape + + @cached_property + def _shape_with_inhalo(self): + """ + Shape of the domain+inhalo region. The inhalo region comprises the + outhalo as well as any additional "ghost" layers for MPI halo + exchanges. Data in the inhalo region are exchanged when running + Operators to maintain consistent values as in sequential runs. + + Notes + ----- + Typically, this property won't be used in user code, but it may come + in handy for testing or debugging + """ + return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) + + @cached_property + def shape_allocated(self): + """ + Shape of the allocated data of the Function type object from which + this PETScArray was derived. It includes the domain and inhalo regions, + as well as any additional padding surrounding the halo. + + Notes + ----- + In an MPI context, this is the *local* with_halo region shape. + """ + return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, + self._padding)], + getters=self.dimensions) + + @cached_property + def _C_ctype(self): + return CustomDtype(self.petsc_type, modifier='*') + + @property + def petsc_type(self): + return dtype_to_petsctype(self._dtype) + + @property + def dtype(self): + return CustomDtype(self.petsc_type) + + @property + def symbolic_shape(self): + field_from_composites = [ + FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] + # Reverse it since DMDA is setup backwards to Devito dimensions. + return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + + @cached_property + def dmda_info(self): + # To access the local grid info via the DM in + # PETSc you use DMDAGetLocalInfo(da, &info) + # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz + return DMDALocalInfo(name='info', liveness='eager') + + @cached_property + def indexed(self): + return PETScIndexedData(self.name, shape=self._shape, function=self.function) + + +class PETScIndexedData(IndexedBase): + pass + + +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] + + +class MatVecEq(Eq): + """ + Represents the mathematical expression of applying a linear + operator to a vector. This is a key component + for running matrix-free solvers. + """ + pass + + +class RHSEq(Eq): + """ + Represents the mathematical expression of building the + rhs of a linear system. + """ + pass + + +class MockEq(Eq): + """ + Represents a mock/placeholder equation to ensure distinct iteration loops. + + For example, the mat-vec action iteration loop is to be isolated from the + expression loop used to build the RHS of the linear system. This separation + facilitates the utilisation of the mat-vec iteration loop in callback functions + created at the IET level. + """ + pass + + +def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): + # TODO: Add check for time dimensions and utilise implicit dimensions. + + is_time_dep = any(dim.is_Time for dim in target.dimensions) + # TODO: Current assumption is rhs is part of pde that remains + # constant at each timestep. Need to insert function to extract this from eq. + y_matvec, x_matvec, b_tmp = [ + PETScArray(name=f'{prefix}_{target.name}', + dtype=target.dtype, + dimensions=target.space_dimensions, + shape=target.grid.shape, + liveness='eager', + halo=target.halo[1:] if is_time_dep else target.halo) + for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] + + # TODO: Extend to rearrange equation for implicit time stepping. + matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), + target=target, solver_parameters=solver_parameters), + subdomain=eq.subdomain) + + # Part of pde that remains constant at each timestep + rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, + solver_parameters=solver_parameters), subdomain=eq.subdomain) + + if not bcs: + return [matvecaction, rhs] + + bcs_for_matvec = [] + for bc in bcs: + # TODO: Insert code to distiguish between essential and natural + # boundary conditions since these are treated differently within + # the solver + # NOTE: May eventually remove the essential bcs from the solve + # (and move to rhs) but for now, they are included since this + # is not trivial to implement when using DMDA + # NOTE: Below is temporary -> Just using this as a palceholder for + # the actual BC implementation for the matvec callback + new_rhs = bc.rhs.subs(target, x_matvec) + bc_rhs = LinearSolveExpr( + new_rhs, target=target, solver_parameters=solver_parameters + ) + bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) + + return [matvecaction] + bcs_for_matvec + [rhs] + + +class LinearSolveExpr(sympy.Function, Reconstructable): + + __rargs__ = ('expr',) + __rkwargs__ = ('target', 'solver_parameters') + + defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'jacobi', + 'ksp_rtol': 1e-7, # Relative tolerance + 'ksp_atol': 1e-50, # Absolute tolerance + 'ksp_divtol': 1e4, # Divergence tolerance + 'ksp_max_it': 10000 # Maximum iterations + } + + def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): + + if solver_parameters is None: + solver_parameters = cls.defaults + else: + for key, val in cls.defaults.items(): + solver_parameters[key] = solver_parameters.get(key, val) + + obj = super().__new__(cls, expr) + obj._expr = expr + obj._target = target + obj._solver_parameters = solver_parameters + return obj + + def __repr__(self): + return "%s(%s)" % (self.__class__.__name__, self.expr) + + __str__ = __repr__ + + def _sympystr(self, printer): + return str(self) + + def __hash__(self): + return hash(self.target) + + @property + def expr(self): + return self._expr + + @property + def target(self): + return self._target + + @property + def solver_parameters(self): + return self._solver_parameters + + func = Reconstructable._rebuild + + +def petsc_lift(clusters): + """ + Lift the iteration space associated with each PETSc equation. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions. + """ + processed = [] + for c in clusters: + + ispace = c.ispace + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + + processed.append(c.rebuild(ispace=ispace)) + + return processed + + +class PETScStruct(CompositeObject): + + __rargs__ = ('name', 'usr_ctx') + + def __init__(self, name, usr_ctx): + pfields = [(i._C_name, dtype_to_ctype(i.dtype)) + for i in usr_ctx if isinstance(i, Symbol)] + self._usr_ctx = usr_ctx + super().__init__(name, 'MatContext', pfields) + + @property + def usr_ctx(self): + return self._usr_ctx + + def _arg_values(self, **kwargs): + values = super()._arg_values(**kwargs) + for i in self.fields: + setattr(values[self.name]._obj, i, kwargs['args'][i]) + return values diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index 2c366a389e..e7b395d436 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,6 +242,9 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) + def _print_PETScRHS(self, expr): + return "%s" % self._print(expr._expr) + _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index 1fab7147d6..e013fccc03 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,7 +35,6 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu -<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -51,16 +50,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev -======= -# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal -RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg -RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list - -RUN apt-get update -y && apt-get dist-upgrade -y && \ - apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ - intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ - libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -87,11 +76,7 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD # ICX OpenMP image -======= -# ICX image ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -108,35 +93,6 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD -======= -# ICX hpc image -############################################################## -FROM oneapi as icx-hpc - -# Install both icc and icx to avoid missing dependencies -RUN apt-get update -y && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic - -# Missig components -# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager -RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ - chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ - rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -# Devito config -ENV DEVITO_ARCH="icx" -ENV DEVITO_LANGUAGE="openmp" -# MPICC compiler for mpi4py -ENV MPICC=mpiicc -ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' - -############################################################## ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -151,12 +107,6 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl -<<<<<<< HEAD -======= -# NOTE: the name of this file ends with ".cpu" but this is a GPU image. -# It then feels a bit akward, so some restructuring might be needed - ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From bd4fc4bb832e87604b6ed1f039f10334768ad906 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 10 Jul 2024 11:31:30 +0100 Subject: [PATCH 109/151] dsl: Dispatch to sympy.Add not both Add and EvalDerivative --- .github/workflows/examples-mpi.yml | 1 + .github/workflows/examples.yml | 1 + .github/workflows/flake8.yml | 1 + .github/workflows/pytest-core-mpi.yml | 1 + .github/workflows/pytest-core-nompi.yml | 1 + .github/workflows/tutorials.yml | 1 + devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 6 +- devito/ir/stree/algorithms.py | 1 - devito/operator/operator.py | 8 +- devito/passes/clusters/petsc.py | 25 ++ devito/passes/iet/__init__.py | 1 + devito/passes/iet/petsc.py | 298 +++++++++++++++++++ devito/petsc/iet.py | 372 ++++++++++++++++++++++++ devito/petsc/types.py | 303 +++++++++++-------- docker/Dockerfile.petsc | 49 ++++ tests/test_petsc.py | 4 + 17 files changed, 939 insertions(+), 136 deletions(-) create mode 100644 devito/passes/clusters/petsc.py create mode 100644 devito/passes/iet/petsc.py create mode 100644 devito/petsc/iet.py create mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 3bbebabb1c..9df9251991 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,6 +17,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 026331c881..4490f47547 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 95a94f6e79..b61e777454 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 4f2853e46e..cc4b8ab683 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 54b5b05b1f..802aa3afb3 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index d27af78fcf..af7f17d255 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 5aa71a0529..f0b0d718ea 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -376,7 +376,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 651aa7d1d3..0d6c7b0ff7 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -227,19 +227,19 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - + # from IPython import embed; embed() # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - + # from IPython import embed; embed() expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + return expr @property diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index c738db645a..a85b93460d 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,7 +163,6 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: - # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index a0c0981d68..6a4f66f28d 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -354,7 +354,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - + # from IPython import embed; embed() return processed # Compilation -- Cluster level @@ -498,6 +498,9 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Target-independent optimizations minimize_symbols(graph) + # If necessary, sort frees into a specific order + sort_frees(graph) + return graph.root, graph # Read-only properties exposed to the outside world @@ -524,7 +527,8 @@ def dimensions(self): @cached_property def input(self): - return tuple(i for i in self.parameters if i.is_Input) + struct_params = derive_struct_inputs(self.parameters) + return tuple(i for i in self.parameters+struct_params if i.is_Input) @cached_property def temporaries(self): diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py new file mode 100644 index 0000000000..d784d99c0d --- /dev/null +++ b/devito/passes/clusters/petsc.py @@ -0,0 +1,25 @@ +from devito.tools import timed_pass +from devito.petsc import LinearSolveExpr + +__all__ = ['petsc_lift'] + + +@timed_pass() +def petsc_lift(clusters): + """ + Lift the iteration space surrounding each PETSc equation to create + distinct iteration loops. This simplifys the movement of the loops + into specific callback functions generated at the IET level. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions, not the ones that stay inside the main kernel. + """ + processed = [] + for c in clusters: + + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index c09db00c9b..cf6a35de90 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,3 +8,4 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa +from .petsc import * # noqa diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py new file mode 100644 index 0000000000..8c8aeafa1c --- /dev/null +++ b/devito/passes/iet/petsc.py @@ -0,0 +1,298 @@ +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, MatVecAction, + Transformer, FindSymbols, LinearSolverExpression, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition) +from devito.types import (PetscMPIInt, PETScStruct, DMDALocalInfo, DM, Mat, + Vec, KSP, PC, SNES, PetscErrorCode) +from devito.symbolics import Byref, Macro, FieldFromPointer +import cgen as c + +__all__ = ['lower_petsc'] + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't + # appear in ccode + + # Check if PETScSolve was used and count occurrences. Each PETScSolve + # will have a unique MatVecAction. + is_petsc = FindNodes(MatVecAction).visit(iet) + + if is_petsc: + + # Collect all solution fields we're solving for + targets = [i.expr.rhs.target for i in is_petsc] + + # Initalize PETSc + init = init_petsc(**kwargs) + + # Create context data struct. + struct = build_struct(iet) + + objs = build_core_objects(targets[-1], struct, **kwargs) + + # Create core PETSc calls required by general linear solves, + # which only need to be generated once e.g create DMDA. + core = core_petsc(targets[-1], objs, **kwargs) + + matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) + + main_mapper = {} + + setup = [] + efuncs = [] + for target in unique_targets: + + solver_objs = build_solver_objs(target) + + matvec_callback_body_iters = [] + + solver_setup = False + + for iter, (matvec,) in matvec_mapper.items(): + + if matvec.expr.rhs.target == target: + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec) + setup.extend(solver) + solver_setup = True + + # Only need to generate solver setup once per target + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + solver_setup = True + + matvec_body = matvec_body_list._rebuild(body=[ + matvec_body_list.body, iter[0]]) + matvec_body_list = matvec_body_list._rebuild(body=matvec_body) + + main_mapper.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, matvec_callback_body_iters, solver_objs, objs, struct) + + setup.append(matvec_op) + setup.append(c.Line()) + + efuncs.append(matvec_callback) + + # Remove the LinSolveExpr from iet and efuncs that were used to carry + # metadata e.g solver_parameters + main_mapper.update(rebuild_expr_mapper(iet)) + efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} + + iet = Transformer(main_mapper).visit(iet) + efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] + + # Replace symbols appearing in each efunc with a pointer to the PETScStruct + efuncs = transform_efuncs(efuncs, struct) + + body = iet.body._rebuild(body=(tuple(init_setup) + iet.body.body)) + iet = iet._rebuild(body=body) + + return iet, {} + + +def init_petsc(**kwargs): + + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve. + # NOTE: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options. + initialize = Call('PetscCall', [Call('PetscInitialize', + arguments=['NULL', 'NULL', + 'NULL', 'NULL'])]) + + return tuple([initialize]) + + +def build_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct. + usr_ctx = [] + + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def core_petsc(target, objs, **kwargs): + # Assumption: all targets are generated from the same Grid, + # so we can use any target. + + # MPI + call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', + arguments=[objs['comm'], + Byref(objs['size'])])]) + # Create DMDA + dmda = create_dmda(target, objs) + dm_setup = Call('PetscCall', [Call('DMSetUp', arguments=[objs['da']])]) + dm_app_ctx = Call('PetscCall', [Call('DMSetApplicationContext', + arguments=[objs['da'], objs['struct']])]) + dm_mat_type = Call('PetscCall', [Call('DMSetMatType', + arguments=[objs['da'], 'MATSHELL'])]) + dm_local_info = Call('PetscCall', [Call('DMDAGetLocalInfo', + arguments=[objs['da'], Byref(objs['info'])])]) + + return tuple([call_mpi, dmda, dm_setup, dm_app_ctx, dm_mat_type, dm_local_info]) + + +def build_core_objects(target, struct, **kwargs): + + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return {'da': DM(name='da'), + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'struct': struct} + + +def create_dmda(target, objs): + + args = [objs['comm']] + + args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] + + # stencil type + args += ['DMDA_STENCIL_BOX'] + + # global dimensions + args += list(target.shape_global)[::-1] + + # no.of processors in each dimension + args += list(target.grid.distributor.topology)[::-1] + + args += [1, target.space_order] + + args += ['NULL' for _ in range(len(target.space_dimensions))] + + args += [Byref(objs['da'])] + + dmda = Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args) + + return dmda + + +def build_solver_objs(target): + + return {'Jac': Mat(name='J_'+str(target.name)), + 'x_global': Vec(name='x_global_'+str(target.name)), + 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), + 'b_global': Vec(name='b_global_'+str(target.name)), + 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), + 'ksp': KSP(name='ksp_'+str(target.name)), + 'pc': PC(name='pc_'+str(target.name)), + 'snes': SNES(name='snes_'+str(target.name)), + 'x': Vec(name='x_'+str(target.name)), + 'y': Vec(name='y_'+str(target.name))} + + +def generate_solver_calls(solver_objs, objs, matvec): + + snes_create = Call('PetscCall', [Call('SNESCreate', arguments=[ + objs['comm'], Byref(solver_objs['snes'])])]) + + snes_set_dm = Call('PetscCall', [Call('SNESSetDM', arguments=[ + solver_objs['snes'], objs['da']])]) + + create_matrix = Call('PetscCall', [Call('DMCreateMatrix', arguments=[ + objs['da'], Byref(solver_objs['Jac'])])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = Call('PetscCall', [Call('SNESSetType', arguments=[ + solver_objs['snes'], 'SNESKSPONLY'])]) + + global_x = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['x_global'])])]) + + local_x = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['x_local'])])]) + + global_b = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['b_global'])])]) + + local_b = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['b_local'])])]) + + snes_get_ksp = Call('PetscCall', [Call('SNESGetKSP', arguments=[ + solver_objs['snes'], Byref(solver_objs['ksp'])])]) + + return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, + global_x, local_x, global_b, local_b, snes_get_ksp]) + + +def create_matvec_callback(target, matvec_callback_body_iters, + solver_objs, objs, struct): + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + defn_struct = Definition(struct) + + get_context = Call('PetscCall', [Call('MatShellGetContext', + arguments=[solver_objs['Jac'], + Byref(struct)])]) + + body = List(body=[defn_struct, + get_context, + matvec_callback_body_iters]) + + matvec_callback = Callable('MyMatShellMult_'+str(target.name), + matvec_body, + retval=objs['err'], + parameters=(solver_objs['Jac'], + solver_objs['x'], + solver_objs['y'])) + + matvec_operation = Call('PetscCall', [ + Call('MatShellSetOperation', arguments=[solver_objs['Jac'], + 'MATOP_MULT', + Callback(matvec_callback.name, + Void, Void)])]) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(iet): + + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for + expr in FindNodes(LinearSolverExpression).visit(iet)} + + +def transform_efuncs(efuncs, struct): + + efuncs_new = [] + for efunc in efuncs: + new_body = efunc.body + for i in struct.usr_ctx: + new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) + efunc_with_new_body = efunc._rebuild(body=new_body) + efuncs_new.append(efunc_with_new_body) + + return efuncs_new + + +Null = Macro('NULL') +Void = Macro('void') + +petsc_call = String('PetscCall') +petsc_call_mpi = String('PetscCallMPI') +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') + +linear_solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} diff --git a/devito/petsc/iet.py b/devito/petsc/iet.py new file mode 100644 index 0000000000..659ac9f94d --- /dev/null +++ b/devito/petsc/iet.py @@ -0,0 +1,372 @@ +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, + Transformer, FindSymbols, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition, BlankLine, PointerCast) +<<<<<<< HEAD +from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, +======= +from devito.petsc import (PetscMPIInt, PETScStruct, DM, Mat, +>>>>>>> b22824875 (fix circular imports) + Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) +from devito.symbolics import Byref, Macro, FieldFromPointer +import cgen as c +from devito.petsc.nodes import MatVecAction, RHSLinearSystem, LinearSolverExpression + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # Check if PETScSolve was used. + petsc_nodes = FindNodes(MatVecAction).visit(iet) + + if not petsc_nodes: + return iet, {} + + else: + # Collect all petsc solution fields + unique_targets = list(set([i.expr.rhs.target for i in petsc_nodes])) + + # Initalize PETSc + init = init_petsc(**kwargs) + + # Create context data struct + struct = build_struct(iet) + + objs = build_core_objects(unique_targets[-1], **kwargs) + + # Create core PETSc calls (not specific to each PETScSolve) + core = core_petsc(unique_targets[-1], struct, objs, **kwargs) + + matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) + + main_mapper = {} + + setup = [] + efuncs = [] + + for target in unique_targets: + + solver_objs = build_solver_objs(target) + + matvec_body_list = List() + + solver_setup = False + + for iter, (matvec,) in matvec_mapper.items(): + + # Skip the MatVecAction if it is not associated with the target + # There will be more than one MatVecAction associated with the target + # e.g interior matvec + BC matvecs + if matvec.expr.rhs.target != target: + continue + + # Only need to generate solver setup once per target + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + solver_setup = True + + matvec_body = matvec_body_list._rebuild(body=[ + matvec_body_list.body, iter[0]]) + matvec_body_list = matvec_body_list._rebuild(body=matvec_body) + + main_mapper.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, matvec_body_list, solver_objs, objs, + struct) + + setup.append(matvec_op) + setup.append(BlankLine) + efuncs.append(matvec_callback) + + # Remove the LinSolveExpr from iet and efuncs that were used to carry + # metadata e.g solver_parameters + main_mapper.update(rebuild_expr_mapper(iet)) + efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} + + iet = Transformer(main_mapper).visit(iet) + efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] + + # Replace symbols appearing in each efunc with a pointer to the PETScStruct + efuncs = transform_efuncs(efuncs, struct) + + body = iet.body._rebuild(init=init, body=core + tuple(setup) + iet.body.body) + iet = iet._rebuild(body=body) + + return iet, {'efuncs': efuncs} + + +def init_petsc(**kwargs): + + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve + # TODO: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options + initialize = Call(petsc_call, [ + Call('PetscInitialize', arguments=[Null, Null, Null, Null])]) + + return tuple([petsc_func_begin_user, initialize]) + + +def build_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct + usr_ctx = [] + + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def core_petsc(target, struct, objs, **kwargs): + + # MPI + call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', + arguments=[objs['comm'], + Byref(objs['size'])])]) + # Create DMDA + dmda = create_dmda(target, objs) + dm_setup = Call(petsc_call, [ + Call('DMSetUp', arguments=[objs['da']])]) + dm_app_ctx = Call(petsc_call, [ + Call('DMSetApplicationContext', arguments=[objs['da'], struct])]) + dm_mat_type = Call(petsc_call, [ + Call('DMSetMatType', arguments=[objs['da'], 'MATSHELL'])]) + + return tuple([petsc_func_begin_user, call_mpi, dmda, dm_setup, + dm_app_ctx, dm_mat_type, BlankLine]) + + +def build_core_objects(target, **kwargs): + + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return {'da': DM(name='da', liveness='eager'), + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'err': PetscErrorCode(name='err')} + + +def create_dmda(target, objs): + + args = [objs['comm']] + + args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] + + # Stencil type + if len(target.space_dimensions) > 1: + args += ['DMDA_STENCIL_BOX'] + + # Global dimensions + args += list(target.shape_global)[::-1] + + # No.of processors in each dimension + if len(target.space_dimensions) > 1: + args += list(target.grid.distributor.topology)[::-1] + + args += [1, target.space_order] + + args += [Null for _ in range(len(target.space_dimensions))] + + args += [Byref(objs['da'])] + + dmda = Call(petsc_call, [ + Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args)]) + + return dmda + + +def build_solver_objs(target): + + return {'Jac': Mat(name='J_'+str(target.name)), + 'x_global': Vec(name='x_global_'+str(target.name)), + 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), + 'b_global': Vec(name='b_global_'+str(target.name)), + 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), + 'ksp': KSP(name='ksp_'+str(target.name)), + 'pc': PC(name='pc_'+str(target.name)), + 'snes': SNES(name='snes_'+str(target.name)), + 'X_global': Vec(name='X_global_'+str(target.name)), + 'Y_global': Vec(name='Y_global_'+str(target.name)), + 'X_local': Vec(name='X_local_'+str(target.name), liveness='eager'), + 'Y_local': Vec(name='Y_local_'+str(target.name), liveness='eager') + } + + +def generate_solver_calls(solver_objs, objs, matvec, target): + + solver_params = matvec.expr.rhs.solver_parameters + + snes_create = Call(petsc_call, [Call('SNESCreate', arguments=[ + objs['comm'], Byref(solver_objs['snes'])])]) + + snes_set_dm = Call(petsc_call, [Call('SNESSetDM', arguments=[ + solver_objs['snes'], objs['da']])]) + + create_matrix = Call(petsc_call, [Call('DMCreateMatrix', arguments=[ + objs['da'], Byref(solver_objs['Jac'])])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = Call(petsc_call, [Call('SNESSetType', arguments=[ + solver_objs['snes'], 'SNESKSPONLY'])]) + + global_x = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['x_global'])])]) + + local_x = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['x_local'])])]) + + global_b = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['b_global'])])]) + + local_b = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['b_local'])])]) + + snes_get_ksp = Call(petsc_call, [Call('SNESGetKSP', arguments=[ + solver_objs['snes'], Byref(solver_objs['ksp'])])]) + + vec_replace_array = Call(petsc_call, [ + Call('VecReplaceArray', arguments=[solver_objs[ + 'x_local'], FieldFromPointer(target._C_field_data, target._C_symbol)])]) + + ksp_set_tols = Call(petsc_call, [Call('KSPSetTolerances', arguments=[ + solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], + solver_params['ksp_divtol'], solver_params['ksp_max_it']])]) + + ksp_set_type = Call(petsc_call, [Call('KSPSetType', arguments=[ + solver_objs['ksp'], linear_solver_mapper[solver_params['ksp_type']]])]) + + ksp_get_pc = Call(petsc_call, [Call('KSPGetPC', arguments=[ + solver_objs['ksp'], Byref(solver_objs['pc'])])]) + + pc_set_type = Call(petsc_call, [Call('PCSetType', arguments=[ + solver_objs['pc'], linear_solver_mapper[solver_params['pc_type']]])]) + + ksp_set_from_ops = Call(petsc_call, [Call('KSPSetFromOptions', arguments=[ + solver_objs['ksp']])]) + + return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, + global_x, local_x, global_b, local_b, snes_get_ksp, + vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, + pc_set_type, ksp_set_from_ops]) + + +def create_matvec_callback(target, body, solver_objs, objs, struct): + + # There will be 2 PETScArrays within the body + petsc_arrays = [i.function for i in FindSymbols('indexedbases').visit(body) + if isinstance(i.function, PETScArray)] + + # There will only be one PETScArray that is written to within this body and + # one PETScArray which corresponds to the 'seed' vector + petsc_arr_write, = FindSymbols('writes').visit(body) + petsc_arr_seed, = [i for i in petsc_arrays if i.function != petsc_arr_write.function] + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + defn_struct = Definition(struct) + + mat_get_dm = Call(petsc_call, [Call('MatGetDM', arguments=[ + solver_objs['Jac'], Byref(objs['da'])])]) + + dm_get_app_context = Call(petsc_call, [Call('DMGetApplicationContext', arguments=[ + objs['da'], Byref(struct._C_symbol)])]) + + dm_get_local_xvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ + objs['da'], Byref(solver_objs['X_local'])])]) + + global_to_local_begin = Call(petsc_call, [Call('DMGlobalToLocalBegin', arguments=[ + objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) + + global_to_local_end = Call(petsc_call, [Call('DMGlobalToLocalEnd', arguments=[ + objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) + + dm_get_local_yvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ + objs['da'], Byref(solver_objs['Y_local'])])]) + + vec_get_array_y = Call(petsc_call, [Call('VecGetArray', arguments=[ + solver_objs['Y_local'], Byref(petsc_arr_write.function)])]) + + vec_get_array_x = Call(petsc_call, [Call('VecGetArray', arguments=[ + solver_objs['X_local'], Byref(petsc_arrays[0])])]) + + dm_get_local_info = Call(petsc_call, [Call('DMDAGetLocalInfo', arguments=[ + objs['da'], Byref(petsc_arrays[0].function.dmda_info)])]) + + casts = [PointerCast(i.function) for i in petsc_arrays] + + vec_restore_array_y = Call(petsc_call, [Call('VecRestoreArray', arguments=[ + solver_objs['Y_local'], Byref(petsc_arr_write[0])])]) + + vec_restore_array_x = Call(petsc_call, [Call('VecRestoreArray', arguments=[ + solver_objs['X_local'], Byref(petsc_arr_seed)])]) + + dm_local_to_global_begin = Call(petsc_call, [Call('DMLocalToGlobalBegin', arguments=[ + objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) + + dm_local_to_global_end = Call(petsc_call, [Call('DMLocalToGlobalEnd', arguments=[ + objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) + + func_return = Call('PetscFunctionReturn', arguments=[0]) + + matvec_body = List(body=[ + petsc_func_begin_user, defn_struct, mat_get_dm, dm_get_app_context, + dm_get_local_xvec, global_to_local_begin, global_to_local_end, + dm_get_local_yvec, vec_get_array_y, vec_get_array_x, dm_get_local_info, + casts, BlankLine, body, vec_restore_array_y, vec_restore_array_x, + dm_local_to_global_begin, dm_local_to_global_end, func_return]) + + matvec_callback = Callable( + 'MyMatShellMult_'+str(target.name), matvec_body, retval=objs['err'], + parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'])) + + matvec_operation = Call(petsc_call, [ + Call('MatShellSetOperation', arguments=[ + solver_objs['Jac'], 'MATOP_MULT', Callback(matvec_callback.name, + Void, Void)])]) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(callable): + + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for + expr in FindNodes(LinearSolverExpression).visit(callable)} + + +def transform_efuncs(efuncs, struct): + + efuncs_new = [] + for efunc in efuncs: + new_body = efunc.body + for i in struct.usr_ctx: + new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) + efunc_with_new_body = efunc._rebuild(body=new_body) + efuncs_new.append(efunc_with_new_body) + + return efuncs_new + + +Null = Macro('NULL') +Void = 'void' + +petsc_call = 'PetscCall' +petsc_call_mpi = 'PetscCallMPI' +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') + +linear_solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} diff --git a/devito/petsc/types.py b/devito/petsc/types.py index 62c69f666a..72a0faa1e8 100644 --- a/devito/petsc/types.py +++ b/devito/petsc/types.py @@ -1,17 +1,19 @@ import sympy -from functools import cached_property -import numpy as np +<<<<<<< HEAD +from devito.tools import Reconstructable, sympy_mutex +======= from devito.tools import CustomDtype from devito.types import LocalObject, Eq, CompositeObject from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction, Symbol +from devito.types.basic import AbstractFunction, AbstractSymbol from devito.finite_differences.tools import fd_weights_registry -from devito.tools import Reconstructable, dtype_to_ctype -from devito.symbolics import FieldFromComposite -from devito.types.basic import IndexedBase +from devito.tools import dtype_to_ctype, Reconstructable, sympy_mutex +from devito.symbolics import FieldFromComposite, Byref + +from devito.petsc.utils import petsc_call class DM(LocalObject): @@ -28,6 +30,15 @@ def __init__(self, *args, stencil_width=None, **kwargs): def stencil_width(self): return self._stencil_width + @property + def info(self): + return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') + + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('DMDestroy', [Byref(self.function)]) + class Mat(LocalObject): """ @@ -35,13 +46,30 @@ class Mat(LocalObject): """ dtype = CustomDtype('Mat') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('MatDestroy', [Byref(self.function)]) + + +class LocalVec(LocalObject): + """ + PETSc Vector object (Vec). + """ + dtype = CustomDtype('Vec') + -class Vec(LocalObject): +class GlobalVec(LocalObject): """ PETSc Vector object (Vec). """ dtype = CustomDtype('Vec') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('VecDestroy', [Byref(self.function)]) + class PetscMPIInt(LocalObject): """ @@ -65,6 +93,11 @@ class SNES(LocalObject): """ dtype = CustomDtype('SNES') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('SNESDestroy', [Byref(self.function)]) + class PC(LocalObject): """ @@ -97,6 +130,10 @@ class PetscErrorCode(LocalObject): dtype = CustomDtype('PetscErrorCode') +class DummyArg(LocalObject): + dtype = CustomDtype('void', modifier='*') + + class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent @@ -114,7 +151,8 @@ class PETScArray(ArrayBasic, Differentiable): _default_fd = 'taylor' __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients')) + ('dimensions', 'shape', 'liveness', 'coefficients', + 'space_order')) def __init_finalize__(self, *args, **kwargs): @@ -126,6 +164,7 @@ def __init_finalize__(self, *args, **kwargs): raise ValueError("coefficients must be one of %s" " not %s" % (str(fd_weights_registry), self._coefficients)) self._shape = kwargs.get('shape') + self._space_order = kwargs.get('space_order', 1) @classmethod def __dtype_setup__(cls, **kwargs): @@ -140,6 +179,10 @@ def coefficients(self): def shape(self): return self._shape + @property + def space_order(self): + return self._space_order + @cached_property def _shape_with_inhalo(self): """ @@ -172,37 +215,28 @@ def shape_allocated(self): @cached_property def _C_ctype(self): - return CustomDtype(self.petsc_type, modifier='*') - - @property - def petsc_type(self): - return dtype_to_petsctype(self._dtype) - - @property - def dtype(self): - return CustomDtype(self.petsc_type) + # NOTE: Reverting to using float/double instead of PetscScalar for + # simplicity when opt='advanced'. Otherwise, Temp objects must also + # be converted to PetscScalar. Additional tests are needed to + # ensure this approach is fine. Previously, issues arose from + # mismatches between precision of Function objects in Devito and the + # precision of the PETSc configuration. + # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables + # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of + # the user's PETSc configuration. + return POINTER(dtype_to_ctype(self.dtype)) @property def symbolic_shape(self): field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] + FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] # Reverse it since DMDA is setup backwards to Devito dimensions. return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - @cached_property - def dmda_info(self): - # To access the local grid info via the DM in - # PETSc you use DMDAGetLocalInfo(da, &info) - # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz - return DMDALocalInfo(name='info', liveness='eager') - - @cached_property - def indexed(self): - return PETScIndexedData(self.name, shape=self._shape, function=self.function) - - -class PETScIndexedData(IndexedBase): - pass + @property + def dmda(self): + name = 'da_so_%s' % self.space_order + return DM(name=name, liveness='eager', stencil_width=self.space_order) def dtype_to_petsctype(dtype): @@ -216,85 +250,25 @@ def dtype_to_petsctype(dtype): }[dtype] -class MatVecEq(Eq): - """ - Represents the mathematical expression of applying a linear - operator to a vector. This is a key component - for running matrix-free solvers. - """ +class InjectSolveEq(Eq): pass - - -class RHSEq(Eq): - """ - Represents the mathematical expression of building the - rhs of a linear system. - """ - pass - - -class MockEq(Eq): - """ - Represents a mock/placeholder equation to ensure distinct iteration loops. - - For example, the mat-vec action iteration loop is to be isolated from the - expression loop used to build the RHS of the linear system. This separation - facilitates the utilisation of the mat-vec iteration loop in callback functions - created at the IET level. - """ - pass - - -def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - # TODO: Add check for time dimensions and utilise implicit dimensions. - - is_time_dep = any(dim.is_Time for dim in target.dimensions) - # TODO: Current assumption is rhs is part of pde that remains - # constant at each timestep. Need to insert function to extract this from eq. - y_matvec, x_matvec, b_tmp = [ - PETScArray(name=f'{prefix}_{target.name}', - dtype=target.dtype, - dimensions=target.space_dimensions, - shape=target.grid.shape, - liveness='eager', - halo=target.halo[1:] if is_time_dep else target.halo) - for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] - - # TODO: Extend to rearrange equation for implicit time stepping. - matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), - target=target, solver_parameters=solver_parameters), - subdomain=eq.subdomain) - - # Part of pde that remains constant at each timestep - rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, - solver_parameters=solver_parameters), subdomain=eq.subdomain) - - if not bcs: - return [matvecaction, rhs] - - bcs_for_matvec = [] - for bc in bcs: - # TODO: Insert code to distiguish between essential and natural - # boundary conditions since these are treated differently within - # the solver - # NOTE: May eventually remove the essential bcs from the solve - # (and move to rhs) but for now, they are included since this - # is not trivial to implement when using DMDA - # NOTE: Below is temporary -> Just using this as a palceholder for - # the actual BC implementation for the matvec callback - new_rhs = bc.rhs.subs(target, x_matvec) - bc_rhs = LinearSolveExpr( - new_rhs, target=target, solver_parameters=solver_parameters - ) - bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) - - return [matvecaction] + bcs_for_matvec + [rhs] +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) class LinearSolveExpr(sympy.Function, Reconstructable): __rargs__ = ('expr',) - __rkwargs__ = ('target', 'solver_parameters') +<<<<<<< HEAD +<<<<<<< HEAD + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays') +======= + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', 'formfuncs', 'formrhs') +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays') +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) defaults = { 'ksp_type': 'gmres', @@ -305,7 +279,16 @@ class LinearSolveExpr(sympy.Function, Reconstructable): 'ksp_max_it': 10000 # Maximum iterations } - def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): + def __new__(cls, expr, target=None, solver_parameters=None, +<<<<<<< HEAD +<<<<<<< HEAD + matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): +======= + matvecs=None, formfuncs=None, formrhs=None, **kwargs): +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) if solver_parameters is None: solver_parameters = cls.defaults @@ -313,10 +296,22 @@ def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): for key, val in cls.defaults.items(): solver_parameters[key] = solver_parameters.get(key, val) - obj = super().__new__(cls, expr) + with sympy_mutex: + obj = sympy.Basic.__new__(cls, expr) obj._expr = expr obj._target = target obj._solver_parameters = solver_parameters + obj._matvecs = matvecs + obj._formfuncs = formfuncs + obj._formrhs = formrhs +<<<<<<< HEAD +<<<<<<< HEAD + obj._arrays = arrays +======= +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + obj._arrays = arrays +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) return obj def __repr__(self): @@ -328,7 +323,12 @@ def _sympystr(self, printer): return str(self) def __hash__(self): - return hash(self.target) + return hash(self.expr) + + def __eq__(self, other): + return (isinstance(other, LinearSolveExpr) and + self.expr == other.expr and + self.target == other.target) @property def expr(self): @@ -342,43 +342,88 @@ def target(self): def solver_parameters(self): return self._solver_parameters - func = Reconstructable._rebuild - + @property + def matvecs(self): + return self._matvecs -def petsc_lift(clusters): - """ - Lift the iteration space associated with each PETSc equation. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions. - """ - processed = [] - for c in clusters: + @property + def formfuncs(self): + return self._formfuncs - ispace = c.ispace - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + @property + def formrhs(self): + return self._formrhs - processed.append(c.rebuild(ispace=ispace)) +<<<<<<< HEAD +<<<<<<< HEAD + @property + def formrhs(self): + return self._formrhs +======= +======= + @property + def arrays(self): + return self._arrays - return processed +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) + func = Reconstructable._rebuild class PETScStruct(CompositeObject): __rargs__ = ('name', 'usr_ctx') + __rkwargs__ = ('liveness',) - def __init__(self, name, usr_ctx): + def __init__(self, name, usr_ctx, liveness='lazy'): pfields = [(i._C_name, dtype_to_ctype(i.dtype)) - for i in usr_ctx if isinstance(i, Symbol)] + for i in usr_ctx if isinstance(i, AbstractSymbol)] self._usr_ctx = usr_ctx super().__init__(name, 'MatContext', pfields) +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) + + assert liveness in ['eager', 'lazy'] + self._liveness = liveness @property - def usr_ctx(self): - return self._usr_ctx + def arrays(self): + return self._arrays +<<<<<<< HEAD + func = Reconstructable._rebuild +======= def _arg_values(self, **kwargs): + # TODO: Handle passing `struct dataobj *restrict g1_vec` and similar + # elements within in the struct. It seemed necessary to convert the Numpy array + # into a pointer to ctypes.Struct within MatContext + # (i.e override _arg_finalize for PETScStruct here?), but this approach seemed + # to result in a NIL pointer in the generated code ... any suggestions? + # Alternative: Allow these pointers to appear in the main kernel arguments + # as usual and use ctx->g1_vec = g1_vec; to assign them to the struct? values = super()._arg_values(**kwargs) - for i in self.fields: - setattr(values[self.name]._obj, i, kwargs['args'][i]) + for i in self.usr_ctx: + setattr(values[self.name]._obj, i.name, kwargs['args'][i.name]) return values + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' +<<<<<<< HEAD +<<<<<<< HEAD +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + +======= + +>>>>>>> 9d611239f (compiler: Fix derive_callback_inputs) + @property + def fields(self): + return self._usr_ctx +>>>>>>> 41157b0bc (compiler: Place dataobj pointers inside ctx struct) diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc new file mode 100644 index 0000000000..782e0ac340 --- /dev/null +++ b/docker/Dockerfile.petsc @@ -0,0 +1,49 @@ +############################################################## +# This Dockerfile builds a base image to run Devito with PETSc +############################################################# + +# Base image +FROM ubuntu:22.04 as base + +ENV DEBIAN_FRONTEND noninteractive + +# Install python +RUN apt-get update && \ + apt-get install -y dh-autoreconf python3-venv python3-dev python3-pip libopenblas-serial-dev git pkgconf mpich + +# Install for basic base not containing it +RUN apt-get install -y vim wget git flex libnuma-dev tmux \ + numactl hwloc curl \ + autoconf libtool build-essential procps + +# Install tmpi +RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi + +# Install OpenGL library, necessary for the installation of GemPy +RUN apt-get install -y libgl1-mesa-glx + +# Install numpy from source +RUN python3 -m venv /venv && \ + /venv/bin/pip install --no-cache-dir --upgrade pip && \ + /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ + rm -rf ~/.cache/pip + +# Create a directory for PETSc +RUN mkdir -p /home/petsc + +# Clone and install PETSc +RUN cd /home/petsc \ + && git clone -b release https://gitlab.com/petsc/petsc.git petsc \ + && cd petsc \ + && git pull \ + && ./configure --with-fortran-bindings=0 --with-openblas-include=$(pkgconf --variable=includedir openblas) --with-openblas-lib=$(pkgconf --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build \ + && make all + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +EXPOSE 8888 +CMD ["/bin/bash"] + +ENV PETSC_ARCH="devito_build" +ENV PETSC_DIR="/home/petsc/petsc" diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8261f69e43..c9ce6b25e5 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -615,7 +615,11 @@ def test_petsc_struct(): @skipif('petsc') +<<<<<<< HEAD @pytest.mark.parallel(mode=[2, 4, 8]) +======= +@pytest.mark.parallel(mode=1) +>>>>>>> 5ab58fee8 (dsl: Dispatch to sympy.Add not both Add and EvalDerivative) def test_apply(mode): grid = Grid(shape=(13, 13), dtype=np.float64) From 09f6fdf8d582f9798c40039e8ea01f00c3a7bc8d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 23 Jul 2024 21:18:45 +0100 Subject: [PATCH 110/151] compiler: Create mixin for ThreadCallable and PETScCallable --- .github/workflows/examples-mpi.yml | 2 +- .github/workflows/examples.yml | 2 +- .github/workflows/flake8.yml | 2 +- .github/workflows/pytest-core-mpi.yml | 2 +- .github/workflows/pytest-core-nompi.yml | 2 +- .github/workflows/tutorials.yml | 2 +- devito/ir/equations/equation.py | 4 +- devito/ir/iet/algorithms.py | 1 - devito/ir/iet/visitors.py | 1 - devito/operator/operator.py | 11 +- devito/passes/clusters/petsc.py | 25 -- devito/passes/iet/__init__.py | 1 - devito/passes/iet/definitions.py | 4 +- devito/passes/iet/petsc.py | 298 ---------------- devito/petsc/clusters.py | 19 +- devito/petsc/iet.py | 372 -------------------- devito/petsc/iet/routines.py | 6 + devito/petsc/types.py | 429 ------------------------ devito/symbolics/printer.py | 3 - devito/types/basic.py | 31 ++ docker/Dockerfile.petsc | 49 --- tests/test_petsc.py | 6 +- 22 files changed, 69 insertions(+), 1203 deletions(-) delete mode 100644 devito/passes/clusters/petsc.py delete mode 100644 devito/passes/iet/petsc.py delete mode 100644 devito/petsc/iet.py delete mode 100644 devito/petsc/types.py delete mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 9df9251991..6a804a27bb 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,7 +17,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 4490f47547..ef14b8450c 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index b61e777454..fe1761f872 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index cc4b8ab683..ab89107047 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 802aa3afb3..a022ed08e1 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index af7f17d255..9e5e1b6cb8 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 0d6c7b0ff7..67ed9269a4 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -227,13 +227,13 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - # from IPython import embed; embed() + # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - # from IPython import embed; embed() + expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index af48dc73c7..9d2db185db 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -24,7 +24,6 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: - # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 48570fa89f..d23f24a23d 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -315,7 +315,6 @@ def _args_call(self, args): def _gen_signature(self, o, is_declaration=False): decls = self._args_decl(o.parameters) - prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),)) if o.attributes: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 6a4f66f28d..f3f4e8565b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -354,7 +354,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - # from IPython import embed; embed() + return processed # Compilation -- Cluster level @@ -389,6 +389,8 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): clusters = cls._specialize_clusters(clusters, **kwargs) + clusters = petsc_project(clusters) + # Operation count after specialization final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) try: @@ -428,7 +430,6 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters - # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) @@ -498,9 +499,6 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Target-independent optimizations minimize_symbols(graph) - # If necessary, sort frees into a specific order - sort_frees(graph) - return graph.root, graph # Read-only properties exposed to the outside world @@ -527,8 +525,7 @@ def dimensions(self): @cached_property def input(self): - struct_params = derive_struct_inputs(self.parameters) - return tuple(i for i in self.parameters+struct_params if i.is_Input) + return tuple(i for i in self.parameters if i.is_Input) @cached_property def temporaries(self): diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py deleted file mode 100644 index d784d99c0d..0000000000 --- a/devito/passes/clusters/petsc.py +++ /dev/null @@ -1,25 +0,0 @@ -from devito.tools import timed_pass -from devito.petsc import LinearSolveExpr - -__all__ = ['petsc_lift'] - - -@timed_pass() -def petsc_lift(clusters): - """ - Lift the iteration space surrounding each PETSc equation to create - distinct iteration loops. This simplifys the movement of the loops - into specific callback functions generated at the IET level. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions, not the ones that stay inside the main kernel. - """ - processed = [] - for c in clusters: - - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - - return processed diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index cf6a35de90..c09db00c9b 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,4 +8,3 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa -from .petsc import * # noqa diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 98fb8cabfe..59de8f95c7 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -456,9 +456,7 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [ - i for i in bases - if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] + bases = [i for i in bases if i.name != i.function._C_name] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py deleted file mode 100644 index 8c8aeafa1c..0000000000 --- a/devito/passes/iet/petsc.py +++ /dev/null @@ -1,298 +0,0 @@ -from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, MatVecAction, - Transformer, FindSymbols, LinearSolverExpression, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition) -from devito.types import (PetscMPIInt, PETScStruct, DMDALocalInfo, DM, Mat, - Vec, KSP, PC, SNES, PetscErrorCode) -from devito.symbolics import Byref, Macro, FieldFromPointer -import cgen as c - -__all__ = ['lower_petsc'] - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't - # appear in ccode - - # Check if PETScSolve was used and count occurrences. Each PETScSolve - # will have a unique MatVecAction. - is_petsc = FindNodes(MatVecAction).visit(iet) - - if is_petsc: - - # Collect all solution fields we're solving for - targets = [i.expr.rhs.target for i in is_petsc] - - # Initalize PETSc - init = init_petsc(**kwargs) - - # Create context data struct. - struct = build_struct(iet) - - objs = build_core_objects(targets[-1], struct, **kwargs) - - # Create core PETSc calls required by general linear solves, - # which only need to be generated once e.g create DMDA. - core = core_petsc(targets[-1], objs, **kwargs) - - matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) - - main_mapper = {} - - setup = [] - efuncs = [] - for target in unique_targets: - - solver_objs = build_solver_objs(target) - - matvec_callback_body_iters = [] - - solver_setup = False - - for iter, (matvec,) in matvec_mapper.items(): - - if matvec.expr.rhs.target == target: - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec) - setup.extend(solver) - solver_setup = True - - # Only need to generate solver setup once per target - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) - solver_setup = True - - matvec_body = matvec_body_list._rebuild(body=[ - matvec_body_list.body, iter[0]]) - matvec_body_list = matvec_body_list._rebuild(body=matvec_body) - - main_mapper.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, matvec_callback_body_iters, solver_objs, objs, struct) - - setup.append(matvec_op) - setup.append(c.Line()) - - efuncs.append(matvec_callback) - - # Remove the LinSolveExpr from iet and efuncs that were used to carry - # metadata e.g solver_parameters - main_mapper.update(rebuild_expr_mapper(iet)) - efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} - - iet = Transformer(main_mapper).visit(iet) - efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] - - # Replace symbols appearing in each efunc with a pointer to the PETScStruct - efuncs = transform_efuncs(efuncs, struct) - - body = iet.body._rebuild(body=(tuple(init_setup) + iet.body.body)) - iet = iet._rebuild(body=body) - - return iet, {} - - -def init_petsc(**kwargs): - - # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve. - # NOTE: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options. - initialize = Call('PetscCall', [Call('PetscInitialize', - arguments=['NULL', 'NULL', - 'NULL', 'NULL'])]) - - return tuple([initialize]) - - -def build_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct. - usr_ctx = [] - - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - -def core_petsc(target, objs, **kwargs): - # Assumption: all targets are generated from the same Grid, - # so we can use any target. - - # MPI - call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', - arguments=[objs['comm'], - Byref(objs['size'])])]) - # Create DMDA - dmda = create_dmda(target, objs) - dm_setup = Call('PetscCall', [Call('DMSetUp', arguments=[objs['da']])]) - dm_app_ctx = Call('PetscCall', [Call('DMSetApplicationContext', - arguments=[objs['da'], objs['struct']])]) - dm_mat_type = Call('PetscCall', [Call('DMSetMatType', - arguments=[objs['da'], 'MATSHELL'])]) - dm_local_info = Call('PetscCall', [Call('DMDAGetLocalInfo', - arguments=[objs['da'], Byref(objs['info'])])]) - - return tuple([call_mpi, dmda, dm_setup, dm_app_ctx, dm_mat_type, dm_local_info]) - - -def build_core_objects(target, struct, **kwargs): - - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' - - return {'da': DM(name='da'), - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'struct': struct} - - -def create_dmda(target, objs): - - args = [objs['comm']] - - args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] - - # stencil type - args += ['DMDA_STENCIL_BOX'] - - # global dimensions - args += list(target.shape_global)[::-1] - - # no.of processors in each dimension - args += list(target.grid.distributor.topology)[::-1] - - args += [1, target.space_order] - - args += ['NULL' for _ in range(len(target.space_dimensions))] - - args += [Byref(objs['da'])] - - dmda = Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args) - - return dmda - - -def build_solver_objs(target): - - return {'Jac': Mat(name='J_'+str(target.name)), - 'x_global': Vec(name='x_global_'+str(target.name)), - 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), - 'b_global': Vec(name='b_global_'+str(target.name)), - 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), - 'ksp': KSP(name='ksp_'+str(target.name)), - 'pc': PC(name='pc_'+str(target.name)), - 'snes': SNES(name='snes_'+str(target.name)), - 'x': Vec(name='x_'+str(target.name)), - 'y': Vec(name='y_'+str(target.name))} - - -def generate_solver_calls(solver_objs, objs, matvec): - - snes_create = Call('PetscCall', [Call('SNESCreate', arguments=[ - objs['comm'], Byref(solver_objs['snes'])])]) - - snes_set_dm = Call('PetscCall', [Call('SNESSetDM', arguments=[ - solver_objs['snes'], objs['da']])]) - - create_matrix = Call('PetscCall', [Call('DMCreateMatrix', arguments=[ - objs['da'], Byref(solver_objs['Jac'])])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = Call('PetscCall', [Call('SNESSetType', arguments=[ - solver_objs['snes'], 'SNESKSPONLY'])]) - - global_x = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['x_global'])])]) - - local_x = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['x_local'])])]) - - global_b = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['b_global'])])]) - - local_b = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['b_local'])])]) - - snes_get_ksp = Call('PetscCall', [Call('SNESGetKSP', arguments=[ - solver_objs['snes'], Byref(solver_objs['ksp'])])]) - - return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, - global_x, local_x, global_b, local_b, snes_get_ksp]) - - -def create_matvec_callback(target, matvec_callback_body_iters, - solver_objs, objs, struct): - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - defn_struct = Definition(struct) - - get_context = Call('PetscCall', [Call('MatShellGetContext', - arguments=[solver_objs['Jac'], - Byref(struct)])]) - - body = List(body=[defn_struct, - get_context, - matvec_callback_body_iters]) - - matvec_callback = Callable('MyMatShellMult_'+str(target.name), - matvec_body, - retval=objs['err'], - parameters=(solver_objs['Jac'], - solver_objs['x'], - solver_objs['y'])) - - matvec_operation = Call('PetscCall', [ - Call('MatShellSetOperation', arguments=[solver_objs['Jac'], - 'MATOP_MULT', - Callback(matvec_callback.name, - Void, Void)])]) - - return matvec_callback, matvec_operation - - -def rebuild_expr_mapper(iet): - - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for - expr in FindNodes(LinearSolverExpression).visit(iet)} - - -def transform_efuncs(efuncs, struct): - - efuncs_new = [] - for efunc in efuncs: - new_body = efunc.body - for i in struct.usr_ctx: - new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) - efunc_with_new_body = efunc._rebuild(body=new_body) - efuncs_new.append(efunc_with_new_body) - - return efuncs_new - - -Null = Macro('NULL') -Void = Macro('void') - -petsc_call = String('PetscCall') -petsc_call_mpi = String('PetscCallMPI') -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') - -linear_solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 0a6679b2dd..143c73f885 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -1,5 +1,5 @@ from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr +from devito.petsc.types import LinearSolveExpr, CallbackExpr @timed_pass() @@ -25,3 +25,20 @@ def petsc_lift(clusters): else: processed.append(c) return processed + + +@timed_pass() +def petsc_project(clusters): + """ + Drop time loop for clusters which appear in PETSc callback functions. + """ + processed = [] + for c in clusters: + if isinstance(c.exprs[0].rhs, CallbackExpr): + time_dims = [d for d in c.ispace.intervals.dimensions if d.is_Time] + ispace = c.ispace.project(lambda d: d not in time_dims) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed diff --git a/devito/petsc/iet.py b/devito/petsc/iet.py deleted file mode 100644 index 659ac9f94d..0000000000 --- a/devito/petsc/iet.py +++ /dev/null @@ -1,372 +0,0 @@ -from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, - Transformer, FindSymbols, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition, BlankLine, PointerCast) -<<<<<<< HEAD -from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, -======= -from devito.petsc import (PetscMPIInt, PETScStruct, DM, Mat, ->>>>>>> b22824875 (fix circular imports) - Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) -from devito.symbolics import Byref, Macro, FieldFromPointer -import cgen as c -from devito.petsc.nodes import MatVecAction, RHSLinearSystem, LinearSolverExpression - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # Check if PETScSolve was used. - petsc_nodes = FindNodes(MatVecAction).visit(iet) - - if not petsc_nodes: - return iet, {} - - else: - # Collect all petsc solution fields - unique_targets = list(set([i.expr.rhs.target for i in petsc_nodes])) - - # Initalize PETSc - init = init_petsc(**kwargs) - - # Create context data struct - struct = build_struct(iet) - - objs = build_core_objects(unique_targets[-1], **kwargs) - - # Create core PETSc calls (not specific to each PETScSolve) - core = core_petsc(unique_targets[-1], struct, objs, **kwargs) - - matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) - - main_mapper = {} - - setup = [] - efuncs = [] - - for target in unique_targets: - - solver_objs = build_solver_objs(target) - - matvec_body_list = List() - - solver_setup = False - - for iter, (matvec,) in matvec_mapper.items(): - - # Skip the MatVecAction if it is not associated with the target - # There will be more than one MatVecAction associated with the target - # e.g interior matvec + BC matvecs - if matvec.expr.rhs.target != target: - continue - - # Only need to generate solver setup once per target - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) - solver_setup = True - - matvec_body = matvec_body_list._rebuild(body=[ - matvec_body_list.body, iter[0]]) - matvec_body_list = matvec_body_list._rebuild(body=matvec_body) - - main_mapper.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, matvec_body_list, solver_objs, objs, - struct) - - setup.append(matvec_op) - setup.append(BlankLine) - efuncs.append(matvec_callback) - - # Remove the LinSolveExpr from iet and efuncs that were used to carry - # metadata e.g solver_parameters - main_mapper.update(rebuild_expr_mapper(iet)) - efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} - - iet = Transformer(main_mapper).visit(iet) - efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] - - # Replace symbols appearing in each efunc with a pointer to the PETScStruct - efuncs = transform_efuncs(efuncs, struct) - - body = iet.body._rebuild(init=init, body=core + tuple(setup) + iet.body.body) - iet = iet._rebuild(body=body) - - return iet, {'efuncs': efuncs} - - -def init_petsc(**kwargs): - - # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve - # TODO: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options - initialize = Call(petsc_call, [ - Call('PetscInitialize', arguments=[Null, Null, Null, Null])]) - - return tuple([petsc_func_begin_user, initialize]) - - -def build_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct - usr_ctx = [] - - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - -def core_petsc(target, struct, objs, **kwargs): - - # MPI - call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', - arguments=[objs['comm'], - Byref(objs['size'])])]) - # Create DMDA - dmda = create_dmda(target, objs) - dm_setup = Call(petsc_call, [ - Call('DMSetUp', arguments=[objs['da']])]) - dm_app_ctx = Call(petsc_call, [ - Call('DMSetApplicationContext', arguments=[objs['da'], struct])]) - dm_mat_type = Call(petsc_call, [ - Call('DMSetMatType', arguments=[objs['da'], 'MATSHELL'])]) - - return tuple([petsc_func_begin_user, call_mpi, dmda, dm_setup, - dm_app_ctx, dm_mat_type, BlankLine]) - - -def build_core_objects(target, **kwargs): - - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' - - return {'da': DM(name='da', liveness='eager'), - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'err': PetscErrorCode(name='err')} - - -def create_dmda(target, objs): - - args = [objs['comm']] - - args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] - - # Stencil type - if len(target.space_dimensions) > 1: - args += ['DMDA_STENCIL_BOX'] - - # Global dimensions - args += list(target.shape_global)[::-1] - - # No.of processors in each dimension - if len(target.space_dimensions) > 1: - args += list(target.grid.distributor.topology)[::-1] - - args += [1, target.space_order] - - args += [Null for _ in range(len(target.space_dimensions))] - - args += [Byref(objs['da'])] - - dmda = Call(petsc_call, [ - Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args)]) - - return dmda - - -def build_solver_objs(target): - - return {'Jac': Mat(name='J_'+str(target.name)), - 'x_global': Vec(name='x_global_'+str(target.name)), - 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), - 'b_global': Vec(name='b_global_'+str(target.name)), - 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), - 'ksp': KSP(name='ksp_'+str(target.name)), - 'pc': PC(name='pc_'+str(target.name)), - 'snes': SNES(name='snes_'+str(target.name)), - 'X_global': Vec(name='X_global_'+str(target.name)), - 'Y_global': Vec(name='Y_global_'+str(target.name)), - 'X_local': Vec(name='X_local_'+str(target.name), liveness='eager'), - 'Y_local': Vec(name='Y_local_'+str(target.name), liveness='eager') - } - - -def generate_solver_calls(solver_objs, objs, matvec, target): - - solver_params = matvec.expr.rhs.solver_parameters - - snes_create = Call(petsc_call, [Call('SNESCreate', arguments=[ - objs['comm'], Byref(solver_objs['snes'])])]) - - snes_set_dm = Call(petsc_call, [Call('SNESSetDM', arguments=[ - solver_objs['snes'], objs['da']])]) - - create_matrix = Call(petsc_call, [Call('DMCreateMatrix', arguments=[ - objs['da'], Byref(solver_objs['Jac'])])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = Call(petsc_call, [Call('SNESSetType', arguments=[ - solver_objs['snes'], 'SNESKSPONLY'])]) - - global_x = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['x_global'])])]) - - local_x = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['x_local'])])]) - - global_b = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['b_global'])])]) - - local_b = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['b_local'])])]) - - snes_get_ksp = Call(petsc_call, [Call('SNESGetKSP', arguments=[ - solver_objs['snes'], Byref(solver_objs['ksp'])])]) - - vec_replace_array = Call(petsc_call, [ - Call('VecReplaceArray', arguments=[solver_objs[ - 'x_local'], FieldFromPointer(target._C_field_data, target._C_symbol)])]) - - ksp_set_tols = Call(petsc_call, [Call('KSPSetTolerances', arguments=[ - solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], - solver_params['ksp_divtol'], solver_params['ksp_max_it']])]) - - ksp_set_type = Call(petsc_call, [Call('KSPSetType', arguments=[ - solver_objs['ksp'], linear_solver_mapper[solver_params['ksp_type']]])]) - - ksp_get_pc = Call(petsc_call, [Call('KSPGetPC', arguments=[ - solver_objs['ksp'], Byref(solver_objs['pc'])])]) - - pc_set_type = Call(petsc_call, [Call('PCSetType', arguments=[ - solver_objs['pc'], linear_solver_mapper[solver_params['pc_type']]])]) - - ksp_set_from_ops = Call(petsc_call, [Call('KSPSetFromOptions', arguments=[ - solver_objs['ksp']])]) - - return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, - global_x, local_x, global_b, local_b, snes_get_ksp, - vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, - pc_set_type, ksp_set_from_ops]) - - -def create_matvec_callback(target, body, solver_objs, objs, struct): - - # There will be 2 PETScArrays within the body - petsc_arrays = [i.function for i in FindSymbols('indexedbases').visit(body) - if isinstance(i.function, PETScArray)] - - # There will only be one PETScArray that is written to within this body and - # one PETScArray which corresponds to the 'seed' vector - petsc_arr_write, = FindSymbols('writes').visit(body) - petsc_arr_seed, = [i for i in petsc_arrays if i.function != petsc_arr_write.function] - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - defn_struct = Definition(struct) - - mat_get_dm = Call(petsc_call, [Call('MatGetDM', arguments=[ - solver_objs['Jac'], Byref(objs['da'])])]) - - dm_get_app_context = Call(petsc_call, [Call('DMGetApplicationContext', arguments=[ - objs['da'], Byref(struct._C_symbol)])]) - - dm_get_local_xvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ - objs['da'], Byref(solver_objs['X_local'])])]) - - global_to_local_begin = Call(petsc_call, [Call('DMGlobalToLocalBegin', arguments=[ - objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) - - global_to_local_end = Call(petsc_call, [Call('DMGlobalToLocalEnd', arguments=[ - objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) - - dm_get_local_yvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ - objs['da'], Byref(solver_objs['Y_local'])])]) - - vec_get_array_y = Call(petsc_call, [Call('VecGetArray', arguments=[ - solver_objs['Y_local'], Byref(petsc_arr_write.function)])]) - - vec_get_array_x = Call(petsc_call, [Call('VecGetArray', arguments=[ - solver_objs['X_local'], Byref(petsc_arrays[0])])]) - - dm_get_local_info = Call(petsc_call, [Call('DMDAGetLocalInfo', arguments=[ - objs['da'], Byref(petsc_arrays[0].function.dmda_info)])]) - - casts = [PointerCast(i.function) for i in petsc_arrays] - - vec_restore_array_y = Call(petsc_call, [Call('VecRestoreArray', arguments=[ - solver_objs['Y_local'], Byref(petsc_arr_write[0])])]) - - vec_restore_array_x = Call(petsc_call, [Call('VecRestoreArray', arguments=[ - solver_objs['X_local'], Byref(petsc_arr_seed)])]) - - dm_local_to_global_begin = Call(petsc_call, [Call('DMLocalToGlobalBegin', arguments=[ - objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) - - dm_local_to_global_end = Call(petsc_call, [Call('DMLocalToGlobalEnd', arguments=[ - objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) - - func_return = Call('PetscFunctionReturn', arguments=[0]) - - matvec_body = List(body=[ - petsc_func_begin_user, defn_struct, mat_get_dm, dm_get_app_context, - dm_get_local_xvec, global_to_local_begin, global_to_local_end, - dm_get_local_yvec, vec_get_array_y, vec_get_array_x, dm_get_local_info, - casts, BlankLine, body, vec_restore_array_y, vec_restore_array_x, - dm_local_to_global_begin, dm_local_to_global_end, func_return]) - - matvec_callback = Callable( - 'MyMatShellMult_'+str(target.name), matvec_body, retval=objs['err'], - parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'])) - - matvec_operation = Call(petsc_call, [ - Call('MatShellSetOperation', arguments=[ - solver_objs['Jac'], 'MATOP_MULT', Callback(matvec_callback.name, - Void, Void)])]) - - return matvec_callback, matvec_operation - - -def rebuild_expr_mapper(callable): - - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for - expr in FindNodes(LinearSolverExpression).visit(callable)} - - -def transform_efuncs(efuncs, struct): - - efuncs_new = [] - for efunc in efuncs: - new_body = efunc.body - for i in struct.usr_ctx: - new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) - efunc_with_new_body = efunc._rebuild(body=new_body) - efuncs_new.append(efunc_with_new_body) - - return efuncs_new - - -Null = Macro('NULL') -Void = 'void' - -petsc_call = 'PetscCall' -petsc_call_mpi = 'PetscCallMPI' -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') - -linear_solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index a516b1bc81..f1595a6531 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -86,6 +86,8 @@ def make_matvec(self, injectsolve, objs, solver_objs): def create_matvec_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs + body = drop_callbackexpr(body) + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] body = uxreplace_time(body, solver_objs) @@ -213,6 +215,8 @@ def make_formfunc(self, injectsolve, objs, solver_objs): def create_formfunc_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs + body = drop_callbackexpr(body) + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] body = uxreplace_time(body, solver_objs) @@ -331,6 +335,8 @@ def make_formrhs(self, injectsolve, objs, solver_objs): def create_formrhs_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs + body = drop_callbackexpr(body) + dmda = objs['da_so_%s' % linsolveexpr.target.space_order] snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) diff --git a/devito/petsc/types.py b/devito/petsc/types.py deleted file mode 100644 index 72a0faa1e8..0000000000 --- a/devito/petsc/types.py +++ /dev/null @@ -1,429 +0,0 @@ -import sympy - -<<<<<<< HEAD -from devito.tools import Reconstructable, sympy_mutex -======= -from devito.tools import CustomDtype -from devito.types import LocalObject, Eq, CompositeObject -from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic -from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction, AbstractSymbol -from devito.finite_differences.tools import fd_weights_registry -from devito.tools import dtype_to_ctype, Reconstructable, sympy_mutex -from devito.symbolics import FieldFromComposite, Byref - -from devito.petsc.utils import petsc_call - - -class DM(LocalObject): - """ - PETSc Data Management object (DM). - """ - dtype = CustomDtype('DM') - - def __init__(self, *args, stencil_width=None, **kwargs): - super().__init__(*args, **kwargs) - self._stencil_width = stencil_width - - @property - def stencil_width(self): - return self._stencil_width - - @property - def info(self): - return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('DMDestroy', [Byref(self.function)]) - - -class Mat(LocalObject): - """ - PETSc Matrix object (Mat). - """ - dtype = CustomDtype('Mat') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('MatDestroy', [Byref(self.function)]) - - -class LocalVec(LocalObject): - """ - PETSc Vector object (Vec). - """ - dtype = CustomDtype('Vec') - - -class GlobalVec(LocalObject): - """ - PETSc Vector object (Vec). - """ - dtype = CustomDtype('Vec') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('VecDestroy', [Byref(self.function)]) - - -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 SNES(LocalObject): - """ - PETSc SNES : Non-Linear Systems Solvers. - """ - dtype = CustomDtype('SNES') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('SNESDestroy', [Byref(self.function)]) - - -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 DMDALocalInfo(LocalObject): - """ - PETSc object - C struct containing information - about the local grid. - """ - dtype = CustomDtype('DMDALocalInfo') - - -class PetscErrorCode(LocalObject): - """ - PETSc datatype used to return PETSc error codes. - https://petsc.org/release/manualpages/Sys/PetscErrorCode/ - """ - dtype = CustomDtype('PetscErrorCode') - - -class DummyArg(LocalObject): - dtype = CustomDtype('void', modifier='*') - - -class PETScArray(ArrayBasic, Differentiable): - """ - 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. - Differentiable enables compatability with standard Function objects, - allowing for the use of the `subs` method. - TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) - and then PETScArray(ArrayBasic). - """ - - _data_alignment = False - - # Default method for the finite difference approximation weights computation. - _default_fd = 'taylor' - - __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients', - 'space_order')) - - def __init_finalize__(self, *args, **kwargs): - - super().__init_finalize__(*args, **kwargs) - - # Symbolic (finite difference) coefficients - self._coefficients = kwargs.get('coefficients', self._default_fd) - if self._coefficients not in fd_weights_registry: - raise ValueError("coefficients must be one of %s" - " not %s" % (str(fd_weights_registry), self._coefficients)) - self._shape = kwargs.get('shape') - self._space_order = kwargs.get('space_order', 1) - - @classmethod - def __dtype_setup__(cls, **kwargs): - return kwargs.get('dtype', np.float32) - - @property - def coefficients(self): - """Form of the coefficients of the function.""" - return self._coefficients - - @property - def shape(self): - return self._shape - - @property - def space_order(self): - return self._space_order - - @cached_property - def _shape_with_inhalo(self): - """ - Shape of the domain+inhalo region. The inhalo region comprises the - outhalo as well as any additional "ghost" layers for MPI halo - exchanges. Data in the inhalo region are exchanged when running - Operators to maintain consistent values as in sequential runs. - - Notes - ----- - Typically, this property won't be used in user code, but it may come - in handy for testing or debugging - """ - return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) - - @cached_property - def shape_allocated(self): - """ - Shape of the allocated data of the Function type object from which - this PETScArray was derived. It includes the domain and inhalo regions, - as well as any additional padding surrounding the halo. - - Notes - ----- - In an MPI context, this is the *local* with_halo region shape. - """ - return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, - self._padding)], - getters=self.dimensions) - - @cached_property - def _C_ctype(self): - # NOTE: Reverting to using float/double instead of PetscScalar for - # simplicity when opt='advanced'. Otherwise, Temp objects must also - # be converted to PetscScalar. Additional tests are needed to - # ensure this approach is fine. Previously, issues arose from - # mismatches between precision of Function objects in Devito and the - # precision of the PETSc configuration. - # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables - # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of - # the user's PETSc configuration. - return POINTER(dtype_to_ctype(self.dtype)) - - @property - def symbolic_shape(self): - field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] - # Reverse it since DMDA is setup backwards to Devito dimensions. - return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - - @property - def dmda(self): - name = 'da_so_%s' % self.space_order - return DM(name=name, liveness='eager', stencil_width=self.space_order) - - -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] - - -class InjectSolveEq(Eq): - pass ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) - - -class LinearSolveExpr(sympy.Function, Reconstructable): - - __rargs__ = ('expr',) -<<<<<<< HEAD -<<<<<<< HEAD - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays') -======= - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', 'formfuncs', 'formrhs') ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays') ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - - defaults = { - 'ksp_type': 'gmres', - 'pc_type': 'jacobi', - 'ksp_rtol': 1e-7, # Relative tolerance - 'ksp_atol': 1e-50, # Absolute tolerance - 'ksp_divtol': 1e4, # Divergence tolerance - 'ksp_max_it': 10000 # Maximum iterations - } - - def __new__(cls, expr, target=None, solver_parameters=None, -<<<<<<< HEAD -<<<<<<< HEAD - matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): -======= - matvecs=None, formfuncs=None, formrhs=None, **kwargs): ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - - if solver_parameters is None: - solver_parameters = cls.defaults - else: - for key, val in cls.defaults.items(): - solver_parameters[key] = solver_parameters.get(key, val) - - with sympy_mutex: - obj = sympy.Basic.__new__(cls, expr) - obj._expr = expr - obj._target = target - obj._solver_parameters = solver_parameters - obj._matvecs = matvecs - obj._formfuncs = formfuncs - obj._formrhs = formrhs -<<<<<<< HEAD -<<<<<<< HEAD - obj._arrays = arrays -======= ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - obj._arrays = arrays ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - return obj - - def __repr__(self): - return "%s(%s)" % (self.__class__.__name__, self.expr) - - __str__ = __repr__ - - def _sympystr(self, printer): - return str(self) - - def __hash__(self): - return hash(self.expr) - - def __eq__(self, other): - return (isinstance(other, LinearSolveExpr) and - self.expr == other.expr and - self.target == other.target) - - @property - def expr(self): - return self._expr - - @property - def target(self): - return self._target - - @property - def solver_parameters(self): - return self._solver_parameters - - @property - def matvecs(self): - return self._matvecs - - @property - def formfuncs(self): - return self._formfuncs - - @property - def formrhs(self): - return self._formrhs - -<<<<<<< HEAD -<<<<<<< HEAD - @property - def formrhs(self): - return self._formrhs -======= -======= - @property - def arrays(self): - return self._arrays - ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - func = Reconstructable._rebuild - - -class PETScStruct(CompositeObject): - - __rargs__ = ('name', 'usr_ctx') - __rkwargs__ = ('liveness',) - - def __init__(self, name, usr_ctx, liveness='lazy'): - pfields = [(i._C_name, dtype_to_ctype(i.dtype)) - for i in usr_ctx if isinstance(i, AbstractSymbol)] - self._usr_ctx = usr_ctx - super().__init__(name, 'MatContext', pfields) ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) - - assert liveness in ['eager', 'lazy'] - self._liveness = liveness - - @property - def arrays(self): - return self._arrays - -<<<<<<< HEAD - func = Reconstructable._rebuild -======= - def _arg_values(self, **kwargs): - # TODO: Handle passing `struct dataobj *restrict g1_vec` and similar - # elements within in the struct. It seemed necessary to convert the Numpy array - # into a pointer to ctypes.Struct within MatContext - # (i.e override _arg_finalize for PETScStruct here?), but this approach seemed - # to result in a NIL pointer in the generated code ... any suggestions? - # Alternative: Allow these pointers to appear in the main kernel arguments - # as usual and use ctx->g1_vec = g1_vec; to assign them to the struct? - values = super()._arg_values(**kwargs) - for i in self.usr_ctx: - setattr(values[self.name]._obj, i.name, kwargs['args'][i.name]) - return values - - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' -<<<<<<< HEAD -<<<<<<< HEAD ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - -======= - ->>>>>>> 9d611239f (compiler: Fix derive_callback_inputs) - @property - def fields(self): - return self._usr_ctx ->>>>>>> 41157b0bc (compiler: Place dataobj pointers inside ctx struct) diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index e7b395d436..2c366a389e 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,9 +242,6 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) - def _print_PETScRHS(self, expr): - return "%s" % self._print(expr._expr) - _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/devito/types/basic.py b/devito/types/basic.py index 14e3949fb8..660ceff437 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1513,6 +1513,37 @@ def __getnewargs_ex__(self): return args, kwargs +class LocalType(Basic): + """ + This is the abstract base class for local types, which are + generated by the compiler in C rather than in Python. + Notes + ----- + Subclasses should setup `_liveness`. + """ + + is_LocalType = True + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' + + """ + A modifier added to the subclass C declaration when it appears + in a function signature. For example, a subclass might define `_C_modifier = '&'` + to impose pass-by-reference semantics. + """ + _C_modifier = None + + # Extended SymPy hierarchy follows, for essentially two reasons: # - To keep track of `function` # - To override SymPy caching behaviour diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc deleted file mode 100644 index 782e0ac340..0000000000 --- a/docker/Dockerfile.petsc +++ /dev/null @@ -1,49 +0,0 @@ -############################################################## -# This Dockerfile builds a base image to run Devito with PETSc -############################################################# - -# Base image -FROM ubuntu:22.04 as base - -ENV DEBIAN_FRONTEND noninteractive - -# Install python -RUN apt-get update && \ - apt-get install -y dh-autoreconf python3-venv python3-dev python3-pip libopenblas-serial-dev git pkgconf mpich - -# Install for basic base not containing it -RUN apt-get install -y vim wget git flex libnuma-dev tmux \ - numactl hwloc curl \ - autoconf libtool build-essential procps - -# Install tmpi -RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi - -# Install OpenGL library, necessary for the installation of GemPy -RUN apt-get install -y libgl1-mesa-glx - -# Install numpy from source -RUN python3 -m venv /venv && \ - /venv/bin/pip install --no-cache-dir --upgrade pip && \ - /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ - rm -rf ~/.cache/pip - -# Create a directory for PETSc -RUN mkdir -p /home/petsc - -# Clone and install PETSc -RUN cd /home/petsc \ - && git clone -b release https://gitlab.com/petsc/petsc.git petsc \ - && cd petsc \ - && git pull \ - && ./configure --with-fortran-bindings=0 --with-openblas-include=$(pkgconf --variable=includedir openblas) --with-openblas-lib=$(pkgconf --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build \ - && make all - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -EXPOSE 8888 -CMD ["/bin/bash"] - -ENV PETSC_ARCH="devito_build" -ENV PETSC_DIR="/home/petsc/petsc" diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c9ce6b25e5..05ff00c70d 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -615,11 +615,7 @@ def test_petsc_struct(): @skipif('petsc') -<<<<<<< HEAD @pytest.mark.parallel(mode=[2, 4, 8]) -======= -@pytest.mark.parallel(mode=1) ->>>>>>> 5ab58fee8 (dsl: Dispatch to sympy.Add not both Add and EvalDerivative) def test_apply(mode): grid = Grid(shape=(13, 13), dtype=np.float64) @@ -786,4 +782,4 @@ def test_time_loop(): body4 = str(op4.body) assert 'ctx.t0 = t0' in body4 - assert body4.count('ctx.t0 = t0') == 1 + assert body4.count('ctx.t0 = t0') == 1 \ No newline at end of file From 09d9180613461d1f242b8671ab480fbfa653ef29 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 3 Jan 2025 11:54:09 +0000 Subject: [PATCH 111/151] types: Edit freesymbols in FieldFromPointer --- devito/ir/equations/algorithms.py | 7 +++++-- devito/operator/operator.py | 1 + devito/petsc/iet/routines.py | 10 +++++++--- devito/petsc/solve.py | 2 +- devito/symbolics/extended_sympy.py | 4 ++++ devito/symbolics/inspection.py | 9 ++++----- tests/test_dimension.py | 25 ++++++++++++++++++++++++- tests/test_petsc.py | 12 ++++++------ tests/test_subdomains.py | 3 ++- tests/test_symbolics.py | 17 ++++++++++++++--- 10 files changed, 68 insertions(+), 22 deletions(-) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 946eb4324b..2a52b1a41d 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -169,10 +169,10 @@ def concretize_subdims(exprs, **kwargs): across `exprs`, such as the thickness symbols. """ sregistry = kwargs.get('sregistry') + mapper = kwargs.get('concretize_mapper') - mapper = {} rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions - + # from IPython import embed; embed() _concretize_subdims(exprs, mapper, rebuilt, sregistry) if not mapper: return exprs @@ -206,6 +206,7 @@ def _(v, mapper, rebuilt, sregistry): @_concretize_subdims.register(Eq) def _(expr, mapper, rebuilt, sregistry): + # from IPython import embed; embed() for d in expr.free_symbols: _concretize_subdims(d, mapper, rebuilt, sregistry) @@ -215,6 +216,7 @@ def _(expr, mapper, rebuilt, sregistry): @_concretize_subdims.register(Thickness) def _(tkn, mapper, rebuilt, sregistry): + # from IPython import embed; embed() if tkn in mapper: # Already have a substitution for this thickness return @@ -224,6 +226,7 @@ def _(tkn, mapper, rebuilt, sregistry): @_concretize_subdims.register(SubDimension) def _(d, mapper, rebuilt, sregistry): + # from IPython import embed; embed() if d in mapper: # Already have a substitution for this dimension return diff --git a/devito/operator/operator.py b/devito/operator/operator.py index f3f4e8565b..73fb6f1450 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -263,6 +263,7 @@ def _lower(cls, expressions, **kwargs): """ # Create a symbol registry kwargs.setdefault('sregistry', SymbolRegistry()) + kwargs.setdefault('concretize_mapper', {}) expressions = as_tuple(expressions) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f1595a6531..9fee7660f6 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -26,6 +26,7 @@ def __new__(cls, rcompile=None, sregistry=None, **kwargs): obj.sregistry = sregistry obj._efuncs = OrderedDict() obj._struct_params = [] + obj.concretize_mapper = kwargs.get('concretize_mapper', {}) return obj @@ -69,7 +70,8 @@ def make_all(self, injectsolve, objs, solver_objs): def make_matvec(self, injectsolve, objs, solver_objs): # Compile matvec `eqns` into an IET via recursive compilation irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, - options={'mpi': False}, sregistry=SymbolRegistry()) + options={'mpi': False}, sregistry=SymbolRegistry(), + concretize_mapper=self.concretize_mapper) body_matvec = self.create_matvec_body(injectsolve, List(body=irs_matvec.uiet.body), solver_objs, objs) @@ -198,7 +200,8 @@ def make_formfunc(self, injectsolve, objs, solver_objs): # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( injectsolve.expr.rhs.formfuncs, - options={'mpi': False}, sregistry=SymbolRegistry() + options={'mpi': False}, sregistry=SymbolRegistry(), + concretize_mapper=self.concretize_mapper ) body_formfunc = self.create_formfunc_body(injectsolve, List(body=irs_formfunc.uiet.body), @@ -318,7 +321,8 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): def make_formrhs(self, injectsolve, objs, solver_objs): # Compile formrhs `eqns` into an IET via recursive compilation irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, - options={'mpi': False}, sregistry=SymbolRegistry()) + options={'mpi': False}, sregistry=SymbolRegistry(), + concretize_mapper=self.concretize_mapper) body_formrhs = self.create_formrhs_body(injectsolve, List(body=irs_formrhs.uiet.body), solver_objs, objs) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index c5df6b859f..da4941f799 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -74,7 +74,7 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): formrhs=formrhs, arrays=arrays, time_mapper=time_mapper, - ), subdomain=eq.subdomain) + )) return [inject_solve] diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index 5b13262ded..9db65d8bb2 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -252,6 +252,10 @@ def __str__(self): def field(self): return self.call + @property + def dtype(self): + return self.field.dtype + __repr__ = __str__ diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 18e2623764..64eeeabb9d 100644 --- a/devito/symbolics/inspection.py +++ b/devito/symbolics/inspection.py @@ -296,9 +296,8 @@ def sympy_dtype(expr, base=None): Infer the dtype of the expression. """ dtypes = {base} - {None} - for i in expr.free_symbols: - try: - dtypes.add(i.dtype) - except AttributeError: - pass + for i in expr.args: + dtype = getattr(i, 'dtype', None) + if dtype: + dtypes.add(dtype) return infer_dtype(dtypes) diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 5ed0bd08a5..0132871a47 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -2064,7 +2064,8 @@ def test_correct_thicknesses(self): ix1 = SubDimension.right('x', x, 2) ix2 = SubDimension.middle('x', x, 2, 2) - rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry()) + rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry(), + concretize_mapper={}) assert rebuilt[0].is_left assert rebuilt[1].is_right @@ -2105,3 +2106,25 @@ def test_condition_concretization(self): assert len({c.condition for c in conditionals}) == 2 for c in conditionals: assert c.condition in expected_conditionals + + def test_repeat_concretization(self): + """ + Ensure that SubDimensions are consistently concretized to the same object + across multiple calls to the function. This is necessary when using + `rcompile` on equations with SubDimensions. + """ + + grid = Grid((2, 2)) + + x = Dimension('x') + ix = SubDimension.middle('ix', x, 2, 2) + + u = Function(name='u', grid=grid) + eq = Eq(u, ix.ltkn) + + kwargs = {'sregistry': SymbolRegistry(), 'concretize_mapper': {}} + + exprs1 = concretize_subdims([eq], **kwargs) + exprs2 = concretize_subdims([eq], **kwargs) + + assert exprs1 == exprs2 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 05ff00c70d..49d93fe745 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -129,12 +129,12 @@ def test_petsc_solve(): rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) assert str(action_expr[-1].expr.rhs) == \ - 'matvec->h_x**(-2)*x_matvec_f[x + 1, y + 2]' + \ - ' - 2.0*matvec->h_x**(-2)*x_matvec_f[x + 2, y + 2]' + \ - ' + matvec->h_x**(-2)*x_matvec_f[x + 3, y + 2]' + \ - ' + matvec->h_y**(-2)*x_matvec_f[x + 2, y + 1]' + \ - ' - 2.0*matvec->h_y**(-2)*x_matvec_f[x + 2, y + 2]' + \ - ' + matvec->h_y**(-2)*x_matvec_f[x + 2, y + 3]' + 'x_matvec_f[x + 1, y + 2]/matvec->h_x**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/matvec->h_x**2' + \ + ' + x_matvec_f[x + 3, y + 2]/matvec->h_x**2' + \ + ' + x_matvec_f[x + 2, y + 1]/matvec->h_y**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/matvec->h_y**2' + \ + ' + x_matvec_f[x + 2, y + 3]/matvec->h_y**2' assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index 784dc1c3ad..db08f60de6 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -37,7 +37,8 @@ def define(self, dimensions): with timed_region('x'): # _lower_exprs expects a SymbolRegistry, so create one expr = Operator._lower_exprs([eq0], options={}, - sregistry=SymbolRegistry())[0] + sregistry=SymbolRegistry(), + concretize_mapper={})[0] assert str(expr.rhs) == 'ix*f[ix + 1, iy + 1] + iy' def test_multiple_middle(self): diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 2bae5679c8..56778879e9 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -14,10 +14,10 @@ CallFromPointer, Cast, DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, ccode, uxreplace, - retrieve_derivatives) + retrieve_derivatives, sympy_dtype) from devito.tools import as_tuple from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, - ComponentAccess, StencilDimension, Symbol as dSymbol) + Symbol as dSymbol, CompositeObject) from devito.types.basic import AbstractSymbol @@ -249,6 +249,17 @@ def test_field_from_pointer(): # Free symbols assert ffp1.free_symbols == {s} + # Test dtype + f = dSymbol('f') + pfields = [(f._C_name, f._C_ctype)] + struct = CompositeObject('s1', 'myStruct', pfields) + ffp4 = FieldFromPointer(f, struct) + assert str(ffp4) == 's1->f' + assert ffp4.dtype == f.dtype + expr = 1/ffp4 + dtype = sympy_dtype(expr) + assert dtype == f.dtype + def test_field_from_composite(): s = Symbol('s') @@ -293,7 +304,7 @@ def test_extended_sympy_arithmetic(): # noncommutative o = Object(name='o', dtype=c_void_p) bar = FieldFromPointer('bar', o) - assert ccode(-1 + bar) == '-1 + o->bar' + assert ccode(-1 + bar) == 'o->bar - 1' def test_integer_abs(): From 8472b77bc4378eab0261843bd0add024577c67ae Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 8 Jan 2025 10:19:44 +0000 Subject: [PATCH 112/151] clean --- devito/ir/equations/algorithms.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 2a52b1a41d..b48e96c711 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -172,7 +172,7 @@ def concretize_subdims(exprs, **kwargs): mapper = kwargs.get('concretize_mapper') rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions - # from IPython import embed; embed() + _concretize_subdims(exprs, mapper, rebuilt, sregistry) if not mapper: return exprs @@ -206,7 +206,6 @@ def _(v, mapper, rebuilt, sregistry): @_concretize_subdims.register(Eq) def _(expr, mapper, rebuilt, sregistry): - # from IPython import embed; embed() for d in expr.free_symbols: _concretize_subdims(d, mapper, rebuilt, sregistry) @@ -216,7 +215,6 @@ def _(expr, mapper, rebuilt, sregistry): @_concretize_subdims.register(Thickness) def _(tkn, mapper, rebuilt, sregistry): - # from IPython import embed; embed() if tkn in mapper: # Already have a substitution for this thickness return @@ -226,7 +224,6 @@ def _(tkn, mapper, rebuilt, sregistry): @_concretize_subdims.register(SubDimension) def _(d, mapper, rebuilt, sregistry): - # from IPython import embed; embed() if d in mapper: # Already have a substitution for this dimension return From 579045f6a978cbc036a74ac3f92d1eed21a8560c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 8 Jan 2025 22:23:49 +0000 Subject: [PATCH 113/151] trigger wfs for branches off this branch --- .github/workflows/examples-mpi.yml | 1 + .github/workflows/examples.yml | 1 + .github/workflows/flake8.yml | 1 + .github/workflows/pytest-core-mpi.yml | 1 + .github/workflows/pytest-core-nompi.yml | 1 + .github/workflows/pytest-petsc.yml | 1 + .github/workflows/tutorials.yml | 1 + 7 files changed, 7 insertions(+) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 6a804a27bb..bbf3105635 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -21,6 +21,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: build: diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index ef14b8450c..85f84f0cd5 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -14,6 +14,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: tutorials: diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index fe1761f872..92556b6be4 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -14,6 +14,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: flake8: diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index ab89107047..2eae16021b 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -14,6 +14,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: test-mpi-basic: diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index a022ed08e1..5cf374f222 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -14,6 +14,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: pytest: diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 6f0a729e95..20bfc645d1 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -13,6 +13,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: pytest: diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index 9e5e1b6cb8..93d64e26ac 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -14,6 +14,7 @@ on: pull_request: branches: - master + - FieldFromPointer jobs: tutorials: From 79cf51a9da277f3bd03f8d695f04650f456e32d0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 13:39:57 +0000 Subject: [PATCH 114/151] rebase leftover --- .github/workflows/examples-mpi.yml | 1 - .github/workflows/examples.yml | 1 - .github/workflows/flake8.yml | 1 - .github/workflows/pytest-core-mpi.yml | 1 - devito/ir/iet/visitors.py | 1 + devito/operator/operator.py | 2 -- devito/petsc/clusters.py | 19 +--------------- devito/petsc/iet/routines.py | 6 ------ devito/types/basic.py | 31 --------------------------- 9 files changed, 2 insertions(+), 61 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index bbf3105635..59cd08eaa5 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,7 +17,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 85f84f0cd5..94335a9e2d 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 92556b6be4..336cedb294 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 2eae16021b..67f841cabd 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index d23f24a23d..48570fa89f 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -315,6 +315,7 @@ def _args_call(self, args): def _gen_signature(self, o, is_declaration=False): decls = self._args_decl(o.parameters) + prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),)) if o.attributes: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 73fb6f1450..23f2b5b3f3 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -390,8 +390,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): clusters = cls._specialize_clusters(clusters, **kwargs) - clusters = petsc_project(clusters) - # Operation count after specialization final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) try: diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 143c73f885..0a6679b2dd 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -1,5 +1,5 @@ from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr, CallbackExpr +from devito.petsc.types import LinearSolveExpr @timed_pass() @@ -25,20 +25,3 @@ def petsc_lift(clusters): else: processed.append(c) return processed - - -@timed_pass() -def petsc_project(clusters): - """ - Drop time loop for clusters which appear in PETSc callback functions. - """ - processed = [] - for c in clusters: - if isinstance(c.exprs[0].rhs, CallbackExpr): - time_dims = [d for d in c.ispace.intervals.dimensions if d.is_Time] - ispace = c.ispace.project(lambda d: d not in time_dims) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - - return processed diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 9fee7660f6..06f848862f 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -88,8 +88,6 @@ def make_matvec(self, injectsolve, objs, solver_objs): def create_matvec_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs - body = drop_callbackexpr(body) - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] body = uxreplace_time(body, solver_objs) @@ -218,8 +216,6 @@ def make_formfunc(self, injectsolve, objs, solver_objs): def create_formfunc_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs - body = drop_callbackexpr(body) - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] body = uxreplace_time(body, solver_objs) @@ -339,8 +335,6 @@ def make_formrhs(self, injectsolve, objs, solver_objs): def create_formrhs_body(self, injectsolve, body, solver_objs, objs): linsolveexpr = injectsolve.expr.rhs - body = drop_callbackexpr(body) - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) diff --git a/devito/types/basic.py b/devito/types/basic.py index 660ceff437..14e3949fb8 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1513,37 +1513,6 @@ def __getnewargs_ex__(self): return args, kwargs -class LocalType(Basic): - """ - This is the abstract base class for local types, which are - generated by the compiler in C rather than in Python. - Notes - ----- - Subclasses should setup `_liveness`. - """ - - is_LocalType = True - - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' - - """ - A modifier added to the subclass C declaration when it appears - in a function signature. For example, a subclass might define `_C_modifier = '&'` - to impose pass-by-reference semantics. - """ - _C_modifier = None - - # Extended SymPy hierarchy follows, for essentially two reasons: # - To keep track of `function` # - To override SymPy caching behaviour From 4272da2e22ca512f90854a9c7d84983fba999023 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 13:48:27 +0000 Subject: [PATCH 115/151] rebase fix and remove concretize mapper --- .github/workflows/pytest-core-nompi.yml | 1 - .github/workflows/tutorials.yml | 1 - devito/ir/equations/algorithms.py | 2 +- devito/operator/operator.py | 1 - devito/petsc/iet/routines.py | 10 +++------- tests/test_dimension.py | 25 +------------------------ tests/test_subdomains.py | 4 ++-- 7 files changed, 7 insertions(+), 37 deletions(-) diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 5cf374f222..df3a050927 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index 93d64e26ac..c792dd6acd 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index b48e96c711..29d7af9614 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -169,7 +169,7 @@ def concretize_subdims(exprs, **kwargs): across `exprs`, such as the thickness symbols. """ sregistry = kwargs.get('sregistry') - mapper = kwargs.get('concretize_mapper') + mapper = {} rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 23f2b5b3f3..fe4b5ab48c 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -263,7 +263,6 @@ def _lower(cls, expressions, **kwargs): """ # Create a symbol registry kwargs.setdefault('sregistry', SymbolRegistry()) - kwargs.setdefault('concretize_mapper', {}) expressions = as_tuple(expressions) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 06f848862f..ac2ba60f7f 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -26,7 +26,6 @@ def __new__(cls, rcompile=None, sregistry=None, **kwargs): obj.sregistry = sregistry obj._efuncs = OrderedDict() obj._struct_params = [] - obj.concretize_mapper = kwargs.get('concretize_mapper', {}) return obj @@ -70,8 +69,7 @@ def make_all(self, injectsolve, objs, solver_objs): def make_matvec(self, injectsolve, objs, solver_objs): # Compile matvec `eqns` into an IET via recursive compilation irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, - options={'mpi': False}, sregistry=SymbolRegistry(), - concretize_mapper=self.concretize_mapper) + options={'mpi': False}, sregistry=self.sregistry) body_matvec = self.create_matvec_body(injectsolve, List(body=irs_matvec.uiet.body), solver_objs, objs) @@ -198,8 +196,7 @@ def make_formfunc(self, injectsolve, objs, solver_objs): # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( injectsolve.expr.rhs.formfuncs, - options={'mpi': False}, sregistry=SymbolRegistry(), - concretize_mapper=self.concretize_mapper + options={'mpi': False}, sregistry=self.sregistry ) body_formfunc = self.create_formfunc_body(injectsolve, List(body=irs_formfunc.uiet.body), @@ -317,8 +314,7 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): def make_formrhs(self, injectsolve, objs, solver_objs): # Compile formrhs `eqns` into an IET via recursive compilation irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, - options={'mpi': False}, sregistry=SymbolRegistry(), - concretize_mapper=self.concretize_mapper) + options={'mpi': False}, sregistry=self.sregistry) body_formrhs = self.create_formrhs_body(injectsolve, List(body=irs_formrhs.uiet.body), solver_objs, objs) diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 0132871a47..5ed0bd08a5 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -2064,8 +2064,7 @@ def test_correct_thicknesses(self): ix1 = SubDimension.right('x', x, 2) ix2 = SubDimension.middle('x', x, 2, 2) - rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry(), - concretize_mapper={}) + rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry()) assert rebuilt[0].is_left assert rebuilt[1].is_right @@ -2106,25 +2105,3 @@ def test_condition_concretization(self): assert len({c.condition for c in conditionals}) == 2 for c in conditionals: assert c.condition in expected_conditionals - - def test_repeat_concretization(self): - """ - Ensure that SubDimensions are consistently concretized to the same object - across multiple calls to the function. This is necessary when using - `rcompile` on equations with SubDimensions. - """ - - grid = Grid((2, 2)) - - x = Dimension('x') - ix = SubDimension.middle('ix', x, 2, 2) - - u = Function(name='u', grid=grid) - eq = Eq(u, ix.ltkn) - - kwargs = {'sregistry': SymbolRegistry(), 'concretize_mapper': {}} - - exprs1 = concretize_subdims([eq], **kwargs) - exprs2 = concretize_subdims([eq], **kwargs) - - assert exprs1 == exprs2 diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index db08f60de6..ffae1dd229 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -37,8 +37,8 @@ def define(self, dimensions): with timed_region('x'): # _lower_exprs expects a SymbolRegistry, so create one expr = Operator._lower_exprs([eq0], options={}, - sregistry=SymbolRegistry(), - concretize_mapper={})[0] + sregistry=SymbolRegistry())[0] + assert str(expr.rhs) == 'ix*f[ix + 1, iy + 1] + iy' def test_multiple_middle(self): From 526449ecc12c0568616de17790c6038dc2dba4b8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 13:50:19 +0000 Subject: [PATCH 116/151] rebase leftover --- devito/ir/equations/algorithms.py | 1 - tests/test_subdomains.py | 1 - 2 files changed, 2 deletions(-) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 29d7af9614..bbbdf71c80 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -170,7 +170,6 @@ def concretize_subdims(exprs, **kwargs): """ sregistry = kwargs.get('sregistry') mapper = {} - rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions _concretize_subdims(exprs, mapper, rebuilt, sregistry) diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index ffae1dd229..784dc1c3ad 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -38,7 +38,6 @@ def define(self, dimensions): # _lower_exprs expects a SymbolRegistry, so create one expr = Operator._lower_exprs([eq0], options={}, sregistry=SymbolRegistry())[0] - assert str(expr.rhs) == 'ix*f[ix + 1, iy + 1] + iy' def test_multiple_middle(self): From 4689508eb2a8854c45debdbfffa61798c38537a8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 13:54:34 +0000 Subject: [PATCH 117/151] flake8 --- devito/ir/equations/algorithms.py | 1 + tests/test_petsc.py | 2 +- tests/test_symbolics.py | 3 ++- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index bbbdf71c80..946eb4324b 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -169,6 +169,7 @@ def concretize_subdims(exprs, **kwargs): across `exprs`, such as the thickness symbols. """ sregistry = kwargs.get('sregistry') + mapper = {} rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 49d93fe745..1c70331066 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -782,4 +782,4 @@ def test_time_loop(): body4 = str(op4.body) assert 'ctx.t0 = t0' in body4 - assert body4.count('ctx.t0 = t0') == 1 \ No newline at end of file + assert body4.count('ctx.t0 = t0') == 1 diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 56778879e9..4ae6fc2cdf 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -17,7 +17,8 @@ retrieve_derivatives, sympy_dtype) from devito.tools import as_tuple from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, - Symbol as dSymbol, CompositeObject) + ComponentAccess, StencilDimension, Symbol as dSymbol, + CompositeObject) from devito.types.basic import AbstractSymbol From 434e4241a6ea5f40115b992eee8ff8c4550aaa53 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 13:57:46 +0000 Subject: [PATCH 118/151] flake8 --- devito/petsc/iet/routines.py | 1 - 1 file changed, 1 deletion(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index ac2ba60f7f..5ef8b2aa5f 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -13,7 +13,6 @@ from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatVecCallback) from devito.petsc.iet.utils import petsc_call, petsc_struct -from devito.ir.support import SymbolRegistry class PETScCallbackBuilder: From cfa72188150080f33aa53902868c94212b7916e9 Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 119/151] misc: Split Dockerfile.cpu into .cpu and .intel --- docker/Dockerfile.intel | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index e013fccc03..1fab7147d6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,6 +35,7 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu +<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -50,6 +51,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +======= +# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal +RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg +RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list + +RUN apt-get update -y && apt-get dist-upgrade -y && \ + apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ + intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ + libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -76,7 +87,11 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD # ICX OpenMP image +======= +# ICX image +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -93,6 +108,35 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD +======= +# ICX hpc image +############################################################## +FROM oneapi as icx-hpc + +# Install both icc and icx to avoid missing dependencies +RUN apt-get update -y && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic + +# Missig components +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager +RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ + chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ + rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# Devito config +ENV DEVITO_ARCH="icx" +ENV DEVITO_LANGUAGE="openmp" +# MPICC compiler for mpi4py +ENV MPICC=mpiicc +ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' + +############################################################## +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -107,6 +151,12 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl +<<<<<<< HEAD +======= +# NOTE: the name of this file ends with ".cpu" but this is a GPU image. +# It then feels a bit akward, so some restructuring might be needed + +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From fd8c53fee458fcd0271c4c143ce9a70f333c4ba5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 120/151] misc: Update file due to incorrect rebase --- docker/Dockerfile.intel | 50 ----------------------------------------- 1 file changed, 50 deletions(-) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index 1fab7147d6..e013fccc03 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,7 +35,6 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu -<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -51,16 +50,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev -======= -# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal -RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg -RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list - -RUN apt-get update -y && apt-get dist-upgrade -y && \ - apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ - intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ - libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -87,11 +76,7 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD # ICX OpenMP image -======= -# ICX image ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -108,35 +93,6 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD -======= -# ICX hpc image -############################################################## -FROM oneapi as icx-hpc - -# Install both icc and icx to avoid missing dependencies -RUN apt-get update -y && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic - -# Missig components -# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager -RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ - chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ - rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -# Devito config -ENV DEVITO_ARCH="icx" -ENV DEVITO_LANGUAGE="openmp" -# MPICC compiler for mpi4py -ENV MPICC=mpiicc -ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' - -############################################################## ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -151,12 +107,6 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl -<<<<<<< HEAD -======= -# NOTE: the name of this file ends with ".cpu" but this is a GPU image. -# It then feels a bit akward, so some restructuring might be needed - ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From d09e654295a09182c438d3634ee72a9c5b57d801 Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 121/151] misc: Split Dockerfile.cpu into .cpu and .intel --- docker/Dockerfile.intel | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index e013fccc03..1fab7147d6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,6 +35,7 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu +<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -50,6 +51,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +======= +# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal +RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg +RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list + +RUN apt-get update -y && apt-get dist-upgrade -y && \ + apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ + intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ + libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -76,7 +87,11 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD # ICX OpenMP image +======= +# ICX image +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -93,6 +108,35 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## +<<<<<<< HEAD +======= +# ICX hpc image +############################################################## +FROM oneapi as icx-hpc + +# Install both icc and icx to avoid missing dependencies +RUN apt-get update -y && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ + apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic + +# Missig components +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager +RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ + chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ + rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# Devito config +ENV DEVITO_ARCH="icx" +ENV DEVITO_LANGUAGE="openmp" +# MPICC compiler for mpi4py +ENV MPICC=mpiicc +ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' + +############################################################## +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -107,6 +151,12 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl +<<<<<<< HEAD +======= +# NOTE: the name of this file ends with ".cpu" but this is a GPU image. +# It then feels a bit akward, so some restructuring might be needed + +>>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 62a4e1cce32fdef78b8fa37449c33c7de0df8277 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 122/151] misc: Update file due to incorrect rebase --- devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 2 +- devito/ir/iet/algorithms.py | 1 + devito/ir/iet/nodes.py | 4 + devito/ir/stree/algorithms.py | 1 + devito/operator/operator.py | 1 + devito/passes/iet/definitions.py | 4 +- devito/petsc/iet/passes.py | 272 ++++------------------ devito/petsc/solve.py | 35 ++- devito/petsc/types.py | 384 +++++++++++++++++++++++++++++++ devito/petsc/utils.py | 13 ++ devito/symbolics/printer.py | 3 + docker/Dockerfile.intel | 50 ---- 13 files changed, 485 insertions(+), 287 deletions(-) create mode 100644 devito/petsc/types.py diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index f0b0d718ea..5aa71a0529 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -376,7 +376,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 67ed9269a4..651aa7d1d3 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -239,7 +239,7 @@ def __new__(cls, *args, **kwargs): expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + return expr @property diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 9d2db185db..af48dc73c7 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -24,6 +24,7 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: + # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index fa7716aaf8..f130334e95 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -29,7 +29,11 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', +<<<<<<< HEAD 'CallableBody', 'Transfer', 'Callback', 'FixedArgsCallable'] +======= + 'CallableBody', 'Transfer', 'Callback'] +>>>>>>> f921beaef (misc: Update file due to incorrect rebase) # First-class IET nodes diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index a85b93460d..c738db645a 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,6 +163,7 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: + # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index fe4b5ab48c..a0c0981d68 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -428,6 +428,7 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters + # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 59de8f95c7..98fb8cabfe 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -456,7 +456,9 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [i for i in bases if i.name != i.function._C_name] + bases = [ + i for i in bases + if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 7d7c156dcb..561f9fd0f4 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,16 +1,13 @@ import cgen as c from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (Transformer, MapNodes, Iteration, List, BlankLine, - DummyExpr, FindNodes, retrieve_iteration_tree, - filter_iterations) -from devito.symbolics import Byref, Macro, FieldFromComposite -from devito.petsc.types import (PetscMPIInt, DM, Mat, LocalVec, GlobalVec, - KSP, PC, SNES, PetscErrorCode, DummyArg, PetscInt, - StartPtr) -from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall -from devito.petsc.utils import solver_mapper, core_metadata -from devito.petsc.iet.routines import PETScCallbackBuilder +from devito.ir.iet import Transformer, MapNodes, Iteration, BlankLine +from devito.symbolics import Byref, Macro +from devito.petsc.types import (PetscMPIInt, PetscErrorCode) +from devito.petsc.iet.nodes import InjectSolveDummy +from devito.petsc.utils import core_metadata +from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup, + Solver, TimeDependent, NonTimeDependent) from devito.petsc.iet.utils import petsc_call, petsc_call_mpi @@ -34,54 +31,36 @@ def lower_petsc(iet, **kwargs): setup = [] subs = {} - - # Create a different DMDA for each target with a unique space order - unique_dmdas = create_dmda_objs(targets) - objs.update(unique_dmdas) - for dmda in unique_dmdas.values(): - setup.extend(create_dmda_calls(dmda, objs)) - - builder = PETScCallbackBuilder(**kwargs) + efuncs = {} for iters, (injectsolve,) in injectsolve_mapper.items(): - solver_objs = build_solver_objs(injectsolve, iters, **kwargs) - # Generate the solver setup for each InjectSolveDummy - solver_setup = generate_solver_setup(solver_objs, objs, injectsolve) - setup.extend(solver_setup) + builder = Builder(injectsolve, objs, iters, **kwargs) - # Generate all PETSc callback functions for the target via recursive compilation - matvec_op, formfunc_op, runsolve = builder.make(injectsolve, - objs, solver_objs) - setup.extend([matvec_op, formfunc_op, BlankLine]) - # Only Transform the spatial iteration loop - space_iter, = spatial_injectsolve_iter(iters, injectsolve) - subs.update({space_iter: List(body=runsolve)}) + setup.extend(builder.solversetup.calls) - # Generate callback to populate main struct object - struct, struct_calls = builder.make_main_struct(unique_dmdas, objs) - setup.extend(struct_calls) + # Transform the spatial iteration loop with the calls to execute the solver + subs.update(builder.solve.mapper) - iet = Transformer(subs).visit(iet) + efuncs.update(builder.cbbuilder.efuncs) - iet = assign_time_iters(iet, struct) + iet = Transformer(subs).visit(iet) body = core + tuple(setup) + (BlankLine,) + iet.body.body body = iet.body._rebuild( init=init, body=body, - frees=(c.Line("PetscCall(PetscFinalize());"),) + frees=(petsc_call('PetscFinalize', []),) ) iet = iet._rebuild(body=body) metadata = core_metadata() - efuncs = tuple(builder.efuncs.values()) - metadata.update({'efuncs': efuncs}) + metadata.update({'efuncs': tuple(efuncs.values())}) return iet, metadata def init_petsc(**kwargs): # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve + # specified via the parameters dict in PETScSolve # TODO: Are users going to be able to use PETSc command line arguments? # In firedrake, they have an options_prefix for each solver, enabling the use # of command line options @@ -110,201 +89,48 @@ def build_core_objects(target, **kwargs): } -def create_dmda_objs(unique_targets): - unique_dmdas = {} - for target in unique_targets: - name = 'da_so_%s' % target.space_order - unique_dmdas[name] = DM(name=name, liveness='eager', - stencil_width=target.space_order) - return unique_dmdas - - -def create_dmda_calls(dmda, objs): - dmda_create = create_dmda(dmda, objs) - dm_setup = petsc_call('DMSetUp', [dmda]) - dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) - dm_get_local_info = petsc_call('DMDAGetLocalInfo', [dmda, Byref(dmda.info)]) - return dmda_create, dm_setup, dm_mat_type, dm_get_local_info, BlankLine - - -def create_dmda(dmda, objs): - no_of_space_dims = len(objs['grid'].dimensions) - - # MPI communicator - args = [objs['comm']] - - # Type of ghost nodes - args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(no_of_space_dims)]) - - # Stencil type - if no_of_space_dims > 1: - args.append('DMDA_STENCIL_BOX') - - # Global dimensions - args.extend(list(objs['grid'].shape)[::-1]) - # No.of processors in each dimension - if no_of_space_dims > 1: - args.extend(list(objs['grid'].distributor.topology)[::-1]) - - # Number of degrees of freedom per node - args.append(1) - # "Stencil width" -> size of overlap - args.append(dmda.stencil_width) - args.extend([Null for _ in range(no_of_space_dims)]) - - # The distributed array object - args.append(Byref(dmda)) - - # The PETSc call used to create the DMDA - dmda = petsc_call('DMDACreate%sd' % no_of_space_dims, args) - - return dmda - - -def build_solver_objs(injectsolve, iters, **kwargs): - target = injectsolve.expr.rhs.target - sreg = kwargs['sregistry'] - return { - 'Jac': Mat(sreg.make_name(prefix='J_')), - 'x_global': GlobalVec(sreg.make_name(prefix='x_global_')), - 'x_local': LocalVec(sreg.make_name(prefix='x_local_'), liveness='eager'), - 'b_global': GlobalVec(sreg.make_name(prefix='b_global_')), - 'b_local': LocalVec(sreg.make_name(prefix='b_local_')), - 'ksp': KSP(sreg.make_name(prefix='ksp_')), - 'pc': PC(sreg.make_name(prefix='pc_')), - 'snes': SNES(sreg.make_name(prefix='snes_')), - 'X_global': GlobalVec(sreg.make_name(prefix='X_global_')), - 'Y_global': GlobalVec(sreg.make_name(prefix='Y_global_')), - 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), - 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), - 'dummy': DummyArg(sreg.make_name(prefix='dummy_')), - 'localsize': PetscInt(sreg.make_name(prefix='localsize_')), - 'start_ptr': StartPtr(sreg.make_name(prefix='start_ptr_'), target.dtype), - 'true_dims': retrieve_time_dims(iters), - 'target': target, - 'time_mapper': injectsolve.expr.rhs.time_mapper, - } - - -def generate_solver_setup(solver_objs, objs, injectsolve): - target = solver_objs['target'] - - dmda = objs['da_so_%s' % target.space_order] - - solver_params = injectsolve.expr.rhs.solver_parameters - - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) - - snes_set_dm = petsc_call('SNESSetDM', [solver_objs['snes'], dmda]) - - create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(solver_objs['Jac'])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) - - snes_set_jac = petsc_call( - 'SNESSetJacobian', [solver_objs['snes'], solver_objs['Jac'], - solver_objs['Jac'], 'MatMFFDComputeJacobian', Null] - ) - - global_x = petsc_call('DMCreateGlobalVector', - [dmda, Byref(solver_objs['x_global'])]) - - global_b = petsc_call('DMCreateGlobalVector', - [dmda, Byref(solver_objs['b_global'])]) - - local_b = petsc_call('DMCreateLocalVector', - [dmda, Byref(solver_objs['b_local'])]) +class Builder: + """ + This class is designed to support future extensions, enabling + different combinations of solver types, preconditioning methods, + and other functionalities as needed. - snes_get_ksp = petsc_call('SNESGetKSP', - [solver_objs['snes'], Byref(solver_objs['ksp'])]) + The class will be extended to accommodate different solver types by + returning subclasses of the objects initialised in __init__, + depending on the properties of `injectsolve`. + """ + def __init__(self, injectsolve, objs, iters, **kwargs): - ksp_set_tols = petsc_call( - 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], - solver_params['ksp_atol'], solver_params['ksp_divtol'], - solver_params['ksp_max_it']] - ) + # Determine the time dependency class + time_mapper = injectsolve.expr.rhs.time_mapper + timedep = TimeDependent if time_mapper else NonTimeDependent + self.timedep = timedep(injectsolve, iters, **kwargs) - ksp_set_type = petsc_call( - 'KSPSetType', [solver_objs['ksp'], solver_mapper[solver_params['ksp_type']]] - ) + # Objects + self.objbuilder = BaseObjectBuilder(injectsolve, **kwargs) + self.solver_objs = self.objbuilder.solver_objs - ksp_get_pc = petsc_call('KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])]) - - # Even though the default will be jacobi, set to PCNONE for now - pc_set_type = petsc_call('PCSetType', [solver_objs['pc'], 'PCNONE']) - - ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) - - return ( - snes_create, - snes_set_dm, - create_matrix, - snes_set_jac, - snes_set_type, - global_x, - global_b, - local_b, - snes_get_ksp, - ksp_set_tols, - ksp_set_type, - ksp_get_pc, - pc_set_type, - ksp_set_from_ops - ) + # Callbacks + self.cbbuilder = CallbackBuilder( + injectsolve, objs, self.solver_objs, timedep=self.timedep, + **kwargs + ) + # Solver setup + self.solversetup = BaseSetup( + self.solver_objs, objs, injectsolve, self.cbbuilder + ) -def assign_time_iters(iet, struct): - """ - Assign time iterators to the struct within loops containing PETScCalls. - Ensure that assignment occurs only once per time loop, if necessary. - Assign only the iterators that are common between the struct fields - and the actual Iteration. - """ - time_iters = [ - i for i in FindNodes(Iteration).visit(iet) - if i.dim.is_Time and FindNodes(PETScCall).visit(i) - ] - - if not time_iters: - return iet - - mapper = {} - for iter in time_iters: - common_dims = [d for d in iter.dimensions if d in struct.fields] - common_dims = [ - DummyExpr(FieldFromComposite(d, struct), d) for d in common_dims - ] - iter_new = iter._rebuild(nodes=List(body=tuple(common_dims)+iter.nodes)) - mapper.update({iter: iter_new}) - - return Transformer(mapper).visit(iet) - - -def retrieve_time_dims(iters): - time_iter = [i for i in iters if any(d.is_Time for d in i.dimensions)] - mapper = {} - if not time_iter: - return mapper - for d in time_iter[0].dimensions: - if d.is_Modulo: - mapper[d.origin] = d - elif d.is_Time: - mapper[d] = d - return mapper - - -def spatial_injectsolve_iter(iter, injectsolve): - spatial_body = [] - for tree in retrieve_iteration_tree(iter[0]): - root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if injectsolve in FindNodes(InjectSolveDummy).visit(root): - spatial_body.append(root) - return spatial_body + # Execute the solver + self.solve = Solver( + self.solver_objs, objs, injectsolve, iters, + self.cbbuilder, timedep=self.timedep + ) Null = Macro('NULL') void = 'void' + # TODO: Don't use c.Line here? petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index da4941f799..4f16ded1f3 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -9,23 +9,22 @@ from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions from devito.tools import as_tuple -from devito.petsc.types import LinearSolveExpr, PETScArray +from devito.petsc.types import LinearSolveExpr, PETScArray, DMDALocalInfo __all__ = ['PETScSolve'] def PETScSolve(eqns, target, solver_parameters=None, **kwargs): - prefixes = ['y_matvec', 'x_matvec', 'y_formfunc', 'x_formfunc', 'b_tmp'] + prefixes = ['y_matvec', 'x_matvec', 'f_formfunc', 'x_formfunc', 'b_tmp'] + + localinfo = DMDALocalInfo(name='info', liveness='eager') arrays = { p: PETScArray(name='%s_%s' % (p, target.name), - dtype=target.dtype, - dimensions=target.space_dimensions, - shape=target.grid.shape, - liveness='eager', - halo=[target.halo[d] for d in target.space_dimensions], - space_order=target.space_order) + target=target, + liveness='eager', + localinfo=localinfo) for p in prefixes } @@ -47,7 +46,7 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): )) formfuncs.append(Eq( - arrays['y_formfunc'], + arrays['f_formfunc'], F_target.subs(targets_to_arrays(arrays['x_formfunc'], targets)), subdomain=eq.subdomain )) @@ -60,8 +59,9 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): funcs = retrieve_functions(eqns) time_mapper = generate_time_mapper(funcs) + matvecs, formfuncs, formrhs = ( - [eq.subs(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs) + [eq.xreplace(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs) ) # Placeholder equation for inserting calls to the solver and generating # correct time loop etc @@ -74,6 +74,7 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): formrhs=formrhs, arrays=arrays, time_mapper=time_mapper, + localinfo=localinfo )) return [inject_solve] @@ -211,12 +212,24 @@ def generate_time_mapper(funcs): Replace time indices with `Symbols` in equations used within PETSc callback functions. These symbols are Uxreplaced at the IET level to align with the `TimeDimension` and `ModuloDimension` objects - present in the inital lowering. + present in the initial lowering. NOTE: All functions used in PETSc callback functions are attached to the `LinearSolveExpr` object, which is passed through the initial lowering (and subsequently dropped and replaced with calls to run the solver). Therefore, the appropriate time loop will always be correctly generated inside the main kernel. + + Examples + -------- + >>> funcs = [ + >>> f1(t + dt, x, y), + >>> g1(t + dt, x, y), + >>> g2(t, x, y), + >>> f1(t, x, y) + >>> ] + >>> generate_time_mapper(funcs) + {t + dt: tau0, t: tau1} + """ time_indices = list({ i if isinstance(d, SteppingDimension) else d diff --git a/devito/petsc/types.py b/devito/petsc/types.py new file mode 100644 index 0000000000..62c69f666a --- /dev/null +++ b/devito/petsc/types.py @@ -0,0 +1,384 @@ +import sympy +from functools import cached_property +import numpy as np + +from devito.tools import CustomDtype +from devito.types import LocalObject, Eq, CompositeObject +from devito.types.utils import DimensionTuple +from devito.types.array import ArrayBasic +from devito.finite_differences import Differentiable +from devito.types.basic import AbstractFunction, Symbol +from devito.finite_differences.tools import fd_weights_registry +from devito.tools import Reconstructable, dtype_to_ctype +from devito.symbolics import FieldFromComposite +from devito.types.basic import IndexedBase + + +class DM(LocalObject): + """ + PETSc Data Management object (DM). + """ + dtype = CustomDtype('DM') + + def __init__(self, *args, stencil_width=None, **kwargs): + super().__init__(*args, **kwargs) + self._stencil_width = stencil_width + + @property + def stencil_width(self): + return self._stencil_width + + +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 SNES(LocalObject): + """ + PETSc SNES : Non-Linear Systems Solvers. + """ + dtype = CustomDtype('SNES') + + +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 DMDALocalInfo(LocalObject): + """ + PETSc object - C struct containing information + about the local grid. + """ + dtype = CustomDtype('DMDALocalInfo') + + +class PetscErrorCode(LocalObject): + """ + PETSc datatype used to return PETSc error codes. + https://petsc.org/release/manualpages/Sys/PetscErrorCode/ + """ + dtype = CustomDtype('PetscErrorCode') + + +class PETScArray(ArrayBasic, Differentiable): + """ + 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. + Differentiable enables compatability with standard Function objects, + allowing for the use of the `subs` method. + TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) + and then PETScArray(ArrayBasic). + """ + + _data_alignment = False + + # Default method for the finite difference approximation weights computation. + _default_fd = 'taylor' + + __rkwargs__ = (AbstractFunction.__rkwargs__ + + ('dimensions', 'shape', 'liveness', 'coefficients')) + + def __init_finalize__(self, *args, **kwargs): + + super().__init_finalize__(*args, **kwargs) + + # Symbolic (finite difference) coefficients + self._coefficients = kwargs.get('coefficients', self._default_fd) + if self._coefficients not in fd_weights_registry: + raise ValueError("coefficients must be one of %s" + " not %s" % (str(fd_weights_registry), self._coefficients)) + self._shape = kwargs.get('shape') + + @classmethod + def __dtype_setup__(cls, **kwargs): + return kwargs.get('dtype', np.float32) + + @property + def coefficients(self): + """Form of the coefficients of the function.""" + return self._coefficients + + @property + def shape(self): + return self._shape + + @cached_property + def _shape_with_inhalo(self): + """ + Shape of the domain+inhalo region. The inhalo region comprises the + outhalo as well as any additional "ghost" layers for MPI halo + exchanges. Data in the inhalo region are exchanged when running + Operators to maintain consistent values as in sequential runs. + + Notes + ----- + Typically, this property won't be used in user code, but it may come + in handy for testing or debugging + """ + return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) + + @cached_property + def shape_allocated(self): + """ + Shape of the allocated data of the Function type object from which + this PETScArray was derived. It includes the domain and inhalo regions, + as well as any additional padding surrounding the halo. + + Notes + ----- + In an MPI context, this is the *local* with_halo region shape. + """ + return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, + self._padding)], + getters=self.dimensions) + + @cached_property + def _C_ctype(self): + return CustomDtype(self.petsc_type, modifier='*') + + @property + def petsc_type(self): + return dtype_to_petsctype(self._dtype) + + @property + def dtype(self): + return CustomDtype(self.petsc_type) + + @property + def symbolic_shape(self): + field_from_composites = [ + FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] + # Reverse it since DMDA is setup backwards to Devito dimensions. + return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + + @cached_property + def dmda_info(self): + # To access the local grid info via the DM in + # PETSc you use DMDAGetLocalInfo(da, &info) + # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz + return DMDALocalInfo(name='info', liveness='eager') + + @cached_property + def indexed(self): + return PETScIndexedData(self.name, shape=self._shape, function=self.function) + + +class PETScIndexedData(IndexedBase): + pass + + +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] + + +class MatVecEq(Eq): + """ + Represents the mathematical expression of applying a linear + operator to a vector. This is a key component + for running matrix-free solvers. + """ + pass + + +class RHSEq(Eq): + """ + Represents the mathematical expression of building the + rhs of a linear system. + """ + pass + + +class MockEq(Eq): + """ + Represents a mock/placeholder equation to ensure distinct iteration loops. + + For example, the mat-vec action iteration loop is to be isolated from the + expression loop used to build the RHS of the linear system. This separation + facilitates the utilisation of the mat-vec iteration loop in callback functions + created at the IET level. + """ + pass + + +def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): + # TODO: Add check for time dimensions and utilise implicit dimensions. + + is_time_dep = any(dim.is_Time for dim in target.dimensions) + # TODO: Current assumption is rhs is part of pde that remains + # constant at each timestep. Need to insert function to extract this from eq. + y_matvec, x_matvec, b_tmp = [ + PETScArray(name=f'{prefix}_{target.name}', + dtype=target.dtype, + dimensions=target.space_dimensions, + shape=target.grid.shape, + liveness='eager', + halo=target.halo[1:] if is_time_dep else target.halo) + for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] + + # TODO: Extend to rearrange equation for implicit time stepping. + matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), + target=target, solver_parameters=solver_parameters), + subdomain=eq.subdomain) + + # Part of pde that remains constant at each timestep + rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, + solver_parameters=solver_parameters), subdomain=eq.subdomain) + + if not bcs: + return [matvecaction, rhs] + + bcs_for_matvec = [] + for bc in bcs: + # TODO: Insert code to distiguish between essential and natural + # boundary conditions since these are treated differently within + # the solver + # NOTE: May eventually remove the essential bcs from the solve + # (and move to rhs) but for now, they are included since this + # is not trivial to implement when using DMDA + # NOTE: Below is temporary -> Just using this as a palceholder for + # the actual BC implementation for the matvec callback + new_rhs = bc.rhs.subs(target, x_matvec) + bc_rhs = LinearSolveExpr( + new_rhs, target=target, solver_parameters=solver_parameters + ) + bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) + + return [matvecaction] + bcs_for_matvec + [rhs] + + +class LinearSolveExpr(sympy.Function, Reconstructable): + + __rargs__ = ('expr',) + __rkwargs__ = ('target', 'solver_parameters') + + defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'jacobi', + 'ksp_rtol': 1e-7, # Relative tolerance + 'ksp_atol': 1e-50, # Absolute tolerance + 'ksp_divtol': 1e4, # Divergence tolerance + 'ksp_max_it': 10000 # Maximum iterations + } + + def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): + + if solver_parameters is None: + solver_parameters = cls.defaults + else: + for key, val in cls.defaults.items(): + solver_parameters[key] = solver_parameters.get(key, val) + + obj = super().__new__(cls, expr) + obj._expr = expr + obj._target = target + obj._solver_parameters = solver_parameters + return obj + + def __repr__(self): + return "%s(%s)" % (self.__class__.__name__, self.expr) + + __str__ = __repr__ + + def _sympystr(self, printer): + return str(self) + + def __hash__(self): + return hash(self.target) + + @property + def expr(self): + return self._expr + + @property + def target(self): + return self._target + + @property + def solver_parameters(self): + return self._solver_parameters + + func = Reconstructable._rebuild + + +def petsc_lift(clusters): + """ + Lift the iteration space associated with each PETSc equation. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions. + """ + processed = [] + for c in clusters: + + ispace = c.ispace + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + + processed.append(c.rebuild(ispace=ispace)) + + return processed + + +class PETScStruct(CompositeObject): + + __rargs__ = ('name', 'usr_ctx') + + def __init__(self, name, usr_ctx): + pfields = [(i._C_name, dtype_to_ctype(i.dtype)) + for i in usr_ctx if isinstance(i, Symbol)] + self._usr_ctx = usr_ctx + super().__init__(name, 'MatContext', pfields) + + @property + def usr_ctx(self): + return self._usr_ctx + + def _arg_values(self, **kwargs): + values = super()._arg_values(**kwargs) + for i in self.fields: + setattr(values[self.name]._obj, i, kwargs['args'][i]) + return values diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index d898db23cb..ccc47a4c3e 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -51,3 +51,16 @@ def core_metadata(): 'lib_dirs': lib_dir, 'ldflags': ('-Wl,-rpath,%s' % lib_dir) } +<<<<<<< HEAD +======= + + +def petsc_call(specific_call, call_args): + general_call = 'PetscCall' + return Call(general_call, [Call(specific_call, arguments=call_args)]) + + +def petsc_call_mpi(specific_call, call_args): + general_call = 'PetscCallMPI' + return Call(general_call, [Call(specific_call, arguments=call_args)]) +>>>>>>> f921beaef (misc: Update file due to incorrect rebase) diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index 2c366a389e..e7b395d436 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,6 +242,9 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) + def _print_PETScRHS(self, expr): + return "%s" % self._print(expr._expr) + _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/docker/Dockerfile.intel b/docker/Dockerfile.intel index 1fab7147d6..e013fccc03 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -35,7 +35,6 @@ RUN apt-get update -y && \ apt-get install -y intel-oneapi-advisor # Drivers mandatory for intel gpu -<<<<<<< HEAD # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" > /etc/apt/sources.list.d/intel-gpu-jammy.list @@ -51,16 +50,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ pciutils libkmod-dev intel-fw-gpu xpu-smi \ # Development packages libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev -======= -# https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-focal.html#ubuntu-20-04-focal -RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg -RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu focal main" > /etc/apt/sources.list.d/intel.list - -RUN apt-get update -y && apt-get dist-upgrade -y && \ - apt-get install -y intel-opencl-icd intel-level-zero-gpu level-zero level-zero-dev \ - intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \ - libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh intel64 && ' ENV MPI4PY_RC_RECV_MPROBE=0 @@ -87,11 +76,7 @@ ENV DEVITO_PLATFORM="intel64" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD # ICX OpenMP image -======= -# ICX image ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) ############################################################## FROM oneapi AS icx @@ -108,35 +93,6 @@ ENV DEVITO_LANGUAGE="openmp" ENV MPICC=mpiicc ############################################################## -<<<<<<< HEAD -======= -# ICX hpc image -############################################################## -FROM oneapi as icx-hpc - -# Install both icc and icx to avoid missing dependencies -RUN apt-get update -y && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp intel-oneapi-mpi-devel && \ - apt-get install -y intel-oneapi-compiler-dpcpp-cpp-and-cpp-classic - -# Missig components -# https://www.intel.com/content/www/us/en/developer/tools/oneapi/hpc-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager -RUN curl -f "https://registrationcenter-download.intel.com/akdlm/IRC_NAS/ebf5d9aa-17a7-46a4-b5df-ace004227c0e/l_dpcpp-cpp-compiler_p_2023.2.1.8.sh" -O && \ - chmod +x l_dpcpp-cpp-compiler_p_2023.2.1.8.sh && ./l_dpcpp-cpp-compiler_p_2023.2.1.8.sh -a -s --eula accept && \ - rm l_dpcpp-cpp-compiler_p_2023.2.1.8.sh - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -# Devito config -ENV DEVITO_ARCH="icx" -ENV DEVITO_LANGUAGE="openmp" -# MPICC compiler for mpi4py -ENV MPICC=mpiicc -ENV MPI4PY_FLAGS='. /opt/intel/oneapi/setvars.sh && CFLAGS="-cc=icx"' - -############################################################## ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # ICX SYCL CPU image ############################################################## FROM icx AS cpu-sycl @@ -151,12 +107,6 @@ ENV DEVITO_PLATFORM="intel64" ############################################################## FROM icx as gpu-sycl -<<<<<<< HEAD -======= -# NOTE: the name of this file ends with ".cpu" but this is a GPU image. -# It then feels a bit akward, so some restructuring might be needed - ->>>>>>> c4373f90e (misc: Split Dockerfile.cpu into .cpu and .intel) # Devito config ENV DEVITO_ARCH="sycl" ENV DEVITO_LANGUAGE="sycl" From 0f862ba0fcdc851b38533dbfeafa7a8b3d9040bf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 10 Jul 2024 11:31:30 +0100 Subject: [PATCH 123/151] dsl: Dispatch to sympy.Add not both Add and EvalDerivative --- .github/workflows/examples-mpi.yml | 1 + .github/workflows/examples.yml | 1 + .github/workflows/flake8.yml | 1 + .github/workflows/pytest-core-mpi.yml | 1 + .github/workflows/pytest-core-nompi.yml | 1 + .github/workflows/tutorials.yml | 1 + devito/ir/clusters/cluster.py | 2 +- devito/ir/equations/equation.py | 6 +- devito/ir/stree/algorithms.py | 1 - devito/operator/operator.py | 8 +- devito/passes/clusters/petsc.py | 25 + devito/passes/iet/__init__.py | 1 + devito/passes/iet/petsc.py | 298 +++++++++ devito/petsc/iet.py | 372 +++++++++++ devito/petsc/iet/routines.py | 831 +++++++++++++++++++----- devito/petsc/types.py | 303 +++++---- devito/petsc/utils.py | 13 - docker/Dockerfile.petsc | 49 ++ 18 files changed, 1590 insertions(+), 325 deletions(-) create mode 100644 devito/passes/clusters/petsc.py create mode 100644 devito/passes/iet/petsc.py create mode 100644 devito/petsc/iet.py create mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 59cd08eaa5..5553e9108d 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,6 +17,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 94335a9e2d..2a94b27393 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 336cedb294..9a0aea32a8 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 67f841cabd..2b36602110 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index df3a050927..b818f2b8f6 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index c792dd6acd..42a2288c56 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,6 +10,7 @@ on: push: branches: - master + - enable-opt pull_request: branches: - master diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 5aa71a0529..f0b0d718ea 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -376,7 +376,7 @@ def dspace(self): d = i.dim try: if i.lower < 0 or \ - i.upper > f._size_nodomain[d].left + f._size_halo[d].right: + i.upper > f._size_nodomain[d].left + f._size_halo[d].right: # It'd mean trying to access a point before the # left halo (test0) or after the right halo (test1) oobs.update(d._defines) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 651aa7d1d3..0d6c7b0ff7 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -227,19 +227,19 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - + # from IPython import embed; embed() # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - + # from IPython import embed; embed() expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) expr._implicit_dims = input_expr.implicit_dims expr._operation = Operation.detect(input_expr) - + return expr @property diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index c738db645a..a85b93460d 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -163,7 +163,6 @@ def preprocess(clusters, options=None, **kwargs): queue = [] processed = [] for c in clusters: - # from IPython import embed; embed() if c.is_halo_touch: hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs) queue.append(c.rebuild(exprs=[], halo_scheme=hs)) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index a0c0981d68..6a4f66f28d 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -354,7 +354,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - + # from IPython import embed; embed() return processed # Compilation -- Cluster level @@ -498,6 +498,9 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Target-independent optimizations minimize_symbols(graph) + # If necessary, sort frees into a specific order + sort_frees(graph) + return graph.root, graph # Read-only properties exposed to the outside world @@ -524,7 +527,8 @@ def dimensions(self): @cached_property def input(self): - return tuple(i for i in self.parameters if i.is_Input) + struct_params = derive_struct_inputs(self.parameters) + return tuple(i for i in self.parameters+struct_params if i.is_Input) @cached_property def temporaries(self): diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py new file mode 100644 index 0000000000..d784d99c0d --- /dev/null +++ b/devito/passes/clusters/petsc.py @@ -0,0 +1,25 @@ +from devito.tools import timed_pass +from devito.petsc import LinearSolveExpr + +__all__ = ['petsc_lift'] + + +@timed_pass() +def petsc_lift(clusters): + """ + Lift the iteration space surrounding each PETSc equation to create + distinct iteration loops. This simplifys the movement of the loops + into specific callback functions generated at the IET level. + TODO: Potentially only need to lift the PETSc equations required + by the callback functions, not the ones that stay inside the main kernel. + """ + processed = [] + for c in clusters: + + if isinstance(c.exprs[0].rhs, LinearSolveExpr): + ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index c09db00c9b..cf6a35de90 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,3 +8,4 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa +from .petsc import * # noqa diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py new file mode 100644 index 0000000000..8c8aeafa1c --- /dev/null +++ b/devito/passes/iet/petsc.py @@ -0,0 +1,298 @@ +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, MatVecAction, + Transformer, FindSymbols, LinearSolverExpression, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition) +from devito.types import (PetscMPIInt, PETScStruct, DMDALocalInfo, DM, Mat, + Vec, KSP, PC, SNES, PetscErrorCode) +from devito.symbolics import Byref, Macro, FieldFromPointer +import cgen as c + +__all__ = ['lower_petsc'] + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't + # appear in ccode + + # Check if PETScSolve was used and count occurrences. Each PETScSolve + # will have a unique MatVecAction. + is_petsc = FindNodes(MatVecAction).visit(iet) + + if is_petsc: + + # Collect all solution fields we're solving for + targets = [i.expr.rhs.target for i in is_petsc] + + # Initalize PETSc + init = init_petsc(**kwargs) + + # Create context data struct. + struct = build_struct(iet) + + objs = build_core_objects(targets[-1], struct, **kwargs) + + # Create core PETSc calls required by general linear solves, + # which only need to be generated once e.g create DMDA. + core = core_petsc(targets[-1], objs, **kwargs) + + matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) + + main_mapper = {} + + setup = [] + efuncs = [] + for target in unique_targets: + + solver_objs = build_solver_objs(target) + + matvec_callback_body_iters = [] + + solver_setup = False + + for iter, (matvec,) in matvec_mapper.items(): + + if matvec.expr.rhs.target == target: + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec) + setup.extend(solver) + solver_setup = True + + # Only need to generate solver setup once per target + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + solver_setup = True + + matvec_body = matvec_body_list._rebuild(body=[ + matvec_body_list.body, iter[0]]) + matvec_body_list = matvec_body_list._rebuild(body=matvec_body) + + main_mapper.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, matvec_callback_body_iters, solver_objs, objs, struct) + + setup.append(matvec_op) + setup.append(c.Line()) + + efuncs.append(matvec_callback) + + # Remove the LinSolveExpr from iet and efuncs that were used to carry + # metadata e.g solver_parameters + main_mapper.update(rebuild_expr_mapper(iet)) + efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} + + iet = Transformer(main_mapper).visit(iet) + efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] + + # Replace symbols appearing in each efunc with a pointer to the PETScStruct + efuncs = transform_efuncs(efuncs, struct) + + body = iet.body._rebuild(body=(tuple(init_setup) + iet.body.body)) + iet = iet._rebuild(body=body) + + return iet, {} + + +def init_petsc(**kwargs): + + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve. + # NOTE: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options. + initialize = Call('PetscCall', [Call('PetscInitialize', + arguments=['NULL', 'NULL', + 'NULL', 'NULL'])]) + + return tuple([initialize]) + + +def build_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct. + usr_ctx = [] + + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def core_petsc(target, objs, **kwargs): + # Assumption: all targets are generated from the same Grid, + # so we can use any target. + + # MPI + call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', + arguments=[objs['comm'], + Byref(objs['size'])])]) + # Create DMDA + dmda = create_dmda(target, objs) + dm_setup = Call('PetscCall', [Call('DMSetUp', arguments=[objs['da']])]) + dm_app_ctx = Call('PetscCall', [Call('DMSetApplicationContext', + arguments=[objs['da'], objs['struct']])]) + dm_mat_type = Call('PetscCall', [Call('DMSetMatType', + arguments=[objs['da'], 'MATSHELL'])]) + dm_local_info = Call('PetscCall', [Call('DMDAGetLocalInfo', + arguments=[objs['da'], Byref(objs['info'])])]) + + return tuple([call_mpi, dmda, dm_setup, dm_app_ctx, dm_mat_type, dm_local_info]) + + +def build_core_objects(target, struct, **kwargs): + + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return {'da': DM(name='da'), + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'struct': struct} + + +def create_dmda(target, objs): + + args = [objs['comm']] + + args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] + + # stencil type + args += ['DMDA_STENCIL_BOX'] + + # global dimensions + args += list(target.shape_global)[::-1] + + # no.of processors in each dimension + args += list(target.grid.distributor.topology)[::-1] + + args += [1, target.space_order] + + args += ['NULL' for _ in range(len(target.space_dimensions))] + + args += [Byref(objs['da'])] + + dmda = Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args) + + return dmda + + +def build_solver_objs(target): + + return {'Jac': Mat(name='J_'+str(target.name)), + 'x_global': Vec(name='x_global_'+str(target.name)), + 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), + 'b_global': Vec(name='b_global_'+str(target.name)), + 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), + 'ksp': KSP(name='ksp_'+str(target.name)), + 'pc': PC(name='pc_'+str(target.name)), + 'snes': SNES(name='snes_'+str(target.name)), + 'x': Vec(name='x_'+str(target.name)), + 'y': Vec(name='y_'+str(target.name))} + + +def generate_solver_calls(solver_objs, objs, matvec): + + snes_create = Call('PetscCall', [Call('SNESCreate', arguments=[ + objs['comm'], Byref(solver_objs['snes'])])]) + + snes_set_dm = Call('PetscCall', [Call('SNESSetDM', arguments=[ + solver_objs['snes'], objs['da']])]) + + create_matrix = Call('PetscCall', [Call('DMCreateMatrix', arguments=[ + objs['da'], Byref(solver_objs['Jac'])])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = Call('PetscCall', [Call('SNESSetType', arguments=[ + solver_objs['snes'], 'SNESKSPONLY'])]) + + global_x = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['x_global'])])]) + + local_x = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['x_local'])])]) + + global_b = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['b_global'])])]) + + local_b = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['b_local'])])]) + + snes_get_ksp = Call('PetscCall', [Call('SNESGetKSP', arguments=[ + solver_objs['snes'], Byref(solver_objs['ksp'])])]) + + return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, + global_x, local_x, global_b, local_b, snes_get_ksp]) + + +def create_matvec_callback(target, matvec_callback_body_iters, + solver_objs, objs, struct): + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + defn_struct = Definition(struct) + + get_context = Call('PetscCall', [Call('MatShellGetContext', + arguments=[solver_objs['Jac'], + Byref(struct)])]) + + body = List(body=[defn_struct, + get_context, + matvec_callback_body_iters]) + + matvec_callback = Callable('MyMatShellMult_'+str(target.name), + matvec_body, + retval=objs['err'], + parameters=(solver_objs['Jac'], + solver_objs['x'], + solver_objs['y'])) + + matvec_operation = Call('PetscCall', [ + Call('MatShellSetOperation', arguments=[solver_objs['Jac'], + 'MATOP_MULT', + Callback(matvec_callback.name, + Void, Void)])]) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(iet): + + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for + expr in FindNodes(LinearSolverExpression).visit(iet)} + + +def transform_efuncs(efuncs, struct): + + efuncs_new = [] + for efunc in efuncs: + new_body = efunc.body + for i in struct.usr_ctx: + new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) + efunc_with_new_body = efunc._rebuild(body=new_body) + efuncs_new.append(efunc_with_new_body) + + return efuncs_new + + +Null = Macro('NULL') +Void = Macro('void') + +petsc_call = String('PetscCall') +petsc_call_mpi = String('PetscCallMPI') +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') + +linear_solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} diff --git a/devito/petsc/iet.py b/devito/petsc/iet.py new file mode 100644 index 0000000000..659ac9f94d --- /dev/null +++ b/devito/petsc/iet.py @@ -0,0 +1,372 @@ +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import (FindNodes, Call, + Transformer, FindSymbols, + MapNodes, Iteration, Callable, Callback, List, Uxreplace, + Definition, BlankLine, PointerCast) +<<<<<<< HEAD +from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, +======= +from devito.petsc import (PetscMPIInt, PETScStruct, DM, Mat, +>>>>>>> b22824875 (fix circular imports) + Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) +from devito.symbolics import Byref, Macro, FieldFromPointer +import cgen as c +from devito.petsc.nodes import MatVecAction, RHSLinearSystem, LinearSolverExpression + + +@iet_pass +def lower_petsc(iet, **kwargs): + + # Check if PETScSolve was used. + petsc_nodes = FindNodes(MatVecAction).visit(iet) + + if not petsc_nodes: + return iet, {} + + else: + # Collect all petsc solution fields + unique_targets = list(set([i.expr.rhs.target for i in petsc_nodes])) + + # Initalize PETSc + init = init_petsc(**kwargs) + + # Create context data struct + struct = build_struct(iet) + + objs = build_core_objects(unique_targets[-1], **kwargs) + + # Create core PETSc calls (not specific to each PETScSolve) + core = core_petsc(unique_targets[-1], struct, objs, **kwargs) + + matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) + + main_mapper = {} + + setup = [] + efuncs = [] + + for target in unique_targets: + + solver_objs = build_solver_objs(target) + + matvec_body_list = List() + + solver_setup = False + + for iter, (matvec,) in matvec_mapper.items(): + + # Skip the MatVecAction if it is not associated with the target + # There will be more than one MatVecAction associated with the target + # e.g interior matvec + BC matvecs + if matvec.expr.rhs.target != target: + continue + + # Only need to generate solver setup once per target + if not solver_setup: + solver = generate_solver_calls(solver_objs, objs, matvec, target) + setup.extend(solver) + solver_setup = True + + matvec_body = matvec_body_list._rebuild(body=[ + matvec_body_list.body, iter[0]]) + matvec_body_list = matvec_body_list._rebuild(body=matvec_body) + + main_mapper.update({iter[0]: None}) + + # Create the matvec callback and operation for each target + matvec_callback, matvec_op = create_matvec_callback( + target, matvec_body_list, solver_objs, objs, + struct) + + setup.append(matvec_op) + setup.append(BlankLine) + efuncs.append(matvec_callback) + + # Remove the LinSolveExpr from iet and efuncs that were used to carry + # metadata e.g solver_parameters + main_mapper.update(rebuild_expr_mapper(iet)) + efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} + + iet = Transformer(main_mapper).visit(iet) + efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] + + # Replace symbols appearing in each efunc with a pointer to the PETScStruct + efuncs = transform_efuncs(efuncs, struct) + + body = iet.body._rebuild(init=init, body=core + tuple(setup) + iet.body.body) + iet = iet._rebuild(body=body) + + return iet, {'efuncs': efuncs} + + +def init_petsc(**kwargs): + + # Initialize PETSc -> for now, assuming all solver options have to be + # specifed via the parameters dict in PETScSolve + # TODO: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options + initialize = Call(petsc_call, [ + Call('PetscInitialize', arguments=[Null, Null, Null, Null])]) + + return tuple([petsc_func_begin_user, initialize]) + + +def build_struct(iet): + # Place all context data required by the shell routines + # into a PETScStruct + usr_ctx = [] + + basics = FindSymbols('basics').visit(iet) + avoid = FindSymbols('dimensions|indexedbases').visit(iet) + usr_ctx.extend(data for data in basics if data not in avoid) + + return PETScStruct('ctx', usr_ctx) + + +def core_petsc(target, struct, objs, **kwargs): + + # MPI + call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', + arguments=[objs['comm'], + Byref(objs['size'])])]) + # Create DMDA + dmda = create_dmda(target, objs) + dm_setup = Call(petsc_call, [ + Call('DMSetUp', arguments=[objs['da']])]) + dm_app_ctx = Call(petsc_call, [ + Call('DMSetApplicationContext', arguments=[objs['da'], struct])]) + dm_mat_type = Call(petsc_call, [ + Call('DMSetMatType', arguments=[objs['da'], 'MATSHELL'])]) + + return tuple([petsc_func_begin_user, call_mpi, dmda, dm_setup, + dm_app_ctx, dm_mat_type, BlankLine]) + + +def build_core_objects(target, **kwargs): + + if kwargs['options']['mpi']: + communicator = target.grid.distributor._obj_comm + else: + communicator = 'PETSC_COMM_SELF' + + return {'da': DM(name='da', liveness='eager'), + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'err': PetscErrorCode(name='err')} + + +def create_dmda(target, objs): + + args = [objs['comm']] + + args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] + + # Stencil type + if len(target.space_dimensions) > 1: + args += ['DMDA_STENCIL_BOX'] + + # Global dimensions + args += list(target.shape_global)[::-1] + + # No.of processors in each dimension + if len(target.space_dimensions) > 1: + args += list(target.grid.distributor.topology)[::-1] + + args += [1, target.space_order] + + args += [Null for _ in range(len(target.space_dimensions))] + + args += [Byref(objs['da'])] + + dmda = Call(petsc_call, [ + Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args)]) + + return dmda + + +def build_solver_objs(target): + + return {'Jac': Mat(name='J_'+str(target.name)), + 'x_global': Vec(name='x_global_'+str(target.name)), + 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), + 'b_global': Vec(name='b_global_'+str(target.name)), + 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), + 'ksp': KSP(name='ksp_'+str(target.name)), + 'pc': PC(name='pc_'+str(target.name)), + 'snes': SNES(name='snes_'+str(target.name)), + 'X_global': Vec(name='X_global_'+str(target.name)), + 'Y_global': Vec(name='Y_global_'+str(target.name)), + 'X_local': Vec(name='X_local_'+str(target.name), liveness='eager'), + 'Y_local': Vec(name='Y_local_'+str(target.name), liveness='eager') + } + + +def generate_solver_calls(solver_objs, objs, matvec, target): + + solver_params = matvec.expr.rhs.solver_parameters + + snes_create = Call(petsc_call, [Call('SNESCreate', arguments=[ + objs['comm'], Byref(solver_objs['snes'])])]) + + snes_set_dm = Call(petsc_call, [Call('SNESSetDM', arguments=[ + solver_objs['snes'], objs['da']])]) + + create_matrix = Call(petsc_call, [Call('DMCreateMatrix', arguments=[ + objs['da'], Byref(solver_objs['Jac'])])]) + + # NOTE: Assumming all solves are linear for now. + snes_set_type = Call(petsc_call, [Call('SNESSetType', arguments=[ + solver_objs['snes'], 'SNESKSPONLY'])]) + + global_x = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['x_global'])])]) + + local_x = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['x_local'])])]) + + global_b = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ + objs['da'], Byref(solver_objs['b_global'])])]) + + local_b = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ + objs['da'], Byref(solver_objs['b_local'])])]) + + snes_get_ksp = Call(petsc_call, [Call('SNESGetKSP', arguments=[ + solver_objs['snes'], Byref(solver_objs['ksp'])])]) + + vec_replace_array = Call(petsc_call, [ + Call('VecReplaceArray', arguments=[solver_objs[ + 'x_local'], FieldFromPointer(target._C_field_data, target._C_symbol)])]) + + ksp_set_tols = Call(petsc_call, [Call('KSPSetTolerances', arguments=[ + solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], + solver_params['ksp_divtol'], solver_params['ksp_max_it']])]) + + ksp_set_type = Call(petsc_call, [Call('KSPSetType', arguments=[ + solver_objs['ksp'], linear_solver_mapper[solver_params['ksp_type']]])]) + + ksp_get_pc = Call(petsc_call, [Call('KSPGetPC', arguments=[ + solver_objs['ksp'], Byref(solver_objs['pc'])])]) + + pc_set_type = Call(petsc_call, [Call('PCSetType', arguments=[ + solver_objs['pc'], linear_solver_mapper[solver_params['pc_type']]])]) + + ksp_set_from_ops = Call(petsc_call, [Call('KSPSetFromOptions', arguments=[ + solver_objs['ksp']])]) + + return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, + global_x, local_x, global_b, local_b, snes_get_ksp, + vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, + pc_set_type, ksp_set_from_ops]) + + +def create_matvec_callback(target, body, solver_objs, objs, struct): + + # There will be 2 PETScArrays within the body + petsc_arrays = [i.function for i in FindSymbols('indexedbases').visit(body) + if isinstance(i.function, PETScArray)] + + # There will only be one PETScArray that is written to within this body and + # one PETScArray which corresponds to the 'seed' vector + petsc_arr_write, = FindSymbols('writes').visit(body) + petsc_arr_seed, = [i for i in petsc_arrays if i.function != petsc_arr_write.function] + + # Struct needs to be defined explicitly here since CompositeObjects + # do not have 'liveness' + defn_struct = Definition(struct) + + mat_get_dm = Call(petsc_call, [Call('MatGetDM', arguments=[ + solver_objs['Jac'], Byref(objs['da'])])]) + + dm_get_app_context = Call(petsc_call, [Call('DMGetApplicationContext', arguments=[ + objs['da'], Byref(struct._C_symbol)])]) + + dm_get_local_xvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ + objs['da'], Byref(solver_objs['X_local'])])]) + + global_to_local_begin = Call(petsc_call, [Call('DMGlobalToLocalBegin', arguments=[ + objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) + + global_to_local_end = Call(petsc_call, [Call('DMGlobalToLocalEnd', arguments=[ + objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) + + dm_get_local_yvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ + objs['da'], Byref(solver_objs['Y_local'])])]) + + vec_get_array_y = Call(petsc_call, [Call('VecGetArray', arguments=[ + solver_objs['Y_local'], Byref(petsc_arr_write.function)])]) + + vec_get_array_x = Call(petsc_call, [Call('VecGetArray', arguments=[ + solver_objs['X_local'], Byref(petsc_arrays[0])])]) + + dm_get_local_info = Call(petsc_call, [Call('DMDAGetLocalInfo', arguments=[ + objs['da'], Byref(petsc_arrays[0].function.dmda_info)])]) + + casts = [PointerCast(i.function) for i in petsc_arrays] + + vec_restore_array_y = Call(petsc_call, [Call('VecRestoreArray', arguments=[ + solver_objs['Y_local'], Byref(petsc_arr_write[0])])]) + + vec_restore_array_x = Call(petsc_call, [Call('VecRestoreArray', arguments=[ + solver_objs['X_local'], Byref(petsc_arr_seed)])]) + + dm_local_to_global_begin = Call(petsc_call, [Call('DMLocalToGlobalBegin', arguments=[ + objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) + + dm_local_to_global_end = Call(petsc_call, [Call('DMLocalToGlobalEnd', arguments=[ + objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) + + func_return = Call('PetscFunctionReturn', arguments=[0]) + + matvec_body = List(body=[ + petsc_func_begin_user, defn_struct, mat_get_dm, dm_get_app_context, + dm_get_local_xvec, global_to_local_begin, global_to_local_end, + dm_get_local_yvec, vec_get_array_y, vec_get_array_x, dm_get_local_info, + casts, BlankLine, body, vec_restore_array_y, vec_restore_array_x, + dm_local_to_global_begin, dm_local_to_global_end, func_return]) + + matvec_callback = Callable( + 'MyMatShellMult_'+str(target.name), matvec_body, retval=objs['err'], + parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'])) + + matvec_operation = Call(petsc_call, [ + Call('MatShellSetOperation', arguments=[ + solver_objs['Jac'], 'MATOP_MULT', Callback(matvec_callback.name, + Void, Void)])]) + + return matvec_callback, matvec_operation + + +def rebuild_expr_mapper(callable): + + return {expr: expr._rebuild( + expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for + expr in FindNodes(LinearSolverExpression).visit(callable)} + + +def transform_efuncs(efuncs, struct): + + efuncs_new = [] + for efunc in efuncs: + new_body = efunc.body + for i in struct.usr_ctx: + new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) + efunc_with_new_body = efunc._rebuild(body=new_body) + efuncs_new.append(efunc_with_new_body) + + return efuncs_new + + +Null = Macro('NULL') +Void = 'void' + +petsc_call = 'PetscCall' +petsc_call_mpi = 'PetscCallMPI' +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') + +linear_solver_mapper = { + 'gmres': 'KSPGMRES', + 'jacobi': 'PCJACOBI', + None: 'PCNONE' +} diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 5ef8b2aa5f..b3dda6acc9 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -3,30 +3,49 @@ import cgen as c from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, - Dereference, DummyExpr, BlankLine, Callable) -from devito.symbolics import Byref, FieldFromPointer, Macro, cast_mapper + Dereference, DummyExpr, BlankLine, Callable, FindNodes, + retrieve_iteration_tree, filter_iterations) +from devito.symbolics import (Byref, FieldFromPointer, Macro, cast_mapper, + FieldFromComposite) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction -from devito.types import ModuloDimension, TimeDimension, Temp +from devito.types import Temp, Symbol from devito.tools import filter_ordered + from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, - MatVecCallback) + MatVecCallback, InjectSolveDummy) from devito.petsc.iet.utils import petsc_call, petsc_struct +from devito.petsc.utils import solver_mapper +from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC, + SNES, DummyArg, PetscInt, StartPtr) -class PETScCallbackBuilder: +class CallbackBuilder: """ Build IET routines to generate PETSc callback functions. """ - def __new__(cls, rcompile=None, sregistry=None, **kwargs): - obj = object.__new__(cls) - obj.rcompile = rcompile - obj.sregistry = sregistry - obj._efuncs = OrderedDict() - obj._struct_params = [] + def __init__(self, injectsolve, objs, solver_objs, + rcompile=None, sregistry=None, timedep=None, **kwargs): + + self.rcompile = rcompile + self.sregistry = sregistry + self.timedep = timedep + self.solver_objs = solver_objs + + self._efuncs = OrderedDict() + self._struct_params = [] + + self._matvec_callback = None + self._formfunc_callback = None + self._formrhs_callback = None + self._struct_callback = None - return obj + self._make_core(injectsolve, objs, solver_objs) + self._main_struct(solver_objs) + self._make_struct_callback(solver_objs, objs) + self._local_struct(solver_objs) + self._efuncs = self._uxreplace_efuncs() @property def efuncs(self): @@ -36,42 +55,38 @@ def efuncs(self): def struct_params(self): return self._struct_params - def make(self, injectsolve, objs, solver_objs): - matvec_callback, formfunc_callback, formrhs_callback = self.make_all( - injectsolve, objs, solver_objs - ) + @property + def filtered_struct_params(self): + return filter_ordered(self.struct_params) - matvec_operation = petsc_call( - 'MatShellSetOperation', [solver_objs['Jac'], 'MATOP_MULT', - MatVecCallback(matvec_callback.name, void, void)] - ) - formfunc_operation = petsc_call( - 'SNESSetFunction', - [solver_objs['snes'], Null, - FormFunctionCallback(formfunc_callback.name, void, void), Null] - ) - runsolve = self.runsolve(solver_objs, objs, formrhs_callback, injectsolve) + @property + def matvec_callback(self): + return self._matvec_callback - return matvec_operation, formfunc_operation, runsolve + @property + def formfunc_callback(self): + return self._formfunc_callback - def make_all(self, injectsolve, objs, solver_objs): - matvec_callback = self.make_matvec(injectsolve, objs, solver_objs) - formfunc_callback = self.make_formfunc(injectsolve, objs, solver_objs) - formrhs_callback = self.make_formrhs(injectsolve, objs, solver_objs) + @property + def formrhs_callback(self): + return self._formrhs_callback - self._efuncs[matvec_callback.name] = matvec_callback - self._efuncs[formfunc_callback.name] = formfunc_callback - self._efuncs[formrhs_callback.name] = formrhs_callback + @property + def struct_callback(self): + return self._struct_callback - return matvec_callback, formfunc_callback, formrhs_callback + def _make_core(self, injectsolve, objs, solver_objs): + self._make_matvec(injectsolve, objs, solver_objs) + self._make_formfunc(injectsolve, objs, solver_objs) + self._make_formrhs(injectsolve, objs, solver_objs) - def make_matvec(self, injectsolve, objs, solver_objs): + def _make_matvec(self, injectsolve, objs, solver_objs): # Compile matvec `eqns` into an IET via recursive compilation irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, options={'mpi': False}, sregistry=self.sregistry) - body_matvec = self.create_matvec_body(injectsolve, - List(body=irs_matvec.uiet.body), - solver_objs, objs) + body_matvec = self._create_matvec_body(injectsolve, + List(body=irs_matvec.uiet.body), + solver_objs, objs) matvec_callback = PETScCallable( self.sregistry.make_name(prefix='MyMatShellMult_'), body_matvec, @@ -80,24 +95,25 @@ def make_matvec(self, injectsolve, objs, solver_objs): solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'] ) ) - return matvec_callback + self._matvec_callback = matvec_callback + self._efuncs[matvec_callback.name] = matvec_callback - def create_matvec_body(self, injectsolve, body, solver_objs, objs): - linsolveexpr = injectsolve.expr.rhs + def _create_matvec_body(self, injectsolve, body, solver_objs, objs): + linsolve_expr = injectsolve.expr.rhs - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + dmda = solver_objs['callbackdm'] - body = uxreplace_time(body, solver_objs) + body = self.timedep.uxreplace_time(body) - struct = build_local_struct(body, 'matvec', liveness='eager') + fields = self._dummy_fields(body, solver_objs) - y_matvec = linsolveexpr.arrays['y_matvec'] - x_matvec = linsolveexpr.arrays['x_matvec'] + y_matvec = linsolve_expr.arrays['y_matvec'] + x_matvec = linsolve_expr.arrays['x_matvec'] mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] ) dm_get_local_xvec = petsc_call( @@ -126,7 +142,7 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): ) dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) vec_restore_array_y = petsc_call( @@ -145,6 +161,14 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] ]) + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['Y_local'])] + ) + # TODO: Some of the calls are placed in the `stacks` argument of the # `CallableBody` to ensure that they precede the `cast` statements. The # 'casts' depend on the calls, so this order is necessary. By doing this, @@ -157,7 +181,9 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): (vec_restore_array_y, vec_restore_array_x, dm_local_to_global_begin, - dm_local_to_global_end) + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) ) stacks = ( @@ -173,8 +199,8 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, struct) for i in - struct.fields if isinstance(i.function, AbstractFunction)] + dereference_funcs = [Dereference(i, dummyctx) for i in + fields if isinstance(i.function, AbstractFunction)] matvec_body = CallableBody( List(body=body), @@ -184,47 +210,48 @@ def create_matvec_body(self, injectsolve, body, solver_objs, objs): ) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for i in struct.fields} + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} matvec_body = Uxreplace(subs).visit(matvec_body) - self._struct_params.extend(struct.fields) + self._struct_params.extend(fields) return matvec_body - def make_formfunc(self, injectsolve, objs, solver_objs): + def _make_formfunc(self, injectsolve, objs, solver_objs): # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( injectsolve.expr.rhs.formfuncs, options={'mpi': False}, sregistry=self.sregistry ) - body_formfunc = self.create_formfunc_body(injectsolve, - List(body=irs_formfunc.uiet.body), - solver_objs, objs) + body_formfunc = self._create_formfunc_body(injectsolve, + List(body=irs_formfunc.uiet.body), + solver_objs, objs) formfunc_callback = PETScCallable( self.sregistry.make_name(prefix='FormFunction_'), body_formfunc, retval=objs['err'], parameters=(solver_objs['snes'], solver_objs['X_global'], - solver_objs['Y_global'], solver_objs['dummy']) + solver_objs['F_global'], dummyptr) ) - return formfunc_callback + self._formfunc_callback = formfunc_callback + self._efuncs[formfunc_callback.name] = formfunc_callback - def create_formfunc_body(self, injectsolve, body, solver_objs, objs): - linsolveexpr = injectsolve.expr.rhs + def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): + linsolve_expr = injectsolve.expr.rhs - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + dmda = solver_objs['callbackdm'] - body = uxreplace_time(body, solver_objs) + body = self.timedep.uxreplace_time(body) - struct = build_local_struct(body, 'formfunc', liveness='eager') + fields = self._dummy_fields(body, solver_objs) - y_formfunc = linsolveexpr.arrays['y_formfunc'] - x_formfunc = linsolveexpr.arrays['x_formfunc'] + f_formfunc = linsolve_expr.arrays['f_formfunc'] + x_formfunc = linsolve_expr.arrays['x_formfunc'] snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] ) dm_get_local_xvec = petsc_call( @@ -241,11 +268,11 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): ]) dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + 'DMGetLocalVector', [dmda, Byref(solver_objs['F_local'])] ) vec_get_array_y = petsc_call( - 'VecGetArray', [solver_objs['Y_local'], Byref(y_formfunc._C_symbol)] + 'VecGetArray', [solver_objs['F_local'], Byref(f_formfunc._C_symbol)] ) vec_get_array_x = petsc_call( @@ -253,11 +280,11 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) vec_restore_array_y = petsc_call( - 'VecRestoreArray', [solver_objs['Y_local'], Byref(y_formfunc._C_symbol)] + 'VecRestoreArray', [solver_objs['F_local'], Byref(f_formfunc._C_symbol)] ) vec_restore_array_x = petsc_call( @@ -265,19 +292,29 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] ]) + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['F_local'])] + ) + body = body._rebuild( body=body.body + (vec_restore_array_y, vec_restore_array_x, dm_local_to_global_begin, - dm_local_to_global_end) + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) ) stacks = ( @@ -293,8 +330,8 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, struct) for i in - struct.fields if isinstance(i.function, AbstractFunction)] + dereference_funcs = [Dereference(i, dummyctx) for i in + fields if isinstance(i.function, AbstractFunction)] formfunc_body = CallableBody( List(body=body), @@ -303,20 +340,20 @@ def create_formfunc_body(self, injectsolve, body, solver_objs, objs): retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for i in struct.fields} + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} formfunc_body = Uxreplace(subs).visit(formfunc_body) - self._struct_params.extend(struct.fields) + self._struct_params.extend(fields) return formfunc_body - def make_formrhs(self, injectsolve, objs, solver_objs): + def _make_formrhs(self, injectsolve, objs, solver_objs): # Compile formrhs `eqns` into an IET via recursive compilation irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, options={'mpi': False}, sregistry=self.sregistry) - body_formrhs = self.create_formrhs_body(injectsolve, - List(body=irs_formrhs.uiet.body), - solver_objs, objs) + body_formrhs = self._create_formrhs_body(injectsolve, + List(body=irs_formrhs.uiet.body), + solver_objs, objs) formrhs_callback = PETScCallable( self.sregistry.make_name(prefix='FormRHS_'), body_formrhs, retval=objs['err'], @@ -324,32 +361,32 @@ def make_formrhs(self, injectsolve, objs, solver_objs): solver_objs['snes'], solver_objs['b_local'] ) ) + self._formrhs_callback = formrhs_callback + self._efuncs[formrhs_callback.name] = formrhs_callback - return formrhs_callback - - def create_formrhs_body(self, injectsolve, body, solver_objs, objs): - linsolveexpr = injectsolve.expr.rhs + def _create_formrhs_body(self, injectsolve, body, solver_objs, objs): + linsolve_expr = injectsolve.expr.rhs - dmda = objs['da_so_%s' % linsolveexpr.target.space_order] + dmda = solver_objs['callbackdm'] snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) - b_arr = linsolveexpr.arrays['b_tmp'] + b_arr = linsolve_expr.arrays['b_tmp'] vec_get_array = petsc_call( 'VecGetArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] ) dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(dmda.info)] + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) - body = uxreplace_time(body, solver_objs) + body = self.timedep.uxreplace_time(body) - struct = build_local_struct(body, 'formrhs', liveness='eager') + fields = self._dummy_fields(body, solver_objs) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(struct._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] ) vec_restore_array = petsc_call( @@ -366,8 +403,8 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, struct) for i in - struct.fields if isinstance(i.function, AbstractFunction)] + dereference_funcs = [Dereference(i, dummyctx) for i in + fields if isinstance(i.function, AbstractFunction)] formrhs_body = CallableBody( List(body=[body]), @@ -377,34 +414,325 @@ def create_formrhs_body(self, injectsolve, body, solver_objs, objs): ) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, struct) for - i in struct.fields if not isinstance(i.function, AbstractFunction)} - + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for + i in fields if not isinstance(i.function, AbstractFunction)} formrhs_body = Uxreplace(subs).visit(formrhs_body) - self._struct_params.extend(struct.fields) + self._struct_params.extend(fields) return formrhs_body - def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): + def _local_struct(self, solver_objs): + """ + This is the struct used within callback functions, + usually accessed via DMGetApplicationContext. + """ + solver_objs['localctx'] = petsc_struct( + dummyctx.name, + self.filtered_struct_params, + solver_objs['Jac'].name+'_ctx', + liveness='eager' + ) + + def _main_struct(self, solver_objs): + """ + This is the struct initialised inside the main kernel and + attached to the DM via DMSetApplicationContext. + """ + solver_objs['mainctx'] = petsc_struct( + self.sregistry.make_name(prefix='ctx'), + self.filtered_struct_params, + solver_objs['Jac'].name+'_ctx' + ) + + def _make_struct_callback(self, solver_objs, objs): + mainctx = solver_objs['mainctx'] + body = [ + DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) + for i in mainctx.callback_fields + ] + struct_callback_body = CallableBody( + List(body=body), init=(petsc_func_begin_user,), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + ) + struct_callback = Callable( + self.sregistry.make_name(prefix='PopulateMatContext_'), + struct_callback_body, objs['err'], + parameters=[mainctx] + ) + self._efuncs[struct_callback.name] = struct_callback + self._struct_callback = struct_callback + + def _dummy_fields(self, iet, solver_objs): + # Place all context data required by the shell routines into a struct + fields = [f.function for f in FindSymbols('basics').visit(iet)] + fields = [f for f in fields if not isinstance(f.function, (PETScArray, Temp))] + fields = [ + f for f in fields if not (f.is_Dimension and not (f.is_Time or f.is_Modulo)) + ] + return fields + + def _uxreplace_efuncs(self): + mapper = {} + visitor = Uxreplace({dummyctx: self.solver_objs['localctx']}) + for k, v in self._efuncs.items(): + mapper.update({k: visitor.visit(v)}) + return mapper + + +class BaseObjectBuilder: + """ + A base class for constructing objects needed for a PETSc solver. + Designed to be extended by subclasses, which can override the `_extend_build` + method to support specific use cases. + """ + + def __init__(self, injectsolve, sregistry=None, **kwargs): + self.sregistry = sregistry + self.solver_objs = self._build(injectsolve) + + def _build(self, injectsolve): + """ + Constructs the core dictionary of solver objects and allows + subclasses to extend or modify it via `_extend_build`. + + Returns: + dict: A dictionary containing the following objects: + - 'Jac' (Mat): A matrix representing the jacobian. + - 'x_global' (GlobalVec): The global solution vector. + - 'x_local' (LocalVec): The local solution vector. + - 'b_global': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. + - 'b_local': (LocalVec) Local RHS vector `b`, where `F(x) = b`. + - 'ksp': (KSP) Krylov solver object that manages the linear solver. + - 'pc': (PC) Preconditioner object. + - 'snes': (SNES) Nonlinear solver object. + - 'F_global': (GlobalVec) Global residual vector `F`, where `F(x) = b`. + - 'F_local': (LocalVec) Local residual vector `F`, where `F(x) = b`. + - 'Y_global': (GlobalVector) The output vector populated by the + matrix-free `MyMatShellMult` callback function. + - 'Y_local': (LocalVector) The output vector populated by the matrix-free + `MyMatShellMult` callback function. + - 'X_global': (GlobalVec) Current guess for the solution, + required by the FormFunction callback. + - 'X_local': (LocalVec) Current guess for the solution, + required by the FormFunction callback. + - 'localsize' (PetscInt): The local length of the solution vector. + - 'start_ptr' (StartPtr): A pointer to the beginning of the solution array + that will be updated at each time step. + - 'dmda' (DM): The DMDA object associated with this solve, linked to + the SNES object via `SNESSetDM`. + - 'callbackdm' (CallbackDM): The DM object accessed within callback + functions via `SNESGetDM`. + """ target = injectsolve.expr.rhs.target + sreg = self.sregistry + base_dict = { + 'Jac': Mat(sreg.make_name(prefix='J_')), + 'x_global': GlobalVec(sreg.make_name(prefix='x_global_')), + 'x_local': LocalVec(sreg.make_name(prefix='x_local_'), liveness='eager'), + 'b_global': GlobalVec(sreg.make_name(prefix='b_global_')), + 'b_local': LocalVec(sreg.make_name(prefix='b_local_')), + 'ksp': KSP(sreg.make_name(prefix='ksp_')), + 'pc': PC(sreg.make_name(prefix='pc_')), + 'snes': SNES(sreg.make_name(prefix='snes_')), + 'F_global': GlobalVec(sreg.make_name(prefix='F_global_')), + 'F_local': LocalVec(sreg.make_name(prefix='F_local_'), liveness='eager'), + 'Y_global': GlobalVec(sreg.make_name(prefix='Y_global_')), + 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), + 'X_global': GlobalVec(sreg.make_name(prefix='X_global_')), + 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), + 'localsize': PetscInt(sreg.make_name(prefix='localsize_')), + 'start_ptr': StartPtr(sreg.make_name(prefix='start_ptr_'), target.dtype), + 'dmda': DM(sreg.make_name(prefix='da_'), liveness='eager', + stencil_width=target.space_order), + 'callbackdm': CallbackDM(sreg.make_name(prefix='dm_'), + liveness='eager', stencil_width=target.space_order), + } + return self._extend_build(base_dict, injectsolve) + + def _extend_build(self, base_dict, injectsolve): + """ + Subclasses can override this method to extend or modify the + base dictionary of solver objects. + """ + return base_dict + + +class BaseSetup: + def __init__(self, solver_objs, objs, injectsolve, cbbuilder): + self.calls = self._setup(solver_objs, objs, injectsolve, cbbuilder) - dmda = objs['da_so_%s' % target.space_order] + def _setup(self, solver_objs, objs, injectsolve, cbbuilder): + dmda = solver_objs['dmda'] + + solver_params = injectsolve.expr.rhs.solver_parameters + + snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) + + snes_set_dm = petsc_call('SNESSetDM', [solver_objs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(solver_objs['Jac'])]) + + # NOTE: Assuming all solves are linear for now. + snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) + + snes_set_jac = petsc_call( + 'SNESSetJacobian', [solver_objs['snes'], solver_objs['Jac'], + solver_objs['Jac'], 'MatMFFDComputeJacobian', Null] + ) + + global_x = petsc_call('DMCreateGlobalVector', + [dmda, Byref(solver_objs['x_global'])]) + + global_b = petsc_call('DMCreateGlobalVector', + [dmda, Byref(solver_objs['b_global'])]) + + local_b = petsc_call('DMCreateLocalVector', + [dmda, Byref(solver_objs['b_local'])]) + + snes_get_ksp = petsc_call('SNESGetKSP', + [solver_objs['snes'], Byref(solver_objs['ksp'])]) + + ksp_set_tols = petsc_call( + 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], + solver_params['ksp_atol'], solver_params['ksp_divtol'], + solver_params['ksp_max_it']] + ) + + ksp_set_type = petsc_call( + 'KSPSetType', [solver_objs['ksp'], solver_mapper[solver_params['ksp_type']]] + ) + + ksp_get_pc = petsc_call( + 'KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])] + ) + + # Even though the default will be jacobi, set to PCNONE for now + pc_set_type = petsc_call('PCSetType', [solver_objs['pc'], 'PCNONE']) + + ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) + + matvec_operation = petsc_call( + 'MatShellSetOperation', + [solver_objs['Jac'], 'MATOP_MULT', + MatVecCallback(cbbuilder.matvec_callback.name, void, void)] + ) + + formfunc_operation = petsc_call( + 'SNESSetFunction', + [solver_objs['snes'], Null, + FormFunctionCallback(cbbuilder.formfunc_callback.name, void, void), Null] + ) + + dmda_calls = self._create_dmda_calls(dmda, objs) + + mainctx = solver_objs['mainctx'] + + call_struct_callback = petsc_call( + cbbuilder.struct_callback.name, [Byref(mainctx)] + ) + calls_set_app_ctx = [ + petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) + ] + calls = [call_struct_callback] + calls_set_app_ctx + [BlankLine] + + base_setup = dmda_calls + ( + snes_create, + snes_set_dm, + create_matrix, + snes_set_jac, + snes_set_type, + global_x, + global_b, + local_b, + snes_get_ksp, + ksp_set_tols, + ksp_set_type, + ksp_get_pc, + pc_set_type, + ksp_set_from_ops, + matvec_operation, + formfunc_operation, + ) + tuple(calls) + + extended_setup = self._extend_setup(solver_objs, objs, injectsolve, cbbuilder) + return base_setup + tuple(extended_setup) + + def _extend_setup(self, solver_objs, objs, injectsolve, cbbuilder): + """ + Hook for subclasses to add additional setup calls. + """ + return [] + + def _create_dmda_calls(self, dmda, objs): + dmda_create = self._create_dmda(dmda, objs) + dm_setup = petsc_call('DMSetUp', [dmda]) + dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) + return dmda_create, dm_setup, dm_mat_type + + def _create_dmda(self, dmda, objs): + grid = objs['grid'] + + nspace_dims = len(grid.dimensions) + + # MPI communicator + args = [objs['comm']] + + # Type of ghost nodes + args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(nspace_dims)]) + + # Stencil type + if nspace_dims > 1: + args.append('DMDA_STENCIL_BOX') + + # Global dimensions + args.extend(list(grid.shape)[::-1]) + # No.of processors in each dimension + if nspace_dims > 1: + args.extend(list(grid.distributor.topology)[::-1]) + + # Number of degrees of freedom per node + args.append(1) + # "Stencil width" -> size of overlap + args.append(dmda.stencil_width) + args.extend([Null]*nspace_dims) + + # The distributed array object + args.append(Byref(dmda)) + + # The PETSc call used to create the DMDA + dmda = petsc_call('DMDACreate%sd' % nspace_dims, args) + + return dmda + + +class Solver: + def __init__(self, solver_objs, objs, injectsolve, iters, cbbuilder, + timedep=None, **kwargs): + self.timedep = timedep + self.calls = self._execute_solve(solver_objs, objs, injectsolve, iters, cbbuilder) + self.spatial_body = self._spatial_loop_nest(iters, injectsolve) + + space_iter, = self.spatial_body + self.mapper = {space_iter: self.calls} + + def _execute_solve(self, solver_objs, objs, injectsolve, iters, cbbuilder): + """ + Assigns the required time iterators to the struct and executes + the necessary calls to execute the SNES solver. + """ + struct_assignment = self.timedep.assign_time_iters(solver_objs['mainctx']) + + rhs_callback = cbbuilder.formrhs_callback + + dmda = solver_objs['dmda'] rhs_call = petsc_call(rhs_callback.name, list(rhs_callback.parameters)) local_x = petsc_call('DMCreateLocalVector', [dmda, Byref(solver_objs['x_local'])]) - if any(i.is_Time for i in target.dimensions): - vec_replace_array = time_dep_replace( - injectsolve, solver_objs, objs, self.sregistry - ) - else: - field_from_ptr = FieldFromPointer(target._C_field_data, target._C_symbol) - vec_replace_array = (petsc_call( - 'VecReplaceArray', [solver_objs['x_local'], field_from_ptr] - ),) + vec_replace_array = self.timedep.replace_array(solver_objs) dm_local_to_global_x = petsc_call( 'DMLocalToGlobal', [dmda, solver_objs['x_local'], 'INSERT_VALUES', @@ -424,7 +752,7 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): dmda, solver_objs['x_global'], 'INSERT_VALUES', solver_objs['x_local']] ) - return ( + run_solver_calls = (struct_assignment,) + ( rhs_call, local_x ) + vec_replace_array + ( @@ -434,90 +762,241 @@ def runsolve(self, solver_objs, objs, rhs_callback, injectsolve): dm_global_to_local_x, BlankLine, ) + return List(body=run_solver_calls) - def make_main_struct(self, unique_dmdas, objs): - struct_main = petsc_struct('ctx', filter_ordered(self.struct_params)) - struct_callback = self.generate_struct_callback(struct_main, objs) - call_struct_callback = petsc_call(struct_callback.name, [Byref(struct_main)]) - calls_set_app_ctx = [ - petsc_call('DMSetApplicationContext', [i, Byref(struct_main)]) - for i in unique_dmdas - ] - calls = [call_struct_callback] + calls_set_app_ctx + def _spatial_loop_nest(self, iters, injectsolve): + spatial_body = [] + for tree in retrieve_iteration_tree(iters[0]): + root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] + if injectsolve in FindNodes(InjectSolveDummy).visit(root): + spatial_body.append(root) + return spatial_body - self._efuncs[struct_callback.name] = struct_callback - return struct_main, calls - def generate_struct_callback(self, struct, objs): - body = [ - DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol) - for i in struct.fields if i not in struct.time_dim_fields - ] - struct_callback_body = CallableBody( - List(body=body), init=tuple([petsc_func_begin_user]), - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) - ) - struct_callback = Callable( - 'PopulateMatContext', struct_callback_body, objs['err'], - parameters=[struct] - ) - return struct_callback +class NonTimeDependent: + def __init__(self, injectsolve, iters, **kwargs): + self.injectsolve = injectsolve + self.iters = iters + self.kwargs = kwargs + self.origin_to_moddim = self._origin_to_moddim_mapper(iters) + self.time_idx_to_symb = injectsolve.expr.rhs.time_mapper + @property + def is_target_time(self): + return False -def build_local_struct(iet, name, liveness): - # Place all context data required by the shell routines into a struct - fields = [ - i.function for i in FindSymbols('basics').visit(iet) - if not isinstance(i.function, (PETScArray, Temp)) - and not (i.is_Dimension and not isinstance(i, (TimeDimension, ModuloDimension))) - ] - return petsc_struct(name, fields, liveness) + @property + def target(self): + return self.injectsolve.expr.rhs.target + def _origin_to_moddim_mapper(self, iters): + return {} -def time_dep_replace(injectsolve, solver_objs, objs, sregistry): - target = injectsolve.expr.lhs - target_time = [ - i for i, d in zip(target.indices, target.dimensions) if d.is_Time - ] - assert len(target_time) == 1 - target_time = target_time.pop() + def uxreplace_time(self, body): + return body - start_ptr = solver_objs['start_ptr'] + def replace_array(self, solver_objs): + """ + VecReplaceArray() is a PETSc function that allows replacing the array + of a `Vec` with a user provided array. + https://petsc.org/release/manualpages/Vec/VecReplaceArray/ - vec_get_size = petsc_call( - 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] - ) + This function is used to replace the array of the PETSc solution `Vec` + with the array from the `Function` object representing the target. - field_from_ptr = FieldFromPointer( - target.function._C_field_data, target.function._C_symbol - ) + Examples + -------- + >>> self.target + f1(x, y) + >>> call = replace_array(solver_objs) + >>> print(call) + PetscCall(VecReplaceArray(x_local_0,f1_vec->data)); + """ + field_from_ptr = FieldFromPointer( + self.target.function._C_field_data, self.target.function._C_symbol + ) + vec_replace_array = (petsc_call( + 'VecReplaceArray', [solver_objs['x_local'], field_from_ptr] + ),) + return vec_replace_array - expr = DummyExpr( - start_ptr, cast_mapper[(target.dtype, '*')](field_from_ptr) + - Mul(target_time, solver_objs['localsize']), init=True - ) + def assign_time_iters(self, struct): + return [] - vec_replace_array = petsc_call('VecReplaceArray', [solver_objs['x_local'], start_ptr]) - return (vec_get_size, expr, vec_replace_array) +class TimeDependent(NonTimeDependent): + """ + A class for managing time-dependent solvers. + + This includes scenarios where the target is not directly a `TimeFunction`, + but depends on other functions that are. + + Outline of time loop abstraction with PETSc: + + - At PETScSolve, time indices are replaced with temporary `Symbol` objects + via a mapper (e.g., {t: tau0, t + dt: tau1}) to prevent the time loop + from being generated in the callback functions. These callbacks, needed + for each `SNESSolve` at every time step, don't require the time loop, but + may still need access to data from other time steps. + - All `Function` objects are passed through the initial lowering via the + `LinearSolveExpr` object, ensuring the correct time loop is generated + in the main kernel. + - Another mapper is created based on the modulo dimensions + generated by the `LinearSolveExpr` object in the main kernel + (e.g., {time: time, t: t0, t + 1: t1}). + - These two mappers are used to generate a final mapper `symb_to_moddim` + (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to + replace the temporary `Symbol` objects in the callback functions with + the correct modulo dimensions. + - Modulo dimensions are updated in the matrix context struct at each time + step and can be accessed in the callback functions where needed. + """ + @property + def is_target_time(self): + return any(i.is_Time for i in self.target.dimensions) -def uxreplace_time(body, solver_objs): - # TODO: Potentially introduce a TimeIteration abstraction to simplify - # all the time processing that is done (searches, replacements, ...) - # "manually" via free functions - time_spacing = solver_objs['target'].grid.stepping_dim.spacing - true_dims = solver_objs['true_dims'] + @property + def time_spacing(self): + return self.target.grid.stepping_dim.spacing - time_mapper = { - v: k.xreplace({time_spacing: 1, -time_spacing: -1}) - for k, v in solver_objs['time_mapper'].items() - } - subs = {symb: true_dims[time_mapper[symb]] for symb in time_mapper} - return Uxreplace(subs).visit(body) + @property + def target_time(self): + target_time = [ + i for i, d in zip(self.target.indices, self.target.dimensions) + if d.is_Time + ] + assert len(target_time) == 1 + target_time = target_time.pop() + return target_time + + @property + def symb_to_moddim(self): + """ + Maps temporary `Symbol` objects created during `PETScSolve` to their + corresponding modulo dimensions (e.g. creates {tau0: t0, tau1: t1}). + """ + mapper = { + v: k.xreplace({self.time_spacing: 1, -self.time_spacing: -1}) + for k, v in self.time_idx_to_symb.items() + } + return {symb: self.origin_to_moddim[mapper[symb]] for symb in mapper} + + def uxreplace_time(self, body): + return Uxreplace(self.symb_to_moddim).visit(body) + + def _origin_to_moddim_mapper(self, iters): + """ + Creates a mapper of the origin of the time dimensions to their corresponding + modulo dimensions from a list of `Iteration` objects. + + Examples + -------- + >>> iters + (, + ) + >>> _origin_to_moddim_mapper(iters) + {time: time, t: t0, t + 1: t1} + """ + time_iter = [i for i in iters if any(d.is_Time for d in i.dimensions)] + mapper = {} + + if not time_iter: + return mapper + + for i in time_iter: + for d in i.dimensions: + if d.is_Modulo: + mapper[d.origin] = d + elif d.is_Time: + mapper[d] = d + return mapper + + def replace_array(self, solver_objs): + """ + In the case that the actual target is time-dependent e.g a `TimeFunction`, + a pointer to the first element in the array that will be updated during + the time step is passed to VecReplaceArray(). + + Examples + -------- + >>> self.target + f1(time + dt, x, y) + >>> calls = replace_array(solver_objs) + >>> print(List(body=calls)) + PetscCall(VecGetSize(x_local_0,&(localsize_0))); + float * start_ptr_0 = (time + 1)*localsize_0 + (float*)(f1_vec->data); + PetscCall(VecReplaceArray(x_local_0,start_ptr_0)); + + >>> self.target + f1(t + dt, x, y) + >>> calls = replace_array(solver_objs) + >>> print(List(body=calls)) + PetscCall(VecGetSize(x_local_0,&(localsize_0))); + float * start_ptr_0 = t1*localsize_0 + (float*)(f1_vec->data); + """ + if self.is_target_time: + mapper = {self.time_spacing: 1, -self.time_spacing: -1} + target_time = self.target_time.xreplace(mapper) + + try: + target_time = self.origin_to_moddim[target_time] + except KeyError: + pass + + start_ptr = solver_objs['start_ptr'] + + vec_get_size = petsc_call( + 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] + ) + + field_from_ptr = FieldFromPointer( + self.target.function._C_field_data, self.target.function._C_symbol + ) + + expr = DummyExpr( + start_ptr, cast_mapper[(self.target.dtype, '*')](field_from_ptr) + + Mul(target_time, solver_objs['localsize']), init=True + ) + + vec_replace_array = petsc_call( + 'VecReplaceArray', [solver_objs['x_local'], start_ptr] + ) + return (vec_get_size, expr, vec_replace_array) + else: + return super().replace_array(solver_objs) + + def assign_time_iters(self, struct): + """ + Assign required time iterators to the struct. + These iterators are updated at each timestep in the main kernel + for use in callback functions. + + Examples + -------- + >>> struct + ctx + >>> struct.fields + [h_x, x_M, x_m, f1(t, x), t0, t1] + >>> assigned = assign_time_iters(struct) + >>> print(assigned[0]) + ctx.t0 = t0; + >>> print(assigned[1]) + ctx.t1 = t1; + """ + to_assign = [ + f for f in struct.fields if (f.is_Dimension and (f.is_Time or f.is_Modulo)) + ] + time_iter_assignments = [ + DummyExpr(FieldFromComposite(field, struct), field) + for field in to_assign + ] + return time_iter_assignments Null = Macro('NULL') void = 'void' +dummyctx = Symbol('lctx') +dummyptr = DummyArg('dummy') # TODO: Don't use c.Line here? diff --git a/devito/petsc/types.py b/devito/petsc/types.py index 62c69f666a..72a0faa1e8 100644 --- a/devito/petsc/types.py +++ b/devito/petsc/types.py @@ -1,17 +1,19 @@ import sympy -from functools import cached_property -import numpy as np +<<<<<<< HEAD +from devito.tools import Reconstructable, sympy_mutex +======= from devito.tools import CustomDtype from devito.types import LocalObject, Eq, CompositeObject from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction, Symbol +from devito.types.basic import AbstractFunction, AbstractSymbol from devito.finite_differences.tools import fd_weights_registry -from devito.tools import Reconstructable, dtype_to_ctype -from devito.symbolics import FieldFromComposite -from devito.types.basic import IndexedBase +from devito.tools import dtype_to_ctype, Reconstructable, sympy_mutex +from devito.symbolics import FieldFromComposite, Byref + +from devito.petsc.utils import petsc_call class DM(LocalObject): @@ -28,6 +30,15 @@ def __init__(self, *args, stencil_width=None, **kwargs): def stencil_width(self): return self._stencil_width + @property + def info(self): + return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') + + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('DMDestroy', [Byref(self.function)]) + class Mat(LocalObject): """ @@ -35,13 +46,30 @@ class Mat(LocalObject): """ dtype = CustomDtype('Mat') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('MatDestroy', [Byref(self.function)]) + + +class LocalVec(LocalObject): + """ + PETSc Vector object (Vec). + """ + dtype = CustomDtype('Vec') + -class Vec(LocalObject): +class GlobalVec(LocalObject): """ PETSc Vector object (Vec). """ dtype = CustomDtype('Vec') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('VecDestroy', [Byref(self.function)]) + class PetscMPIInt(LocalObject): """ @@ -65,6 +93,11 @@ class SNES(LocalObject): """ dtype = CustomDtype('SNES') + @property + def _C_free(self): + # from devito.petsc.utils import petsc_call + return petsc_call('SNESDestroy', [Byref(self.function)]) + class PC(LocalObject): """ @@ -97,6 +130,10 @@ class PetscErrorCode(LocalObject): dtype = CustomDtype('PetscErrorCode') +class DummyArg(LocalObject): + dtype = CustomDtype('void', modifier='*') + + class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent @@ -114,7 +151,8 @@ class PETScArray(ArrayBasic, Differentiable): _default_fd = 'taylor' __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients')) + ('dimensions', 'shape', 'liveness', 'coefficients', + 'space_order')) def __init_finalize__(self, *args, **kwargs): @@ -126,6 +164,7 @@ def __init_finalize__(self, *args, **kwargs): raise ValueError("coefficients must be one of %s" " not %s" % (str(fd_weights_registry), self._coefficients)) self._shape = kwargs.get('shape') + self._space_order = kwargs.get('space_order', 1) @classmethod def __dtype_setup__(cls, **kwargs): @@ -140,6 +179,10 @@ def coefficients(self): def shape(self): return self._shape + @property + def space_order(self): + return self._space_order + @cached_property def _shape_with_inhalo(self): """ @@ -172,37 +215,28 @@ def shape_allocated(self): @cached_property def _C_ctype(self): - return CustomDtype(self.petsc_type, modifier='*') - - @property - def petsc_type(self): - return dtype_to_petsctype(self._dtype) - - @property - def dtype(self): - return CustomDtype(self.petsc_type) + # NOTE: Reverting to using float/double instead of PetscScalar for + # simplicity when opt='advanced'. Otherwise, Temp objects must also + # be converted to PetscScalar. Additional tests are needed to + # ensure this approach is fine. Previously, issues arose from + # mismatches between precision of Function objects in Devito and the + # precision of the PETSc configuration. + # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables + # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of + # the user's PETSc configuration. + return POINTER(dtype_to_ctype(self.dtype)) @property def symbolic_shape(self): field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] + FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] # Reverse it since DMDA is setup backwards to Devito dimensions. return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - @cached_property - def dmda_info(self): - # To access the local grid info via the DM in - # PETSc you use DMDAGetLocalInfo(da, &info) - # and then access the local no.of grid points via info.gmx, info.gmy, info.gmz - return DMDALocalInfo(name='info', liveness='eager') - - @cached_property - def indexed(self): - return PETScIndexedData(self.name, shape=self._shape, function=self.function) - - -class PETScIndexedData(IndexedBase): - pass + @property + def dmda(self): + name = 'da_so_%s' % self.space_order + return DM(name=name, liveness='eager', stencil_width=self.space_order) def dtype_to_petsctype(dtype): @@ -216,85 +250,25 @@ def dtype_to_petsctype(dtype): }[dtype] -class MatVecEq(Eq): - """ - Represents the mathematical expression of applying a linear - operator to a vector. This is a key component - for running matrix-free solvers. - """ +class InjectSolveEq(Eq): pass - - -class RHSEq(Eq): - """ - Represents the mathematical expression of building the - rhs of a linear system. - """ - pass - - -class MockEq(Eq): - """ - Represents a mock/placeholder equation to ensure distinct iteration loops. - - For example, the mat-vec action iteration loop is to be isolated from the - expression loop used to build the RHS of the linear system. This separation - facilitates the utilisation of the mat-vec iteration loop in callback functions - created at the IET level. - """ - pass - - -def PETScSolve(eq, target, bcs=None, solver_parameters=None, **kwargs): - # TODO: Add check for time dimensions and utilise implicit dimensions. - - is_time_dep = any(dim.is_Time for dim in target.dimensions) - # TODO: Current assumption is rhs is part of pde that remains - # constant at each timestep. Need to insert function to extract this from eq. - y_matvec, x_matvec, b_tmp = [ - PETScArray(name=f'{prefix}_{target.name}', - dtype=target.dtype, - dimensions=target.space_dimensions, - shape=target.grid.shape, - liveness='eager', - halo=target.halo[1:] if is_time_dep else target.halo) - for prefix in ['y_matvec', 'x_matvec', 'b_tmp']] - - # TODO: Extend to rearrange equation for implicit time stepping. - matvecaction = MatVecEq(y_matvec, LinearSolveExpr(eq.lhs.subs(target, x_matvec), - target=target, solver_parameters=solver_parameters), - subdomain=eq.subdomain) - - # Part of pde that remains constant at each timestep - rhs = RHSEq(b_tmp, LinearSolveExpr(eq.rhs, target=target, - solver_parameters=solver_parameters), subdomain=eq.subdomain) - - if not bcs: - return [matvecaction, rhs] - - bcs_for_matvec = [] - for bc in bcs: - # TODO: Insert code to distiguish between essential and natural - # boundary conditions since these are treated differently within - # the solver - # NOTE: May eventually remove the essential bcs from the solve - # (and move to rhs) but for now, they are included since this - # is not trivial to implement when using DMDA - # NOTE: Below is temporary -> Just using this as a palceholder for - # the actual BC implementation for the matvec callback - new_rhs = bc.rhs.subs(target, x_matvec) - bc_rhs = LinearSolveExpr( - new_rhs, target=target, solver_parameters=solver_parameters - ) - bcs_for_matvec.append(MatVecEq(y_matvec, bc_rhs, subdomain=bc.subdomain)) - - return [matvecaction] + bcs_for_matvec + [rhs] +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) class LinearSolveExpr(sympy.Function, Reconstructable): __rargs__ = ('expr',) - __rkwargs__ = ('target', 'solver_parameters') +<<<<<<< HEAD +<<<<<<< HEAD + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays') +======= + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', 'formfuncs', 'formrhs') +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays') +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) defaults = { 'ksp_type': 'gmres', @@ -305,7 +279,16 @@ class LinearSolveExpr(sympy.Function, Reconstructable): 'ksp_max_it': 10000 # Maximum iterations } - def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): + def __new__(cls, expr, target=None, solver_parameters=None, +<<<<<<< HEAD +<<<<<<< HEAD + matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): +======= + matvecs=None, formfuncs=None, formrhs=None, **kwargs): +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) if solver_parameters is None: solver_parameters = cls.defaults @@ -313,10 +296,22 @@ def __new__(cls, expr, target=None, solver_parameters=None, **kwargs): for key, val in cls.defaults.items(): solver_parameters[key] = solver_parameters.get(key, val) - obj = super().__new__(cls, expr) + with sympy_mutex: + obj = sympy.Basic.__new__(cls, expr) obj._expr = expr obj._target = target obj._solver_parameters = solver_parameters + obj._matvecs = matvecs + obj._formfuncs = formfuncs + obj._formrhs = formrhs +<<<<<<< HEAD +<<<<<<< HEAD + obj._arrays = arrays +======= +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + obj._arrays = arrays +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) return obj def __repr__(self): @@ -328,7 +323,12 @@ def _sympystr(self, printer): return str(self) def __hash__(self): - return hash(self.target) + return hash(self.expr) + + def __eq__(self, other): + return (isinstance(other, LinearSolveExpr) and + self.expr == other.expr and + self.target == other.target) @property def expr(self): @@ -342,43 +342,88 @@ def target(self): def solver_parameters(self): return self._solver_parameters - func = Reconstructable._rebuild - + @property + def matvecs(self): + return self._matvecs -def petsc_lift(clusters): - """ - Lift the iteration space associated with each PETSc equation. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions. - """ - processed = [] - for c in clusters: + @property + def formfuncs(self): + return self._formfuncs - ispace = c.ispace - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) + @property + def formrhs(self): + return self._formrhs - processed.append(c.rebuild(ispace=ispace)) +<<<<<<< HEAD +<<<<<<< HEAD + @property + def formrhs(self): + return self._formrhs +======= +======= + @property + def arrays(self): + return self._arrays - return processed +>>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) + func = Reconstructable._rebuild class PETScStruct(CompositeObject): __rargs__ = ('name', 'usr_ctx') + __rkwargs__ = ('liveness',) - def __init__(self, name, usr_ctx): + def __init__(self, name, usr_ctx, liveness='lazy'): pfields = [(i._C_name, dtype_to_ctype(i.dtype)) - for i in usr_ctx if isinstance(i, Symbol)] + for i in usr_ctx if isinstance(i, AbstractSymbol)] self._usr_ctx = usr_ctx super().__init__(name, 'MatContext', pfields) +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) + + assert liveness in ['eager', 'lazy'] + self._liveness = liveness @property - def usr_ctx(self): - return self._usr_ctx + def arrays(self): + return self._arrays +<<<<<<< HEAD + func = Reconstructable._rebuild +======= def _arg_values(self, **kwargs): + # TODO: Handle passing `struct dataobj *restrict g1_vec` and similar + # elements within in the struct. It seemed necessary to convert the Numpy array + # into a pointer to ctypes.Struct within MatContext + # (i.e override _arg_finalize for PETScStruct here?), but this approach seemed + # to result in a NIL pointer in the generated code ... any suggestions? + # Alternative: Allow these pointers to appear in the main kernel arguments + # as usual and use ctx->g1_vec = g1_vec; to assign them to the struct? values = super()._arg_values(**kwargs) - for i in self.fields: - setattr(values[self.name]._obj, i, kwargs['args'][i]) + for i in self.usr_ctx: + setattr(values[self.name]._obj, i.name, kwargs['args'][i.name]) return values + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' +<<<<<<< HEAD +<<<<<<< HEAD +>>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) +======= + +======= + +>>>>>>> 9d611239f (compiler: Fix derive_callback_inputs) + @property + def fields(self): + return self._usr_ctx +>>>>>>> 41157b0bc (compiler: Place dataobj pointers inside ctx struct) diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index ccc47a4c3e..d898db23cb 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -51,16 +51,3 @@ def core_metadata(): 'lib_dirs': lib_dir, 'ldflags': ('-Wl,-rpath,%s' % lib_dir) } -<<<<<<< HEAD -======= - - -def petsc_call(specific_call, call_args): - general_call = 'PetscCall' - return Call(general_call, [Call(specific_call, arguments=call_args)]) - - -def petsc_call_mpi(specific_call, call_args): - general_call = 'PetscCallMPI' - return Call(general_call, [Call(specific_call, arguments=call_args)]) ->>>>>>> f921beaef (misc: Update file due to incorrect rebase) diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc new file mode 100644 index 0000000000..782e0ac340 --- /dev/null +++ b/docker/Dockerfile.petsc @@ -0,0 +1,49 @@ +############################################################## +# This Dockerfile builds a base image to run Devito with PETSc +############################################################# + +# Base image +FROM ubuntu:22.04 as base + +ENV DEBIAN_FRONTEND noninteractive + +# Install python +RUN apt-get update && \ + apt-get install -y dh-autoreconf python3-venv python3-dev python3-pip libopenblas-serial-dev git pkgconf mpich + +# Install for basic base not containing it +RUN apt-get install -y vim wget git flex libnuma-dev tmux \ + numactl hwloc curl \ + autoconf libtool build-essential procps + +# Install tmpi +RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi + +# Install OpenGL library, necessary for the installation of GemPy +RUN apt-get install -y libgl1-mesa-glx + +# Install numpy from source +RUN python3 -m venv /venv && \ + /venv/bin/pip install --no-cache-dir --upgrade pip && \ + /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ + rm -rf ~/.cache/pip + +# Create a directory for PETSc +RUN mkdir -p /home/petsc + +# Clone and install PETSc +RUN cd /home/petsc \ + && git clone -b release https://gitlab.com/petsc/petsc.git petsc \ + && cd petsc \ + && git pull \ + && ./configure --with-fortran-bindings=0 --with-openblas-include=$(pkgconf --variable=includedir openblas) --with-openblas-lib=$(pkgconf --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build \ + && make all + +RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +EXPOSE 8888 +CMD ["/bin/bash"] + +ENV PETSC_ARCH="devito_build" +ENV PETSC_DIR="/home/petsc/petsc" From b8de6145200b82611e6b1f1b3eb0d76345708cc4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 23 Jul 2024 21:18:45 +0100 Subject: [PATCH 124/151] compiler: Create mixin for ThreadCallable and PETScCallable --- .github/workflows/examples-mpi.yml | 2 +- .github/workflows/examples.yml | 2 +- .github/workflows/flake8.yml | 2 +- .github/workflows/pytest-core-mpi.yml | 2 +- .github/workflows/pytest-core-nompi.yml | 2 +- .github/workflows/tutorials.yml | 2 +- devito/ir/equations/equation.py | 4 +- devito/ir/iet/algorithms.py | 1 - devito/ir/iet/nodes.py | 4 - devito/ir/iet/visitors.py | 1 - devito/operator/operator.py | 11 +- devito/passes/clusters/petsc.py | 25 -- devito/passes/iet/__init__.py | 1 - devito/passes/iet/definitions.py | 4 +- devito/passes/iet/petsc.py | 298 ---------------- devito/petsc/clusters.py | 19 +- devito/petsc/iet.py | 372 -------------------- devito/petsc/iet/utils.py | 4 +- devito/petsc/types.py | 429 ------------------------ devito/petsc/types/object.py | 48 ++- devito/petsc/types/types.py | 45 ++- devito/symbolics/printer.py | 3 - devito/types/basic.py | 31 ++ docker/Dockerfile.petsc | 49 --- 24 files changed, 145 insertions(+), 1216 deletions(-) delete mode 100644 devito/passes/clusters/petsc.py delete mode 100644 devito/passes/iet/petsc.py delete mode 100644 devito/petsc/iet.py delete mode 100644 devito/petsc/types.py delete mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 5553e9108d..bbf3105635 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,7 +17,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 2a94b27393..85f84f0cd5 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 9a0aea32a8..92556b6be4 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 2b36602110..2eae16021b 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index b818f2b8f6..5cf374f222 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index 42a2288c56..93d64e26ac 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,7 +10,7 @@ on: push: branches: - master - - enable-opt + - time_loop_branch pull_request: branches: - master diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 0d6c7b0ff7..67ed9269a4 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -227,13 +227,13 @@ def __new__(cls, *args, **kwargs): expr = uxreplace(expr, {d: IntDiv(index, d.factor)}) conditionals = frozendict(conditionals) - # from IPython import embed; embed() + # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) # Finally create the LoweredEq with all metadata attached expr = super().__new__(cls, expr.lhs, rhs, evaluate=False) - # from IPython import embed; embed() + expr._ispace = ispace expr._conditionals = conditionals expr._reads, expr._writes = detect_io(expr) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index af48dc73c7..9d2db185db 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -24,7 +24,6 @@ def iet_build(stree): elif i.is_Exprs: exprs = [] for e in i.exprs: - # from IPython import embed; embed() if e.is_Increment: exprs.append(Increment(e)) elif isinstance(e.rhs, LinearSolveExpr): diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index f130334e95..fa7716aaf8 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -29,11 +29,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', -<<<<<<< HEAD 'CallableBody', 'Transfer', 'Callback', 'FixedArgsCallable'] -======= - 'CallableBody', 'Transfer', 'Callback'] ->>>>>>> f921beaef (misc: Update file due to incorrect rebase) # First-class IET nodes diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 48570fa89f..d23f24a23d 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -315,7 +315,6 @@ def _args_call(self, args): def _gen_signature(self, o, is_declaration=False): decls = self._args_decl(o.parameters) - prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),)) if o.attributes: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 6a4f66f28d..f3f4e8565b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -354,7 +354,7 @@ def _lower_exprs(cls, expressions, **kwargs): expressions = concretize_subdims(expressions, **kwargs) processed = [LoweredEq(i) for i in expressions] - # from IPython import embed; embed() + return processed # Compilation -- Cluster level @@ -389,6 +389,8 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): clusters = cls._specialize_clusters(clusters, **kwargs) + clusters = petsc_project(clusters) + # Operation count after specialization final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) try: @@ -428,7 +430,6 @@ def _lower_stree(cls, clusters, **kwargs): * Derive sections for performance profiling """ # Build a ScheduleTree from a sequence of Clusters - # from IPython import embed; embed() stree = stree_build(clusters, **kwargs) stree = cls._specialize_stree(stree) @@ -498,9 +499,6 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Target-independent optimizations minimize_symbols(graph) - # If necessary, sort frees into a specific order - sort_frees(graph) - return graph.root, graph # Read-only properties exposed to the outside world @@ -527,8 +525,7 @@ def dimensions(self): @cached_property def input(self): - struct_params = derive_struct_inputs(self.parameters) - return tuple(i for i in self.parameters+struct_params if i.is_Input) + return tuple(i for i in self.parameters if i.is_Input) @cached_property def temporaries(self): diff --git a/devito/passes/clusters/petsc.py b/devito/passes/clusters/petsc.py deleted file mode 100644 index d784d99c0d..0000000000 --- a/devito/passes/clusters/petsc.py +++ /dev/null @@ -1,25 +0,0 @@ -from devito.tools import timed_pass -from devito.petsc import LinearSolveExpr - -__all__ = ['petsc_lift'] - - -@timed_pass() -def petsc_lift(clusters): - """ - Lift the iteration space surrounding each PETSc equation to create - distinct iteration loops. This simplifys the movement of the loops - into specific callback functions generated at the IET level. - TODO: Potentially only need to lift the PETSc equations required - by the callback functions, not the ones that stay inside the main kernel. - """ - processed = [] - for c in clusters: - - if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.dimensions) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - - return processed diff --git a/devito/passes/iet/__init__.py b/devito/passes/iet/__init__.py index cf6a35de90..c09db00c9b 100644 --- a/devito/passes/iet/__init__.py +++ b/devito/passes/iet/__init__.py @@ -8,4 +8,3 @@ from .instrument import * # noqa from .languages import * # noqa from .errors import * # noqa -from .petsc import * # noqa diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 98fb8cabfe..59de8f95c7 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -456,9 +456,7 @@ def place_casts(self, iet, **kwargs): # Some objects don't distinguish their _C_symbol because they are known, # by construction, not to require it, thus making the generated code # cleaner. These objects don't need a cast - bases = [ - i for i in bases - if i.name != i.function._C_name and not isinstance(i.function, PETScArray)] + bases = [i for i in bases if i.name != i.function._C_name] # Create and attach the type casts casts = tuple(self.lang.PointerCast(i.function, obj=i) for i in bases diff --git a/devito/passes/iet/petsc.py b/devito/passes/iet/petsc.py deleted file mode 100644 index 8c8aeafa1c..0000000000 --- a/devito/passes/iet/petsc.py +++ /dev/null @@ -1,298 +0,0 @@ -from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, MatVecAction, - Transformer, FindSymbols, LinearSolverExpression, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition) -from devito.types import (PetscMPIInt, PETScStruct, DMDALocalInfo, DM, Mat, - Vec, KSP, PC, SNES, PetscErrorCode) -from devito.symbolics import Byref, Macro, FieldFromPointer -import cgen as c - -__all__ = ['lower_petsc'] - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # TODO: Drop the LinearSolveExpr's using .args[0] so that _rebuild doesn't - # appear in ccode - - # Check if PETScSolve was used and count occurrences. Each PETScSolve - # will have a unique MatVecAction. - is_petsc = FindNodes(MatVecAction).visit(iet) - - if is_petsc: - - # Collect all solution fields we're solving for - targets = [i.expr.rhs.target for i in is_petsc] - - # Initalize PETSc - init = init_petsc(**kwargs) - - # Create context data struct. - struct = build_struct(iet) - - objs = build_core_objects(targets[-1], struct, **kwargs) - - # Create core PETSc calls required by general linear solves, - # which only need to be generated once e.g create DMDA. - core = core_petsc(targets[-1], objs, **kwargs) - - matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) - - main_mapper = {} - - setup = [] - efuncs = [] - for target in unique_targets: - - solver_objs = build_solver_objs(target) - - matvec_callback_body_iters = [] - - solver_setup = False - - for iter, (matvec,) in matvec_mapper.items(): - - if matvec.expr.rhs.target == target: - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec) - setup.extend(solver) - solver_setup = True - - # Only need to generate solver setup once per target - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) - solver_setup = True - - matvec_body = matvec_body_list._rebuild(body=[ - matvec_body_list.body, iter[0]]) - matvec_body_list = matvec_body_list._rebuild(body=matvec_body) - - main_mapper.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, matvec_callback_body_iters, solver_objs, objs, struct) - - setup.append(matvec_op) - setup.append(c.Line()) - - efuncs.append(matvec_callback) - - # Remove the LinSolveExpr from iet and efuncs that were used to carry - # metadata e.g solver_parameters - main_mapper.update(rebuild_expr_mapper(iet)) - efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} - - iet = Transformer(main_mapper).visit(iet) - efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] - - # Replace symbols appearing in each efunc with a pointer to the PETScStruct - efuncs = transform_efuncs(efuncs, struct) - - body = iet.body._rebuild(body=(tuple(init_setup) + iet.body.body)) - iet = iet._rebuild(body=body) - - return iet, {} - - -def init_petsc(**kwargs): - - # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve. - # NOTE: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options. - initialize = Call('PetscCall', [Call('PetscInitialize', - arguments=['NULL', 'NULL', - 'NULL', 'NULL'])]) - - return tuple([initialize]) - - -def build_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct. - usr_ctx = [] - - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - -def core_petsc(target, objs, **kwargs): - # Assumption: all targets are generated from the same Grid, - # so we can use any target. - - # MPI - call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', - arguments=[objs['comm'], - Byref(objs['size'])])]) - # Create DMDA - dmda = create_dmda(target, objs) - dm_setup = Call('PetscCall', [Call('DMSetUp', arguments=[objs['da']])]) - dm_app_ctx = Call('PetscCall', [Call('DMSetApplicationContext', - arguments=[objs['da'], objs['struct']])]) - dm_mat_type = Call('PetscCall', [Call('DMSetMatType', - arguments=[objs['da'], 'MATSHELL'])]) - dm_local_info = Call('PetscCall', [Call('DMDAGetLocalInfo', - arguments=[objs['da'], Byref(objs['info'])])]) - - return tuple([call_mpi, dmda, dm_setup, dm_app_ctx, dm_mat_type, dm_local_info]) - - -def build_core_objects(target, struct, **kwargs): - - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' - - return {'da': DM(name='da'), - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'struct': struct} - - -def create_dmda(target, objs): - - args = [objs['comm']] - - args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] - - # stencil type - args += ['DMDA_STENCIL_BOX'] - - # global dimensions - args += list(target.shape_global)[::-1] - - # no.of processors in each dimension - args += list(target.grid.distributor.topology)[::-1] - - args += [1, target.space_order] - - args += ['NULL' for _ in range(len(target.space_dimensions))] - - args += [Byref(objs['da'])] - - dmda = Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args) - - return dmda - - -def build_solver_objs(target): - - return {'Jac': Mat(name='J_'+str(target.name)), - 'x_global': Vec(name='x_global_'+str(target.name)), - 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), - 'b_global': Vec(name='b_global_'+str(target.name)), - 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), - 'ksp': KSP(name='ksp_'+str(target.name)), - 'pc': PC(name='pc_'+str(target.name)), - 'snes': SNES(name='snes_'+str(target.name)), - 'x': Vec(name='x_'+str(target.name)), - 'y': Vec(name='y_'+str(target.name))} - - -def generate_solver_calls(solver_objs, objs, matvec): - - snes_create = Call('PetscCall', [Call('SNESCreate', arguments=[ - objs['comm'], Byref(solver_objs['snes'])])]) - - snes_set_dm = Call('PetscCall', [Call('SNESSetDM', arguments=[ - solver_objs['snes'], objs['da']])]) - - create_matrix = Call('PetscCall', [Call('DMCreateMatrix', arguments=[ - objs['da'], Byref(solver_objs['Jac'])])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = Call('PetscCall', [Call('SNESSetType', arguments=[ - solver_objs['snes'], 'SNESKSPONLY'])]) - - global_x = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['x_global'])])]) - - local_x = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['x_local'])])]) - - global_b = Call('PetscCall', [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['b_global'])])]) - - local_b = Call('PetscCall', [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['b_local'])])]) - - snes_get_ksp = Call('PetscCall', [Call('SNESGetKSP', arguments=[ - solver_objs['snes'], Byref(solver_objs['ksp'])])]) - - return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, - global_x, local_x, global_b, local_b, snes_get_ksp]) - - -def create_matvec_callback(target, matvec_callback_body_iters, - solver_objs, objs, struct): - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - defn_struct = Definition(struct) - - get_context = Call('PetscCall', [Call('MatShellGetContext', - arguments=[solver_objs['Jac'], - Byref(struct)])]) - - body = List(body=[defn_struct, - get_context, - matvec_callback_body_iters]) - - matvec_callback = Callable('MyMatShellMult_'+str(target.name), - matvec_body, - retval=objs['err'], - parameters=(solver_objs['Jac'], - solver_objs['x'], - solver_objs['y'])) - - matvec_operation = Call('PetscCall', [ - Call('MatShellSetOperation', arguments=[solver_objs['Jac'], - 'MATOP_MULT', - Callback(matvec_callback.name, - Void, Void)])]) - - return matvec_callback, matvec_operation - - -def rebuild_expr_mapper(iet): - - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for - expr in FindNodes(LinearSolverExpression).visit(iet)} - - -def transform_efuncs(efuncs, struct): - - efuncs_new = [] - for efunc in efuncs: - new_body = efunc.body - for i in struct.usr_ctx: - new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) - efunc_with_new_body = efunc._rebuild(body=new_body) - efuncs_new.append(efunc_with_new_body) - - return efuncs_new - - -Null = Macro('NULL') -Void = Macro('void') - -petsc_call = String('PetscCall') -petsc_call_mpi = String('PetscCallMPI') -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') - -linear_solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 0a6679b2dd..143c73f885 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -1,5 +1,5 @@ from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr +from devito.petsc.types import LinearSolveExpr, CallbackExpr @timed_pass() @@ -25,3 +25,20 @@ def petsc_lift(clusters): else: processed.append(c) return processed + + +@timed_pass() +def petsc_project(clusters): + """ + Drop time loop for clusters which appear in PETSc callback functions. + """ + processed = [] + for c in clusters: + if isinstance(c.exprs[0].rhs, CallbackExpr): + time_dims = [d for d in c.ispace.intervals.dimensions if d.is_Time] + ispace = c.ispace.project(lambda d: d not in time_dims) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + + return processed diff --git a/devito/petsc/iet.py b/devito/petsc/iet.py deleted file mode 100644 index 659ac9f94d..0000000000 --- a/devito/petsc/iet.py +++ /dev/null @@ -1,372 +0,0 @@ -from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (FindNodes, Call, - Transformer, FindSymbols, - MapNodes, Iteration, Callable, Callback, List, Uxreplace, - Definition, BlankLine, PointerCast) -<<<<<<< HEAD -from devito.petsc.types import (PetscMPIInt, PETScStruct, DM, Mat, -======= -from devito.petsc import (PetscMPIInt, PETScStruct, DM, Mat, ->>>>>>> b22824875 (fix circular imports) - Vec, KSP, PC, SNES, PetscErrorCode, PETScArray) -from devito.symbolics import Byref, Macro, FieldFromPointer -import cgen as c -from devito.petsc.nodes import MatVecAction, RHSLinearSystem, LinearSolverExpression - - -@iet_pass -def lower_petsc(iet, **kwargs): - - # Check if PETScSolve was used. - petsc_nodes = FindNodes(MatVecAction).visit(iet) - - if not petsc_nodes: - return iet, {} - - else: - # Collect all petsc solution fields - unique_targets = list(set([i.expr.rhs.target for i in petsc_nodes])) - - # Initalize PETSc - init = init_petsc(**kwargs) - - # Create context data struct - struct = build_struct(iet) - - objs = build_core_objects(unique_targets[-1], **kwargs) - - # Create core PETSc calls (not specific to each PETScSolve) - core = core_petsc(unique_targets[-1], struct, objs, **kwargs) - - matvec_mapper = MapNodes(Iteration, MatVecAction, 'groupby').visit(iet) - - main_mapper = {} - - setup = [] - efuncs = [] - - for target in unique_targets: - - solver_objs = build_solver_objs(target) - - matvec_body_list = List() - - solver_setup = False - - for iter, (matvec,) in matvec_mapper.items(): - - # Skip the MatVecAction if it is not associated with the target - # There will be more than one MatVecAction associated with the target - # e.g interior matvec + BC matvecs - if matvec.expr.rhs.target != target: - continue - - # Only need to generate solver setup once per target - if not solver_setup: - solver = generate_solver_calls(solver_objs, objs, matvec, target) - setup.extend(solver) - solver_setup = True - - matvec_body = matvec_body_list._rebuild(body=[ - matvec_body_list.body, iter[0]]) - matvec_body_list = matvec_body_list._rebuild(body=matvec_body) - - main_mapper.update({iter[0]: None}) - - # Create the matvec callback and operation for each target - matvec_callback, matvec_op = create_matvec_callback( - target, matvec_body_list, solver_objs, objs, - struct) - - setup.append(matvec_op) - setup.append(BlankLine) - efuncs.append(matvec_callback) - - # Remove the LinSolveExpr from iet and efuncs that were used to carry - # metadata e.g solver_parameters - main_mapper.update(rebuild_expr_mapper(iet)) - efunc_mapper = {efunc: rebuild_expr_mapper(efunc) for efunc in efuncs} - - iet = Transformer(main_mapper).visit(iet) - efuncs = [Transformer(efunc_mapper[efunc]).visit(efunc) for efunc in efuncs] - - # Replace symbols appearing in each efunc with a pointer to the PETScStruct - efuncs = transform_efuncs(efuncs, struct) - - body = iet.body._rebuild(init=init, body=core + tuple(setup) + iet.body.body) - iet = iet._rebuild(body=body) - - return iet, {'efuncs': efuncs} - - -def init_petsc(**kwargs): - - # Initialize PETSc -> for now, assuming all solver options have to be - # specifed via the parameters dict in PETScSolve - # TODO: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options - initialize = Call(petsc_call, [ - Call('PetscInitialize', arguments=[Null, Null, Null, Null])]) - - return tuple([petsc_func_begin_user, initialize]) - - -def build_struct(iet): - # Place all context data required by the shell routines - # into a PETScStruct - usr_ctx = [] - - basics = FindSymbols('basics').visit(iet) - avoid = FindSymbols('dimensions|indexedbases').visit(iet) - usr_ctx.extend(data for data in basics if data not in avoid) - - return PETScStruct('ctx', usr_ctx) - - -def core_petsc(target, struct, objs, **kwargs): - - # MPI - call_mpi = Call(petsc_call_mpi, [Call('MPI_Comm_size', - arguments=[objs['comm'], - Byref(objs['size'])])]) - # Create DMDA - dmda = create_dmda(target, objs) - dm_setup = Call(petsc_call, [ - Call('DMSetUp', arguments=[objs['da']])]) - dm_app_ctx = Call(petsc_call, [ - Call('DMSetApplicationContext', arguments=[objs['da'], struct])]) - dm_mat_type = Call(petsc_call, [ - Call('DMSetMatType', arguments=[objs['da'], 'MATSHELL'])]) - - return tuple([petsc_func_begin_user, call_mpi, dmda, dm_setup, - dm_app_ctx, dm_mat_type, BlankLine]) - - -def build_core_objects(target, **kwargs): - - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' - - return {'da': DM(name='da', liveness='eager'), - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'err': PetscErrorCode(name='err')} - - -def create_dmda(target, objs): - - args = [objs['comm']] - - args += ['DM_BOUNDARY_GHOSTED' for _ in range(len(target.space_dimensions))] - - # Stencil type - if len(target.space_dimensions) > 1: - args += ['DMDA_STENCIL_BOX'] - - # Global dimensions - args += list(target.shape_global)[::-1] - - # No.of processors in each dimension - if len(target.space_dimensions) > 1: - args += list(target.grid.distributor.topology)[::-1] - - args += [1, target.space_order] - - args += [Null for _ in range(len(target.space_dimensions))] - - args += [Byref(objs['da'])] - - dmda = Call(petsc_call, [ - Call(f'DMDACreate{len(target.space_dimensions)}d', arguments=args)]) - - return dmda - - -def build_solver_objs(target): - - return {'Jac': Mat(name='J_'+str(target.name)), - 'x_global': Vec(name='x_global_'+str(target.name)), - 'x_local': Vec(name='x_local_'+str(target.name), liveness='eager'), - 'b_global': Vec(name='b_global_'+str(target.name)), - 'b_local': Vec(name='b_local_'+str(target.name), liveness='eager'), - 'ksp': KSP(name='ksp_'+str(target.name)), - 'pc': PC(name='pc_'+str(target.name)), - 'snes': SNES(name='snes_'+str(target.name)), - 'X_global': Vec(name='X_global_'+str(target.name)), - 'Y_global': Vec(name='Y_global_'+str(target.name)), - 'X_local': Vec(name='X_local_'+str(target.name), liveness='eager'), - 'Y_local': Vec(name='Y_local_'+str(target.name), liveness='eager') - } - - -def generate_solver_calls(solver_objs, objs, matvec, target): - - solver_params = matvec.expr.rhs.solver_parameters - - snes_create = Call(petsc_call, [Call('SNESCreate', arguments=[ - objs['comm'], Byref(solver_objs['snes'])])]) - - snes_set_dm = Call(petsc_call, [Call('SNESSetDM', arguments=[ - solver_objs['snes'], objs['da']])]) - - create_matrix = Call(petsc_call, [Call('DMCreateMatrix', arguments=[ - objs['da'], Byref(solver_objs['Jac'])])]) - - # NOTE: Assumming all solves are linear for now. - snes_set_type = Call(petsc_call, [Call('SNESSetType', arguments=[ - solver_objs['snes'], 'SNESKSPONLY'])]) - - global_x = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['x_global'])])]) - - local_x = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['x_local'])])]) - - global_b = Call(petsc_call, [Call('DMCreateGlobalVector', arguments=[ - objs['da'], Byref(solver_objs['b_global'])])]) - - local_b = Call(petsc_call, [Call('DMCreateLocalVector', arguments=[ - objs['da'], Byref(solver_objs['b_local'])])]) - - snes_get_ksp = Call(petsc_call, [Call('SNESGetKSP', arguments=[ - solver_objs['snes'], Byref(solver_objs['ksp'])])]) - - vec_replace_array = Call(petsc_call, [ - Call('VecReplaceArray', arguments=[solver_objs[ - 'x_local'], FieldFromPointer(target._C_field_data, target._C_symbol)])]) - - ksp_set_tols = Call(petsc_call, [Call('KSPSetTolerances', arguments=[ - solver_objs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], - solver_params['ksp_divtol'], solver_params['ksp_max_it']])]) - - ksp_set_type = Call(petsc_call, [Call('KSPSetType', arguments=[ - solver_objs['ksp'], linear_solver_mapper[solver_params['ksp_type']]])]) - - ksp_get_pc = Call(petsc_call, [Call('KSPGetPC', arguments=[ - solver_objs['ksp'], Byref(solver_objs['pc'])])]) - - pc_set_type = Call(petsc_call, [Call('PCSetType', arguments=[ - solver_objs['pc'], linear_solver_mapper[solver_params['pc_type']]])]) - - ksp_set_from_ops = Call(petsc_call, [Call('KSPSetFromOptions', arguments=[ - solver_objs['ksp']])]) - - return tuple([snes_create, snes_set_dm, create_matrix, snes_set_type, - global_x, local_x, global_b, local_b, snes_get_ksp, - vec_replace_array, ksp_set_tols, ksp_set_type, ksp_get_pc, - pc_set_type, ksp_set_from_ops]) - - -def create_matvec_callback(target, body, solver_objs, objs, struct): - - # There will be 2 PETScArrays within the body - petsc_arrays = [i.function for i in FindSymbols('indexedbases').visit(body) - if isinstance(i.function, PETScArray)] - - # There will only be one PETScArray that is written to within this body and - # one PETScArray which corresponds to the 'seed' vector - petsc_arr_write, = FindSymbols('writes').visit(body) - petsc_arr_seed, = [i for i in petsc_arrays if i.function != petsc_arr_write.function] - - # Struct needs to be defined explicitly here since CompositeObjects - # do not have 'liveness' - defn_struct = Definition(struct) - - mat_get_dm = Call(petsc_call, [Call('MatGetDM', arguments=[ - solver_objs['Jac'], Byref(objs['da'])])]) - - dm_get_app_context = Call(petsc_call, [Call('DMGetApplicationContext', arguments=[ - objs['da'], Byref(struct._C_symbol)])]) - - dm_get_local_xvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ - objs['da'], Byref(solver_objs['X_local'])])]) - - global_to_local_begin = Call(petsc_call, [Call('DMGlobalToLocalBegin', arguments=[ - objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) - - global_to_local_end = Call(petsc_call, [Call('DMGlobalToLocalEnd', arguments=[ - objs['da'], solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local']])]) - - dm_get_local_yvec = Call(petsc_call, [Call('DMGetLocalVector', arguments=[ - objs['da'], Byref(solver_objs['Y_local'])])]) - - vec_get_array_y = Call(petsc_call, [Call('VecGetArray', arguments=[ - solver_objs['Y_local'], Byref(petsc_arr_write.function)])]) - - vec_get_array_x = Call(petsc_call, [Call('VecGetArray', arguments=[ - solver_objs['X_local'], Byref(petsc_arrays[0])])]) - - dm_get_local_info = Call(petsc_call, [Call('DMDAGetLocalInfo', arguments=[ - objs['da'], Byref(petsc_arrays[0].function.dmda_info)])]) - - casts = [PointerCast(i.function) for i in petsc_arrays] - - vec_restore_array_y = Call(petsc_call, [Call('VecRestoreArray', arguments=[ - solver_objs['Y_local'], Byref(petsc_arr_write[0])])]) - - vec_restore_array_x = Call(petsc_call, [Call('VecRestoreArray', arguments=[ - solver_objs['X_local'], Byref(petsc_arr_seed)])]) - - dm_local_to_global_begin = Call(petsc_call, [Call('DMLocalToGlobalBegin', arguments=[ - objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) - - dm_local_to_global_end = Call(petsc_call, [Call('DMLocalToGlobalEnd', arguments=[ - objs['da'], solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global']])]) - - func_return = Call('PetscFunctionReturn', arguments=[0]) - - matvec_body = List(body=[ - petsc_func_begin_user, defn_struct, mat_get_dm, dm_get_app_context, - dm_get_local_xvec, global_to_local_begin, global_to_local_end, - dm_get_local_yvec, vec_get_array_y, vec_get_array_x, dm_get_local_info, - casts, BlankLine, body, vec_restore_array_y, vec_restore_array_x, - dm_local_to_global_begin, dm_local_to_global_end, func_return]) - - matvec_callback = Callable( - 'MyMatShellMult_'+str(target.name), matvec_body, retval=objs['err'], - parameters=(solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'])) - - matvec_operation = Call(petsc_call, [ - Call('MatShellSetOperation', arguments=[ - solver_objs['Jac'], 'MATOP_MULT', Callback(matvec_callback.name, - Void, Void)])]) - - return matvec_callback, matvec_operation - - -def rebuild_expr_mapper(callable): - - return {expr: expr._rebuild( - expr=expr.expr._rebuild(rhs=expr.expr.rhs.expr)) for - expr in FindNodes(LinearSolverExpression).visit(callable)} - - -def transform_efuncs(efuncs, struct): - - efuncs_new = [] - for efunc in efuncs: - new_body = efunc.body - for i in struct.usr_ctx: - new_body = Uxreplace({i: FieldFromPointer(i, struct)}).visit(new_body) - efunc_with_new_body = efunc._rebuild(body=new_body) - efuncs_new.append(efunc_with_new_body) - - return efuncs_new - - -Null = Macro('NULL') -Void = 'void' - -petsc_call = 'PetscCall' -petsc_call_mpi = 'PetscCallMPI' -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') - -linear_solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index a7855fbb36..adcf709eab 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -10,10 +10,10 @@ def petsc_call_mpi(specific_call, call_args): return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)]) -def petsc_struct(name, fields, liveness='lazy'): +def petsc_struct(name, fields, pname, liveness='lazy'): # TODO: Fix this circular import from devito.petsc.types.object import PETScStruct - return PETScStruct(name=name, pname='MatContext', + return PETScStruct(name=name, pname=pname, fields=fields, liveness=liveness) diff --git a/devito/petsc/types.py b/devito/petsc/types.py deleted file mode 100644 index 72a0faa1e8..0000000000 --- a/devito/petsc/types.py +++ /dev/null @@ -1,429 +0,0 @@ -import sympy - -<<<<<<< HEAD -from devito.tools import Reconstructable, sympy_mutex -======= -from devito.tools import CustomDtype -from devito.types import LocalObject, Eq, CompositeObject -from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic -from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction, AbstractSymbol -from devito.finite_differences.tools import fd_weights_registry -from devito.tools import dtype_to_ctype, Reconstructable, sympy_mutex -from devito.symbolics import FieldFromComposite, Byref - -from devito.petsc.utils import petsc_call - - -class DM(LocalObject): - """ - PETSc Data Management object (DM). - """ - dtype = CustomDtype('DM') - - def __init__(self, *args, stencil_width=None, **kwargs): - super().__init__(*args, **kwargs) - self._stencil_width = stencil_width - - @property - def stencil_width(self): - return self._stencil_width - - @property - def info(self): - return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('DMDestroy', [Byref(self.function)]) - - -class Mat(LocalObject): - """ - PETSc Matrix object (Mat). - """ - dtype = CustomDtype('Mat') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('MatDestroy', [Byref(self.function)]) - - -class LocalVec(LocalObject): - """ - PETSc Vector object (Vec). - """ - dtype = CustomDtype('Vec') - - -class GlobalVec(LocalObject): - """ - PETSc Vector object (Vec). - """ - dtype = CustomDtype('Vec') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('VecDestroy', [Byref(self.function)]) - - -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 SNES(LocalObject): - """ - PETSc SNES : Non-Linear Systems Solvers. - """ - dtype = CustomDtype('SNES') - - @property - def _C_free(self): - # from devito.petsc.utils import petsc_call - return petsc_call('SNESDestroy', [Byref(self.function)]) - - -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 DMDALocalInfo(LocalObject): - """ - PETSc object - C struct containing information - about the local grid. - """ - dtype = CustomDtype('DMDALocalInfo') - - -class PetscErrorCode(LocalObject): - """ - PETSc datatype used to return PETSc error codes. - https://petsc.org/release/manualpages/Sys/PetscErrorCode/ - """ - dtype = CustomDtype('PetscErrorCode') - - -class DummyArg(LocalObject): - dtype = CustomDtype('void', modifier='*') - - -class PETScArray(ArrayBasic, Differentiable): - """ - 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. - Differentiable enables compatability with standard Function objects, - allowing for the use of the `subs` method. - TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) - and then PETScArray(ArrayBasic). - """ - - _data_alignment = False - - # Default method for the finite difference approximation weights computation. - _default_fd = 'taylor' - - __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients', - 'space_order')) - - def __init_finalize__(self, *args, **kwargs): - - super().__init_finalize__(*args, **kwargs) - - # Symbolic (finite difference) coefficients - self._coefficients = kwargs.get('coefficients', self._default_fd) - if self._coefficients not in fd_weights_registry: - raise ValueError("coefficients must be one of %s" - " not %s" % (str(fd_weights_registry), self._coefficients)) - self._shape = kwargs.get('shape') - self._space_order = kwargs.get('space_order', 1) - - @classmethod - def __dtype_setup__(cls, **kwargs): - return kwargs.get('dtype', np.float32) - - @property - def coefficients(self): - """Form of the coefficients of the function.""" - return self._coefficients - - @property - def shape(self): - return self._shape - - @property - def space_order(self): - return self._space_order - - @cached_property - def _shape_with_inhalo(self): - """ - Shape of the domain+inhalo region. The inhalo region comprises the - outhalo as well as any additional "ghost" layers for MPI halo - exchanges. Data in the inhalo region are exchanged when running - Operators to maintain consistent values as in sequential runs. - - Notes - ----- - Typically, this property won't be used in user code, but it may come - in handy for testing or debugging - """ - return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) - - @cached_property - def shape_allocated(self): - """ - Shape of the allocated data of the Function type object from which - this PETScArray was derived. It includes the domain and inhalo regions, - as well as any additional padding surrounding the halo. - - Notes - ----- - In an MPI context, this is the *local* with_halo region shape. - """ - return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, - self._padding)], - getters=self.dimensions) - - @cached_property - def _C_ctype(self): - # NOTE: Reverting to using float/double instead of PetscScalar for - # simplicity when opt='advanced'. Otherwise, Temp objects must also - # be converted to PetscScalar. Additional tests are needed to - # ensure this approach is fine. Previously, issues arose from - # mismatches between precision of Function objects in Devito and the - # precision of the PETSc configuration. - # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables - # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of - # the user's PETSc configuration. - return POINTER(dtype_to_ctype(self.dtype)) - - @property - def symbolic_shape(self): - field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] - # Reverse it since DMDA is setup backwards to Devito dimensions. - return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - - @property - def dmda(self): - name = 'da_so_%s' % self.space_order - return DM(name=name, liveness='eager', stencil_width=self.space_order) - - -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] - - -class InjectSolveEq(Eq): - pass ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) - - -class LinearSolveExpr(sympy.Function, Reconstructable): - - __rargs__ = ('expr',) -<<<<<<< HEAD -<<<<<<< HEAD - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays') -======= - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', 'formfuncs', 'formrhs') ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays') ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - - defaults = { - 'ksp_type': 'gmres', - 'pc_type': 'jacobi', - 'ksp_rtol': 1e-7, # Relative tolerance - 'ksp_atol': 1e-50, # Absolute tolerance - 'ksp_divtol': 1e4, # Divergence tolerance - 'ksp_max_it': 10000 # Maximum iterations - } - - def __new__(cls, expr, target=None, solver_parameters=None, -<<<<<<< HEAD -<<<<<<< HEAD - matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): -======= - matvecs=None, formfuncs=None, formrhs=None, **kwargs): ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - - if solver_parameters is None: - solver_parameters = cls.defaults - else: - for key, val in cls.defaults.items(): - solver_parameters[key] = solver_parameters.get(key, val) - - with sympy_mutex: - obj = sympy.Basic.__new__(cls, expr) - obj._expr = expr - obj._target = target - obj._solver_parameters = solver_parameters - obj._matvecs = matvecs - obj._formfuncs = formfuncs - obj._formrhs = formrhs -<<<<<<< HEAD -<<<<<<< HEAD - obj._arrays = arrays -======= ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - obj._arrays = arrays ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - return obj - - def __repr__(self): - return "%s(%s)" % (self.__class__.__name__, self.expr) - - __str__ = __repr__ - - def _sympystr(self, printer): - return str(self) - - def __hash__(self): - return hash(self.expr) - - def __eq__(self, other): - return (isinstance(other, LinearSolveExpr) and - self.expr == other.expr and - self.target == other.target) - - @property - def expr(self): - return self._expr - - @property - def target(self): - return self._target - - @property - def solver_parameters(self): - return self._solver_parameters - - @property - def matvecs(self): - return self._matvecs - - @property - def formfuncs(self): - return self._formfuncs - - @property - def formrhs(self): - return self._formrhs - -<<<<<<< HEAD -<<<<<<< HEAD - @property - def formrhs(self): - return self._formrhs -======= -======= - @property - def arrays(self): - return self._arrays - ->>>>>>> b50c18397 (compiler: Clean up access to petsc arrays in each callback) - func = Reconstructable._rebuild - - -class PETScStruct(CompositeObject): - - __rargs__ = ('name', 'usr_ctx') - __rkwargs__ = ('liveness',) - - def __init__(self, name, usr_ctx, liveness='lazy'): - pfields = [(i._C_name, dtype_to_ctype(i.dtype)) - for i in usr_ctx if isinstance(i, AbstractSymbol)] - self._usr_ctx = usr_ctx - super().__init__(name, 'MatContext', pfields) ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) - - assert liveness in ['eager', 'lazy'] - self._liveness = liveness - - @property - def arrays(self): - return self._arrays - -<<<<<<< HEAD - func = Reconstructable._rebuild -======= - def _arg_values(self, **kwargs): - # TODO: Handle passing `struct dataobj *restrict g1_vec` and similar - # elements within in the struct. It seemed necessary to convert the Numpy array - # into a pointer to ctypes.Struct within MatContext - # (i.e override _arg_finalize for PETScStruct here?), but this approach seemed - # to result in a NIL pointer in the generated code ... any suggestions? - # Alternative: Allow these pointers to appear in the main kernel arguments - # as usual and use ctx->g1_vec = g1_vec; to assign them to the struct? - values = super()._arg_values(**kwargs) - for i in self.usr_ctx: - setattr(values[self.name]._obj, i.name, kwargs['args'][i.name]) - return values - - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' -<<<<<<< HEAD -<<<<<<< HEAD ->>>>>>> 29ada0831 (compiler: form rhs in callback function and remove manual petsc casts) -======= - -======= - ->>>>>>> 9d611239f (compiler: Fix derive_callback_inputs) - @property - def fields(self): - return self._usr_ctx ->>>>>>> 41157b0bc (compiler: Place dataobj pointers inside ctx struct) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 9f8cbe4cbb..1bcfb3a6cf 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -9,7 +9,9 @@ class DM(LocalObject): """ - PETSc Data Management object (DM). + PETSc Data Management object (DM). This is the primary DM instance + created within the main kernel and linked to the SNES + solver using `SNESSetDM`. """ dtype = CustomDtype('DM') @@ -21,10 +23,6 @@ def __init__(self, *args, stencil_width=None, **kwargs): def stencil_width(self): return self._stencil_width - @property - def info(self): - return DMDALocalInfo(name='%s_info' % self.name, liveness='eager') - @property def _C_free(self): return petsc_call('DMDestroy', [Byref(self.function)]) @@ -34,6 +32,22 @@ def _C_free_priority(self): return 3 +class CallbackDM(LocalObject): + """ + PETSc Data Management object (DM). This is the DM instance + accessed within the callback functions via `SNESGetDM`. + """ + dtype = CustomDtype('DM') + + def __init__(self, *args, stencil_width=None, **kwargs): + super().__init__(*args, **kwargs) + self._stencil_width = stencil_width + + @property + def stencil_width(self): + return self._stencil_width + + class Mat(LocalObject): """ PETSc Matrix object (Mat). @@ -51,14 +65,18 @@ def _C_free_priority(self): class LocalVec(LocalObject): """ - PETSc Vector object (Vec). + PETSc local vector object (Vec). + A local vector has ghost locations that contain values that are + owned by other MPI ranks. """ dtype = CustomDtype('Vec') class GlobalVec(LocalObject): """ - PETSc Vector object (Vec). + PETSc global vector object (Vec). + A global vector is a parallel vector that has no duplicate values + between MPI ranks. A global vector has no ghost locations. """ dtype = CustomDtype('Vec') @@ -142,6 +160,10 @@ class PetscErrorCode(LocalObject): class DummyArg(LocalObject): + """ + A void pointer used to satisfy the function + signature of the `FormFunction` callback. + """ dtype = CustomDtype('void', modifier='*') @@ -160,9 +182,21 @@ def fields(self): @property def time_dim_fields(self): + """ + Fields within the struct that are updated during the time loop. + These are not set in the `PopulateMatContext` callback. + """ return [f for f in self.fields if isinstance(f, (ModuloDimension, TimeDimension))] + @property + def callback_fields(self): + """ + Fields within the struct that are initialized in the `PopulateMatContext` + callback. These fields are not updated in the time loop. + """ + return [f for f in self.fields if f not in self.time_dim_fields] + @property def _C_ctype(self): return POINTER(self.dtype) if self.liveness == \ diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index eda2fa40d4..1a4a778c9e 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -4,23 +4,53 @@ class LinearSolveExpr(sympy.Function, Reconstructable): + """ + A symbolic expression passed through the Operator, containing the metadata + needed to execute a linear solver. Linear problems are handled with + `SNESSetType(snes, KSPONLY)`, enabling a unified interface for both + linear and nonlinear solvers. + + # TODO: extend this + defaults: + - 'ksp_type': String with the name of the PETSc Krylov method. + Default is 'gmres' (Generalized Minimal Residual Method). + https://petsc.org/main/manualpages/KSP/KSPType/ + + - 'pc_type': String with the name of the PETSc preconditioner. + Default is 'jacobi' (i.e diagonal scaling preconditioning). + https://petsc.org/main/manualpages/PC/PCType/ + + KSP tolerances: + https://petsc.org/release/manualpages/KSP/KSPSetTolerances/ + + - 'ksp_rtol': Relative convergence tolerance. Default + is 1e-5. + - 'ksp_atol': Absolute convergence for tolerance. Default + is 1e-50. + - 'ksp_divtol': Divergence tolerance, amount residual norm can + increase before `KSPConvergedDefault()` concludes + that the method is diverging. Default is 1e5. + - 'ksp_max_it': Maximum number of iterations to use. Default + is 1e4. + """ __rargs__ = ('expr',) __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays', 'time_mapper') + 'formfuncs', 'formrhs', 'arrays', 'time_mapper', + 'localinfo') defaults = { 'ksp_type': 'gmres', 'pc_type': 'jacobi', - 'ksp_rtol': 1e-7, # Relative tolerance + 'ksp_rtol': 1e-5, # Relative tolerance 'ksp_atol': 1e-50, # Absolute tolerance - 'ksp_divtol': 1e4, # Divergence tolerance - 'ksp_max_it': 10000 # Maximum iterations + 'ksp_divtol': 1e5, # Divergence tolerance + 'ksp_max_it': 1e4 # Maximum iterations } def __new__(cls, expr, target=None, solver_parameters=None, matvecs=None, formfuncs=None, formrhs=None, - arrays=None, time_mapper=None, **kwargs): + arrays=None, time_mapper=None, localinfo=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -39,6 +69,7 @@ def __new__(cls, expr, target=None, solver_parameters=None, obj._formrhs = formrhs obj._arrays = arrays obj._time_mapper = time_mapper + obj._localinfo = localinfo return obj def __repr__(self): @@ -89,6 +120,10 @@ def arrays(self): def time_mapper(self): return self._time_mapper + @property + def localinfo(self): + return self._localinfo + @classmethod def eval(cls, *args): return None diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py index e7b395d436..2c366a389e 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -242,9 +242,6 @@ def _print_Float(self, expr): def _print_Differentiable(self, expr): return "(%s)" % self._print(expr._expr) - def _print_PETScRHS(self, expr): - return "%s" % self._print(expr._expr) - _print_EvalDerivative = C99CodePrinter._print_Add def _print_CallFromPointer(self, expr): diff --git a/devito/types/basic.py b/devito/types/basic.py index 14e3949fb8..660ceff437 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1513,6 +1513,37 @@ def __getnewargs_ex__(self): return args, kwargs +class LocalType(Basic): + """ + This is the abstract base class for local types, which are + generated by the compiler in C rather than in Python. + Notes + ----- + Subclasses should setup `_liveness`. + """ + + is_LocalType = True + + @property + def liveness(self): + return self._liveness + + @property + def _mem_internal_eager(self): + return self._liveness == 'eager' + + @property + def _mem_internal_lazy(self): + return self._liveness == 'lazy' + + """ + A modifier added to the subclass C declaration when it appears + in a function signature. For example, a subclass might define `_C_modifier = '&'` + to impose pass-by-reference semantics. + """ + _C_modifier = None + + # Extended SymPy hierarchy follows, for essentially two reasons: # - To keep track of `function` # - To override SymPy caching behaviour diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc deleted file mode 100644 index 782e0ac340..0000000000 --- a/docker/Dockerfile.petsc +++ /dev/null @@ -1,49 +0,0 @@ -############################################################## -# This Dockerfile builds a base image to run Devito with PETSc -############################################################# - -# Base image -FROM ubuntu:22.04 as base - -ENV DEBIAN_FRONTEND noninteractive - -# Install python -RUN apt-get update && \ - apt-get install -y dh-autoreconf python3-venv python3-dev python3-pip libopenblas-serial-dev git pkgconf mpich - -# Install for basic base not containing it -RUN apt-get install -y vim wget git flex libnuma-dev tmux \ - numactl hwloc curl \ - autoconf libtool build-essential procps - -# Install tmpi -RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi - -# Install OpenGL library, necessary for the installation of GemPy -RUN apt-get install -y libgl1-mesa-glx - -# Install numpy from source -RUN python3 -m venv /venv && \ - /venv/bin/pip install --no-cache-dir --upgrade pip && \ - /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ - rm -rf ~/.cache/pip - -# Create a directory for PETSc -RUN mkdir -p /home/petsc - -# Clone and install PETSc -RUN cd /home/petsc \ - && git clone -b release https://gitlab.com/petsc/petsc.git petsc \ - && cd petsc \ - && git pull \ - && ./configure --with-fortran-bindings=0 --with-openblas-include=$(pkgconf --variable=includedir openblas) --with-openblas-lib=$(pkgconf --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build \ - && make all - -RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -EXPOSE 8888 -CMD ["/bin/bash"] - -ENV PETSC_ARCH="devito_build" -ENV PETSC_DIR="/home/petsc/petsc" From 82b434b4da36c95a13cc2458bdc1e282e7249404 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 3 Jan 2025 11:54:09 +0000 Subject: [PATCH 125/151] types: Edit freesymbols in FieldFromPointer --- devito/ir/equations/algorithms.py | 7 +++++-- devito/operator/operator.py | 1 + tests/test_dimension.py | 25 ++++++++++++++++++++++++- tests/test_subdomains.py | 3 ++- 4 files changed, 32 insertions(+), 4 deletions(-) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 946eb4324b..2a52b1a41d 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -169,10 +169,10 @@ def concretize_subdims(exprs, **kwargs): across `exprs`, such as the thickness symbols. """ sregistry = kwargs.get('sregistry') + mapper = kwargs.get('concretize_mapper') - mapper = {} rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions - + # from IPython import embed; embed() _concretize_subdims(exprs, mapper, rebuilt, sregistry) if not mapper: return exprs @@ -206,6 +206,7 @@ def _(v, mapper, rebuilt, sregistry): @_concretize_subdims.register(Eq) def _(expr, mapper, rebuilt, sregistry): + # from IPython import embed; embed() for d in expr.free_symbols: _concretize_subdims(d, mapper, rebuilt, sregistry) @@ -215,6 +216,7 @@ def _(expr, mapper, rebuilt, sregistry): @_concretize_subdims.register(Thickness) def _(tkn, mapper, rebuilt, sregistry): + # from IPython import embed; embed() if tkn in mapper: # Already have a substitution for this thickness return @@ -224,6 +226,7 @@ def _(tkn, mapper, rebuilt, sregistry): @_concretize_subdims.register(SubDimension) def _(d, mapper, rebuilt, sregistry): + # from IPython import embed; embed() if d in mapper: # Already have a substitution for this dimension return diff --git a/devito/operator/operator.py b/devito/operator/operator.py index f3f4e8565b..73fb6f1450 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -263,6 +263,7 @@ def _lower(cls, expressions, **kwargs): """ # Create a symbol registry kwargs.setdefault('sregistry', SymbolRegistry()) + kwargs.setdefault('concretize_mapper', {}) expressions = as_tuple(expressions) diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 5ed0bd08a5..0132871a47 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -2064,7 +2064,8 @@ def test_correct_thicknesses(self): ix1 = SubDimension.right('x', x, 2) ix2 = SubDimension.middle('x', x, 2, 2) - rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry()) + rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry(), + concretize_mapper={}) assert rebuilt[0].is_left assert rebuilt[1].is_right @@ -2105,3 +2106,25 @@ def test_condition_concretization(self): assert len({c.condition for c in conditionals}) == 2 for c in conditionals: assert c.condition in expected_conditionals + + def test_repeat_concretization(self): + """ + Ensure that SubDimensions are consistently concretized to the same object + across multiple calls to the function. This is necessary when using + `rcompile` on equations with SubDimensions. + """ + + grid = Grid((2, 2)) + + x = Dimension('x') + ix = SubDimension.middle('ix', x, 2, 2) + + u = Function(name='u', grid=grid) + eq = Eq(u, ix.ltkn) + + kwargs = {'sregistry': SymbolRegistry(), 'concretize_mapper': {}} + + exprs1 = concretize_subdims([eq], **kwargs) + exprs2 = concretize_subdims([eq], **kwargs) + + assert exprs1 == exprs2 diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index 784dc1c3ad..db08f60de6 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -37,7 +37,8 @@ def define(self, dimensions): with timed_region('x'): # _lower_exprs expects a SymbolRegistry, so create one expr = Operator._lower_exprs([eq0], options={}, - sregistry=SymbolRegistry())[0] + sregistry=SymbolRegistry(), + concretize_mapper={})[0] assert str(expr.rhs) == 'ix*f[ix + 1, iy + 1] + iy' def test_multiple_middle(self): From 8cb18d3bc924ce6572d602331880b937f6364ca6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 9 Jan 2025 09:28:21 +0000 Subject: [PATCH 126/151] compiler: Use classes --- devito/petsc/types/array.py | 48 +++++++++++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 10 deletions(-) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index a150ea3247..db955dde43 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -9,8 +9,9 @@ from devito.finite_differences.tools import fd_weights_registry from devito.tools import dtype_to_ctype from devito.symbolics import FieldFromComposite +from devito.types import Symbol -from .object import DM +from .object import DM, DMDALocalInfo class PETScArray(ArrayBasic, Differentiable): @@ -30,7 +31,7 @@ class PETScArray(ArrayBasic, Differentiable): __rkwargs__ = (AbstractFunction.__rkwargs__ + ('dimensions', 'shape', 'liveness', 'coefficients', - 'space_order')) + 'space_order', 'symbolic_shape')) def __init_finalize__(self, *args, **kwargs): @@ -43,6 +44,7 @@ def __init_finalize__(self, *args, **kwargs): " not %s" % (str(fd_weights_registry), self._coefficients)) self._shape = kwargs.get('shape') self._space_order = kwargs.get('space_order', 1) + self._symbolic_shape = kwargs.get('symbolic_shape') @classmethod def __dtype_setup__(cls, **kwargs): @@ -104,14 +106,40 @@ def _C_ctype(self): # the user's PETSc configuration. return POINTER(dtype_to_ctype(self.dtype)) + # @property + # def symbolic_shape(self): + # # from IPython import embed; embed() + # if self.dmda: + # field_from_composites = [ + # FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] + # # Reverse it since DMDA is setup backwards to Devito dimensions. + # return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + # else: + # return DimensionTuple(*self.dimensions) + + # @property + # def symbolic_shape(self): + # # info = DMDALocalInfo(name='info', liveness='eager') + + # field_from_composites = [ + # FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] + + # # Reverse it since DMDA is setup backwards to Devito dimensions. + # return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + + # @cached_property + # def dmda(self): + # name = 'da_so_%s' % self.space_order + # return DM(name=name, liveness='eager', stencil_width=self.space_order) + + # @cached_property + # def dmda(self): + # return self._dmda + @property def symbolic_shape(self): - field_from_composites = [ - FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] - # Reverse it since DMDA is setup backwards to Devito dimensions. - return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + return self._symbolic_shape - @cached_property - def dmda(self): - name = 'da_so_%s' % self.space_order - return DM(name=name, liveness='eager', stencil_width=self.space_order) + # @property + # def dmda_info(self): + # return DMDALocalInfo(name='info', liveness='eager') From da2bbf38c2da2467f6e36030a5aaff7c330222e1 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 9 Jan 2025 14:15:42 +0000 Subject: [PATCH 127/151] compiler: Switch to using a separate struct and dm for each SNES solve --- tests/test_petsc.py | 90 +++++++++++++++++++++------------------------ 1 file changed, 42 insertions(+), 48 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 1c70331066..45706c5806 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -129,12 +129,12 @@ def test_petsc_solve(): rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) assert str(action_expr[-1].expr.rhs) == \ - 'x_matvec_f[x + 1, y + 2]/matvec->h_x**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/matvec->h_x**2' + \ - ' + x_matvec_f[x + 3, y + 2]/matvec->h_x**2' + \ - ' + x_matvec_f[x + 2, y + 1]/matvec->h_y**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/matvec->h_y**2' + \ - ' + x_matvec_f[x + 2, y + 3]/matvec->h_y**2' + 'x_matvec_f[x + 1, y + 2]/lctx->h_x**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_x**2' + \ + ' + x_matvec_f[x + 3, y + 2]/lctx->h_x**2' + \ + ' + x_matvec_f[x + 2, y + 1]/lctx->h_y**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_y**2' + \ + ' + x_matvec_f[x + 2, y + 3]/lctx->h_y**2' assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' @@ -178,6 +178,8 @@ def test_multiple_petsc_solves(): # One PopulateMatContext for all solves assert len(callable_roots) == 7 + # Check unique dmda is created per SNES solve + @skipif('petsc') def test_petsc_cast(): @@ -188,29 +190,29 @@ def test_petsc_cast(): g1 = Grid((2, 2)) g2 = Grid((2, 2, 2)) - arr0 = PETScArray(name='arr0', dimensions=g0.dimensions, shape=g0.shape) - arr1 = PETScArray(name='arr1', dimensions=g1.dimensions, shape=g1.shape) - arr2 = PETScArray(name='arr2', dimensions=g2.dimensions, shape=g2.shape) + f0 = Function(name='f0', grid=g0) + f1 = Function(name='f1', grid=g1) + f2 = Function(name='f2', grid=g2) - arr3 = PETScArray(name='arr3', dimensions=g1.dimensions, - shape=g1.shape, space_order=4) + arr0 = PETScArray(name='arr0', dimensions=f0.dimensions, + shape=f0.shape, symbolic_shape=f0.symbolic_shape) + arr1 = PETScArray(name='arr1', dimensions=f1.dimensions, + shape=f1.shape, symbolic_shape=f1.symbolic_shape) + arr2 = PETScArray(name='arr2', dimensions=f2.dimensions, + shape=f2.shape, symbolic_shape=f2.symbolic_shape) cast0 = PointerCast(arr0) cast1 = PointerCast(arr1) cast2 = PointerCast(arr2) - cast3 = PointerCast(arr3) assert str(cast0) == \ 'float (*restrict arr0) = (float (*)) arr0_vec;' assert str(cast1) == \ - 'float (*restrict arr1)[da_so_1_info.gxm] = ' + \ - '(float (*)[da_so_1_info.gxm]) arr1_vec;' + 'float (*restrict arr1)[f1_vec->size[1]] = ' + \ + '(float (*)[f1_vec->size[1]]) arr1_vec;' assert str(cast2) == \ - 'float (*restrict arr2)[da_so_1_info.gym][da_so_1_info.gxm] = ' + \ - '(float (*)[da_so_1_info.gym][da_so_1_info.gxm]) arr2_vec;' - assert str(cast3) == \ - 'float (*restrict arr3)[da_so_4_info.gxm] = ' + \ - '(float (*)[da_so_4_info.gxm]) arr3_vec;' + 'float (*restrict arr2)[f2_vec->size[1]][f2_vec->size[2]] = ' + \ + '(float (*)[f2_vec->size[1]][f2_vec->size[2]]) arr2_vec;' @skipif('petsc') @@ -258,23 +260,15 @@ def test_dmda_create(): op3 = Operator(petsc3, opt='noop') assert 'PetscCall(DMDACreate1d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ - '2,1,2,NULL,&(da_so_2)));' in str(op1) + '2,1,2,NULL,&(da_0)));' in str(op1) assert 'PetscCall(DMDACreate2d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_so_4)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_0)));' \ in str(op2) assert 'PetscCall(DMDACreate3d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ - ',1,1,1,1,6,NULL,NULL,NULL,&(da_so_6)));' in str(op3) - - # Check unique DMDA is created per grid, per space_order - f4 = Function(name='f4', grid=grid2, space_order=6) - eqn4 = Eq(f4.laplace, 10) - petsc4 = PETScSolve(eqn4, f4) - with switchconfig(openmp=False): - op4 = Operator(petsc2+petsc2+petsc4, opt='noop') - assert str(op4).count('DMDACreate2d') == 2 + ',1,1,1,1,6,NULL,NULL,NULL,&(da_0)));' in str(op3) @skipif('petsc') @@ -614,31 +608,31 @@ def test_petsc_struct(): assert all(not isinstance(i, CCompositeObject) for i in op.parameters) -@skipif('petsc') -@pytest.mark.parallel(mode=[2, 4, 8]) -def test_apply(mode): +# @skipif('petsc') +# @pytest.mark.parallel(mode=[2, 4, 8]) +# def test_apply(mode): - grid = Grid(shape=(13, 13), dtype=np.float64) +# grid = Grid(shape=(13, 13), dtype=np.float64) - pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) - rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) - mu = Constant(name='mu', value=2.0) +# pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) +# rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) +# mu = Constant(name='mu', value=2.0) - eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) +# eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) - petsc = PETScSolve(eqn, pn) +# petsc = PETScSolve(eqn, pn) - # Build the op - with switchconfig(openmp=False, mpi=True): - op = Operator(petsc) +# # Build the op +# with switchconfig(openmp=False, mpi=True): +# op = Operator(petsc) - # Check the Operator runs without errors. Not verifying output for - # now. Need to consolidate BC implementation - op.apply() +# # Check the Operator runs without errors. Not verifying output for +# # now. Need to consolidate BC implementation +# op.apply() - # Verify that users can override `mu` - mu_new = Constant(name='mu_new', value=4.0) - op.apply(mu=mu_new) +# # Verify that users can override `mu` +# mu_new = Constant(name='mu_new', value=4.0) +# op.apply(mu=mu_new) @skipif('petsc') From c445140357ce61a9b5f779f02e0cad75e3f6e7b4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 9 Jan 2025 16:50:14 +0000 Subject: [PATCH 128/151] compiler: Rework petscarrays --- devito/petsc/types/array.py | 87 ++++++++++++++++++++++--------------- tests/test_petsc.py | 66 ++++++++++++---------------- 2 files changed, 79 insertions(+), 74 deletions(-) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index db955dde43..8b7cb02ad4 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -7,19 +7,26 @@ from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry -from devito.tools import dtype_to_ctype +from devito.tools import dtype_to_ctype, as_tuple from devito.symbolics import FieldFromComposite from devito.types import Symbol +from devito.data import FULL from .object import DM, DMDALocalInfo + class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent a customised variant of ArrayBasic. Differentiable enables compatability with standard Function objects, allowing for the use of the `subs` method. + + PETScArray objects represent vector objects within PETSc. + They correspond to the spatial domain of a Function-like object passed + into PETScSolve from a user. + TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) and then PETScArray(ArrayBasic). """ @@ -30,11 +37,14 @@ class PETScArray(ArrayBasic, Differentiable): _default_fd = 'taylor' __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('dimensions', 'shape', 'liveness', 'coefficients', - 'space_order', 'symbolic_shape')) + ('target', 'liveness', 'coefficients')) def __init_finalize__(self, *args, **kwargs): + self._target = kwargs.get('target') + self._ndim = kwargs['ndim'] = len(self._target.space_dimensions) + self._dimensions = kwargs['dimensions'] = self._target.space_dimensions + super().__init_finalize__(*args, **kwargs) # Symbolic (finite difference) coefficients @@ -42,13 +52,31 @@ def __init_finalize__(self, *args, **kwargs): if self._coefficients not in fd_weights_registry: raise ValueError("coefficients must be one of %s" " not %s" % (str(fd_weights_registry), self._coefficients)) - self._shape = kwargs.get('shape') - self._space_order = kwargs.get('space_order', 1) - self._symbolic_shape = kwargs.get('symbolic_shape') + + @property + def ndim(self): + return self._ndim @classmethod def __dtype_setup__(cls, **kwargs): - return kwargs.get('dtype', np.float32) + return kwargs['target'].dtype + + @classmethod + def __indices_setup__(cls, *args, **kwargs): + dimensions = kwargs['target'].space_dimensions + if args: + indices = args + else: + indices = dimensions + return as_tuple(dimensions), as_tuple(indices) + + @property + def dimensions(self): + return self._dimensions + + @property + def target(self): + return self._target @property def coefficients(self): @@ -57,41 +85,19 @@ def coefficients(self): @property def shape(self): - return self._shape + return self.target.grid.shape @property def space_order(self): - return self._space_order + return self.target.space_order @cached_property def _shape_with_inhalo(self): - """ - Shape of the domain+inhalo region. The inhalo region comprises the - outhalo as well as any additional "ghost" layers for MPI halo - exchanges. Data in the inhalo region are exchanged when running - Operators to maintain consistent values as in sequential runs. - - Notes - ----- - Typically, this property won't be used in user code, but it may come - in handy for testing or debugging - """ - return tuple(j + i + k for i, (j, k) in zip(self.shape, self._halo)) + return self.target.shape_with_inhalo @cached_property def shape_allocated(self): - """ - Shape of the allocated data of the Function type object from which - this PETScArray was derived. It includes the domain and inhalo regions, - as well as any additional padding surrounding the halo. - - Notes - ----- - In an MPI context, this is the *local* with_halo region shape. - """ - return DimensionTuple(*[j + i + k for i, (j, k) in zip(self._shape_with_inhalo, - self._padding)], - getters=self.dimensions) + return self.target.shape_allocated @cached_property def _C_ctype(self): @@ -106,6 +112,16 @@ def _C_ctype(self): # the user's PETSc configuration. return POINTER(dtype_to_ctype(self.dtype)) + @property + def halo(self): + return [self.target.halo[d] for d in self.target.space_dimensions] + + def __halo_setup__(self, **kwargs): + target = kwargs['target'] + halo = [target.halo[d] for d in target.space_dimensions] + # halo = tuple(kwargs.get('halo', ((0, 0),)*self.ndim)) + return DimensionTuple(*halo, getters=target.space_dimensions) + # @property # def symbolic_shape(self): # # from IPython import embed; embed() @@ -138,8 +154,9 @@ def _C_ctype(self): @property def symbolic_shape(self): - return self._symbolic_shape + # TODO: double check if this should be reversed for dmda + return tuple(self.target._C_get_field(FULL, d).size for d in self.dimensions) # @property # def dmda_info(self): - # return DMDALocalInfo(name='info', liveness='eager') + # return DMDALocalInfo(name='info', liveness='eager') \ No newline at end of file diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 45706c5806..cc66badcbe 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -51,28 +51,22 @@ def test_petsc_functions(): 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) + f0 = Function(name='f', grid=grid, space_order=2, dtype=np.float32) + f1 = Function(name='f', grid=grid, space_order=2, dtype=np.float64) + + ptr0 = PETScArray(name='ptr0', target=f0) + ptr1 = PETScArray(name='ptr1', target=f0, is_const=True) + ptr2 = PETScArray(name='ptr2', target=f1, 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) == 'float *restrict ptr0_vec;' assert str(defn1) == 'const float *restrict ptr1_vec;' assert str(defn2) == 'const double *restrict ptr2_vec;' - assert str(defn3) == 'int *restrict ptr3_vec;' - assert str(defn4) == 'const long *restrict ptr4_vec;' assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' @@ -86,7 +80,7 @@ def test_petsc_subs(): f1 = Function(name='f1', grid=grid, space_order=2) f2 = Function(name='f2', grid=grid, space_order=2) - arr = PETScArray(name='arr', dimensions=f2.dimensions, dtype=f2.dtype) + arr = PETScArray(name='arr', target=f2) eqn = Eq(f1, f2.laplace) eqn_subs = eqn.subs(f2, arr) @@ -174,11 +168,8 @@ def test_multiple_petsc_solves(): callable_roots = [meta_call.root for meta_call in op._func_table.values()] - # One FormRHS, one MatShellMult and one FormFunction per solve - # One PopulateMatContext for all solves - assert len(callable_roots) == 7 - - # Check unique dmda is created per SNES solve + # One FormRHS, MatShellMult, FormFunction, PopulateMatContext per solve + assert len(callable_roots) == 8 @skipif('petsc') @@ -194,12 +185,9 @@ def test_petsc_cast(): f1 = Function(name='f1', grid=g1) f2 = Function(name='f2', grid=g2) - arr0 = PETScArray(name='arr0', dimensions=f0.dimensions, - shape=f0.shape, symbolic_shape=f0.symbolic_shape) - arr1 = PETScArray(name='arr1', dimensions=f1.dimensions, - shape=f1.shape, symbolic_shape=f1.symbolic_shape) - arr2 = PETScArray(name='arr2', dimensions=f2.dimensions, - shape=f2.shape, symbolic_shape=f2.symbolic_shape) + arr0 = PETScArray(name='arr0', target=f0) + arr1 = PETScArray(name='arr1', target=f1) + arr2 = PETScArray(name='arr2', target=f2) cast0 = PointerCast(arr0) cast1 = PointerCast(arr1) @@ -296,8 +284,8 @@ def test_cinterface_petsc_struct(): assert 'include "%s.h"' % name in ccode # The public `struct MatContext` only appears in the header file - assert 'struct MatContext\n{' not in ccode - assert 'struct MatContext\n{' in hcode + assert 'struct J_0_ctx\n{' not in ccode + assert 'struct J_0_ctx\n{' in hcode @skipif('petsc') @@ -656,7 +644,7 @@ def test_petsc_frees(): assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_0)));' assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_0)));' assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_0)));' - assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_so_2)));' + assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_0)));' @skipif('petsc') @@ -733,10 +721,10 @@ def test_time_loop(): body1 = str(op1.body) rhs1 = str(op1._func_table['FormRHS_0'].root.ccode) - assert 'ctx.t0 = t0' in body1 - assert 'ctx.t1 = t1' not in body1 - assert 'formrhs->t0' in rhs1 - assert 'formrhs->t1' not in rhs1 + assert 'ctx0.t0 = t0' in body1 + assert 'ctx0.t1 = t1' not in body1 + assert 'lctx->t0' in rhs1 + assert 'lctx->t1' not in rhs1 # Non-modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) @@ -748,8 +736,8 @@ def test_time_loop(): body2 = str(op2.body) rhs2 = str(op2._func_table['FormRHS_0'].root.ccode) - assert 'ctx.time = time' in body2 - assert 'formrhs->time' in rhs2 + assert 'ctx0.time = time' in body2 + assert 'lctx->time' in rhs2 # Modulo time stepping with more than one time step # used in one of the callback functions @@ -760,10 +748,10 @@ def test_time_loop(): body3 = str(op3.body) rhs3 = str(op3._func_table['FormRHS_0'].root.ccode) - assert 'ctx.t0 = t0' in body3 - assert 'ctx.t1 = t1' in body3 - assert 'formrhs->t0' in rhs3 - assert 'formrhs->t1' in rhs3 + assert 'ctx0.t0 = t0' in body3 + assert 'ctx0.t1 = t1' in body3 + assert 'lctx->t0' in rhs3 + assert 'lctx->t1' in rhs3 # Multiple petsc solves within the same time loop v2 = Function(name='v2', grid=grid, space_order=2) @@ -775,5 +763,5 @@ def test_time_loop(): op4 = Operator(petsc4 + petsc5) body4 = str(op4.body) - assert 'ctx.t0 = t0' in body4 - assert body4.count('ctx.t0 = t0') == 1 + assert 'ctx0.t0 = t0' in body4 + assert body4.count('ctx0.t0 = t0') == 1 From 670e1d2ff4270147b1b423598c2546c6c6dec7cb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 9 Jan 2025 20:51:19 +0000 Subject: [PATCH 129/151] misc: Clean up --- devito/petsc/types/array.py | 51 ++++--------------------------------- tests/test_petsc.py | 36 +++++++++++++------------- 2 files changed, 23 insertions(+), 64 deletions(-) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 8b7cb02ad4..96fc442ecc 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,5 +1,4 @@ from functools import cached_property -import numpy as np from ctypes import POINTER from devito.types.utils import DimensionTuple @@ -8,13 +7,8 @@ from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry from devito.tools import dtype_to_ctype, as_tuple -from devito.symbolics import FieldFromComposite -from devito.types import Symbol from devito.data import FULL -from .object import DM, DMDALocalInfo - - class PETScArray(ArrayBasic, Differentiable): """ @@ -70,6 +64,11 @@ def __indices_setup__(cls, *args, **kwargs): indices = dimensions return as_tuple(dimensions), as_tuple(indices) + def __halo_setup__(self, **kwargs): + target = kwargs['target'] + halo = [target.halo[d] for d in target.space_dimensions] + return DimensionTuple(*halo, getters=target.space_dimensions) + @property def dimensions(self): return self._dimensions @@ -116,47 +115,7 @@ def _C_ctype(self): def halo(self): return [self.target.halo[d] for d in self.target.space_dimensions] - def __halo_setup__(self, **kwargs): - target = kwargs['target'] - halo = [target.halo[d] for d in target.space_dimensions] - # halo = tuple(kwargs.get('halo', ((0, 0),)*self.ndim)) - return DimensionTuple(*halo, getters=target.space_dimensions) - - # @property - # def symbolic_shape(self): - # # from IPython import embed; embed() - # if self.dmda: - # field_from_composites = [ - # FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions] - # # Reverse it since DMDA is setup backwards to Devito dimensions. - # return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - # else: - # return DimensionTuple(*self.dimensions) - - # @property - # def symbolic_shape(self): - # # info = DMDALocalInfo(name='info', liveness='eager') - - # field_from_composites = [ - # FieldFromComposite('g%sm' % d.name, self.dmda_info) for d in self.dimensions] - - # # Reverse it since DMDA is setup backwards to Devito dimensions. - # return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) - - # @cached_property - # def dmda(self): - # name = 'da_so_%s' % self.space_order - # return DM(name=name, liveness='eager', stencil_width=self.space_order) - - # @cached_property - # def dmda(self): - # return self._dmda - @property def symbolic_shape(self): # TODO: double check if this should be reversed for dmda return tuple(self.target._C_get_field(FULL, d).size for d in self.dimensions) - - # @property - # def dmda_info(self): - # return DMDALocalInfo(name='info', liveness='eager') \ No newline at end of file diff --git a/tests/test_petsc.py b/tests/test_petsc.py index cc66badcbe..65a3ee4ab4 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -596,31 +596,31 @@ def test_petsc_struct(): assert all(not isinstance(i, CCompositeObject) for i in op.parameters) -# @skipif('petsc') -# @pytest.mark.parallel(mode=[2, 4, 8]) -# def test_apply(mode): +@skipif('petsc') +@pytest.mark.parallel(mode=[2, 4, 8]) +def test_apply(mode): -# grid = Grid(shape=(13, 13), dtype=np.float64) + grid = Grid(shape=(13, 13), dtype=np.float64) -# pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) -# rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) -# mu = Constant(name='mu', value=2.0) + pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) + rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + mu = Constant(name='mu', value=2.0) -# eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) + eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) -# petsc = PETScSolve(eqn, pn) + petsc = PETScSolve(eqn, pn) -# # Build the op -# with switchconfig(openmp=False, mpi=True): -# op = Operator(petsc) + # Build the op + with switchconfig(openmp=False, mpi=True): + op = Operator(petsc) -# # Check the Operator runs without errors. Not verifying output for -# # now. Need to consolidate BC implementation -# op.apply() + # Check the Operator runs without errors. Not verifying output for + # now. Need to consolidate BC implementation + op.apply() -# # Verify that users can override `mu` -# mu_new = Constant(name='mu_new', value=4.0) -# op.apply(mu=mu_new) + # Verify that users can override `mu` + mu_new = Constant(name='mu_new', value=4.0) + op.apply(mu=mu_new) @skipif('petsc') From ad5fb499e6b721599fb9331ef847c195c6211b19 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 10 Jan 2025 00:08:30 +0000 Subject: [PATCH 130/151] compiler: Timedep abstraction --- devito/petsc/types/array.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 96fc442ecc..7027eb7f77 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -18,8 +18,8 @@ class PETScArray(ArrayBasic, Differentiable): allowing for the use of the `subs` method. PETScArray objects represent vector objects within PETSc. - They correspond to the spatial domain of a Function-like object passed - into PETScSolve from a user. + They correspond to the spatial domain of a Function-like object + provided by the user, which is passed to PETScSolve as the target. TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) and then PETScArray(ArrayBasic). @@ -111,10 +111,6 @@ def _C_ctype(self): # the user's PETSc configuration. return POINTER(dtype_to_ctype(self.dtype)) - @property - def halo(self): - return [self.target.halo[d] for d in self.target.space_dimensions] - @property def symbolic_shape(self): # TODO: double check if this should be reversed for dmda From 89cc06e2b5315b74a733ce76b6aac7a926035564 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 12 Jan 2025 11:15:31 +0000 Subject: [PATCH 131/151] More cleanup, add callbackdm etc --- devito/petsc/solve.py | 2 ++ devito/petsc/types/array.py | 23 +++++++++------- tests/test_petsc.py | 53 +++++++++++++++++++++---------------- 3 files changed, 45 insertions(+), 33 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 4f16ded1f3..8f1cdffee8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -20,6 +20,8 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): localinfo = DMDALocalInfo(name='info', liveness='eager') + localinfo = DMDALocalInfo(name='info', liveness='eager') + arrays = { p: PETScArray(name='%s_%s' % (p, target.name), target=target, diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 7027eb7f77..f2c2b7f05d 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -7,7 +7,7 @@ from devito.types.basic import AbstractFunction from devito.finite_differences.tools import fd_weights_registry from devito.tools import dtype_to_ctype, as_tuple -from devito.data import FULL +from devito.symbolics import FieldFromComposite class PETScArray(ArrayBasic, Differentiable): @@ -31,7 +31,7 @@ class PETScArray(ArrayBasic, Differentiable): _default_fd = 'taylor' __rkwargs__ = (AbstractFunction.__rkwargs__ + - ('target', 'liveness', 'coefficients')) + ('target', 'liveness', 'coefficients', 'localinfo')) def __init_finalize__(self, *args, **kwargs): @@ -47,6 +47,8 @@ def __init_finalize__(self, *args, **kwargs): raise ValueError("coefficients must be one of %s" " not %s" % (str(fd_weights_registry), self._coefficients)) + self._localinfo = kwargs.get('localinfo', None) + @property def ndim(self): return self._ndim @@ -90,6 +92,10 @@ def shape(self): def space_order(self): return self.target.space_order + @property + def localinfo(self): + return self._localinfo + @cached_property def _shape_with_inhalo(self): return self.target.shape_with_inhalo @@ -100,12 +106,7 @@ def shape_allocated(self): @cached_property def _C_ctype(self): - # NOTE: Reverting to using float/double instead of PetscScalar for - # simplicity when opt='advanced'. Otherwise, Temp objects must also - # be converted to PetscScalar. Additional tests are needed to - # ensure this approach is fine. Previously, issues arose from - # mismatches between precision of Function objects in Devito and the - # precision of the PETSc configuration. + # TODO: swith to using PetsccScalar instead of float/double # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of # the user's PETSc configuration. @@ -113,5 +114,7 @@ def _C_ctype(self): @property def symbolic_shape(self): - # TODO: double check if this should be reversed for dmda - return tuple(self.target._C_get_field(FULL, d).size for d in self.dimensions) + field_from_composites = [ + FieldFromComposite('g%sm' % d.name, self.localinfo) for d in self.dimensions] + # Reverse it since DMDA is setup backwards to Devito dimensions. + return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 65a3ee4ab4..d98cb637eb 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -5,7 +5,7 @@ from conftest import skipif from devito import Grid, Function, TimeFunction, Eq, Operator, switchconfig from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, - FindNodes, PointerCast, retrieve_iteration_tree) + FindNodes, retrieve_iteration_tree) from devito.types import Constant, CCompositeObject from devito.passes.iet.languages.C import CDataManager from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, @@ -177,30 +177,37 @@ def test_petsc_cast(): """ Test casting of PETScArray. """ - g0 = Grid((2)) - g1 = Grid((2, 2)) - g2 = Grid((2, 2, 2)) + grid1 = Grid((2)) + grid2 = Grid((2, 2)) + grid3 = Grid((4, 5, 6)) + + f1 = Function(name='f1', grid=grid1, space_order=2) + f2 = Function(name='f2', grid=grid2, space_order=4) + f3 = Function(name='f3', grid=grid3, space_order=6) + + eqn1 = Eq(f1.laplace, 10) + eqn2 = Eq(f2.laplace, 10) + eqn3 = Eq(f3.laplace, 10) - f0 = Function(name='f0', grid=g0) - f1 = Function(name='f1', grid=g1) - f2 = Function(name='f2', grid=g2) + petsc1 = PETScSolve(eqn1, f1) + petsc2 = PETScSolve(eqn2, f2) + petsc3 = PETScSolve(eqn3, f3) - arr0 = PETScArray(name='arr0', target=f0) - arr1 = PETScArray(name='arr1', target=f1) - arr2 = PETScArray(name='arr2', target=f2) + with switchconfig(openmp=False): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + op3 = Operator(petsc3, opt='noop') - cast0 = PointerCast(arr0) - cast1 = PointerCast(arr1) - cast2 = PointerCast(arr2) + cb1 = [meta_call.root for meta_call in op1._func_table.values()] + cb2 = [meta_call.root for meta_call in op2._func_table.values()] + cb3 = [meta_call.root for meta_call in op3._func_table.values()] - assert str(cast0) == \ - 'float (*restrict arr0) = (float (*)) arr0_vec;' - assert str(cast1) == \ - 'float (*restrict arr1)[f1_vec->size[1]] = ' + \ - '(float (*)[f1_vec->size[1]]) arr1_vec;' - assert str(cast2) == \ - 'float (*restrict arr2)[f2_vec->size[1]][f2_vec->size[2]] = ' + \ - '(float (*)[f2_vec->size[1]][f2_vec->size[2]]) arr2_vec;' + assert 'float (*restrict x_matvec_f1) = ' + \ + '(float (*)) x_matvec_f1_vec;' in str(cb1[0]) + assert 'float (*restrict x_matvec_f2)[info.gxm] = ' + \ + '(float (*)[info.gxm]) x_matvec_f2_vec;' in str(cb2[0]) + assert 'float (*restrict x_matvec_f3)[info.gym][info.gxm] = ' + \ + '(float (*)[info.gym][info.gxm]) x_matvec_f3_vec;' in str(cb3[0]) @skipif('petsc') @@ -219,8 +226,8 @@ def test_LinearSolveExpr(): assert linsolveexpr.target == f # Check the solver parameters assert linsolveexpr.solver_parameters == \ - {'ksp_type': 'gmres', 'pc_type': 'jacobi', 'ksp_rtol': 1e-07, - 'ksp_atol': 1e-50, 'ksp_divtol': 10000.0, 'ksp_max_it': 10000} + {'ksp_type': 'gmres', 'pc_type': 'jacobi', 'ksp_rtol': 1e-05, + 'ksp_atol': 1e-50, 'ksp_divtol': 100000.0, 'ksp_max_it': 10000} @skipif('petsc') From 2c4a56814ec74d4c82e8bb16c18c5b0e0ece8fb8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 12 Jan 2025 12:38:19 +0000 Subject: [PATCH 132/151] Address some comments --- devito/petsc/types/array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index f2c2b7f05d..2b4be17da1 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -106,7 +106,7 @@ def shape_allocated(self): @cached_property def _C_ctype(self): - # TODO: swith to using PetsccScalar instead of float/double + # TODO: Switch to using PetscScalar instead of float/double # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of # the user's PETSc configuration. From b83ec923ff6ddd83a951585fa566efe8cd8abb9a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 23 Jan 2025 12:54:22 +0000 Subject: [PATCH 133/151] compiler: Move uxreplace_efuncs into Builder --- tests/test_petsc.py | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d98cb637eb..e57b5c7291 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -603,31 +603,31 @@ def test_petsc_struct(): assert all(not isinstance(i, CCompositeObject) for i in op.parameters) -@skipif('petsc') -@pytest.mark.parallel(mode=[2, 4, 8]) -def test_apply(mode): +# @skipif('petsc') +# @pytest.mark.parallel(mode=[2, 4, 8]) +# def test_apply(mode): - grid = Grid(shape=(13, 13), dtype=np.float64) +# grid = Grid(shape=(13, 13), dtype=np.float64) - pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) - rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) - mu = Constant(name='mu', value=2.0) +# pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) +# rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) +# mu = Constant(name='mu', value=2.0) - eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) +# eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) - petsc = PETScSolve(eqn, pn) +# petsc = PETScSolve(eqn, pn) - # Build the op - with switchconfig(openmp=False, mpi=True): - op = Operator(petsc) +# # Build the op +# with switchconfig(openmp=False, mpi=True): +# op = Operator(petsc) - # Check the Operator runs without errors. Not verifying output for - # now. Need to consolidate BC implementation - op.apply() +# # Check the Operator runs without errors. Not verifying output for +# # now. Need to consolidate BC implementation +# op.apply() - # Verify that users can override `mu` - mu_new = Constant(name='mu_new', value=4.0) - op.apply(mu=mu_new) +# # Verify that users can override `mu` +# mu_new = Constant(name='mu_new', value=4.0) +# op.apply(mu=mu_new) @skipif('petsc') From 4d75af8491a235620af3d289f1689a62f3d18947 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 23 Jan 2025 18:38:38 +0000 Subject: [PATCH 134/151] misc: Clean up --- devito/petsc/solve.py | 2 -- tests/test_petsc.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 8f1cdffee8..4f16ded1f3 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -20,8 +20,6 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): localinfo = DMDALocalInfo(name='info', liveness='eager') - localinfo = DMDALocalInfo(name='info', liveness='eager') - arrays = { p: PETScArray(name='%s_%s' % (p, target.name), target=target, diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e57b5c7291..a1956e42d2 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -567,7 +567,7 @@ def test_callback_arguments(): assert len(ff.parameters) == 4 assert str(mv.parameters) == '(J_0, X_global_0, Y_global_0)' - assert str(ff.parameters) == '(snes_0, X_global_0, Y_global_0, dummy_0)' + assert str(ff.parameters) == '(snes_0, X_global_0, F_global_0, dummy_0)' @skipif('petsc') From 5f25b0f517fa134f7ee70db2db76d19326eef447 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sat, 25 Jan 2025 17:53:26 +0000 Subject: [PATCH 135/151] misc: Add timeloop abstraction docs --- tests/test_petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index a1956e42d2..11a8932e17 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -567,7 +567,7 @@ def test_callback_arguments(): assert len(ff.parameters) == 4 assert str(mv.parameters) == '(J_0, X_global_0, Y_global_0)' - assert str(ff.parameters) == '(snes_0, X_global_0, F_global_0, dummy_0)' + assert str(ff.parameters) == '(snes_0, X_global_0, F_global_0, dummy)' @skipif('petsc') From 37b62b3178d93fbf9ae374c79b14cf3b09f85e36 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 26 Jan 2025 22:45:05 +0000 Subject: [PATCH 136/151] misc: Clean up and docs --- tests/test_petsc.py | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 11a8932e17..13f0d064e5 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -603,31 +603,31 @@ def test_petsc_struct(): assert all(not isinstance(i, CCompositeObject) for i in op.parameters) -# @skipif('petsc') -# @pytest.mark.parallel(mode=[2, 4, 8]) -# def test_apply(mode): +@skipif('petsc') +@pytest.mark.parallel(mode=[2, 4, 8]) +def test_apply(mode): -# grid = Grid(shape=(13, 13), dtype=np.float64) + grid = Grid(shape=(13, 13), dtype=np.float64) -# pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) -# rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) -# mu = Constant(name='mu', value=2.0) + pn = Function(name='pn', grid=grid, space_order=2, dtype=np.float64) + rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + mu = Constant(name='mu', value=2.0) -# eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) + eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) -# petsc = PETScSolve(eqn, pn) + petsc = PETScSolve(eqn, pn) -# # Build the op -# with switchconfig(openmp=False, mpi=True): -# op = Operator(petsc) + # Build the op + with switchconfig(openmp=False, mpi=True): + op = Operator(petsc) -# # Check the Operator runs without errors. Not verifying output for -# # now. Need to consolidate BC implementation -# op.apply() + # Check the Operator runs without errors. Not verifying output for + # now. Need to consolidate BC implementation + op.apply() -# # Verify that users can override `mu` -# mu_new = Constant(name='mu_new', value=4.0) -# op.apply(mu=mu_new) + # Verify that users can override `mu` + mu_new = Constant(name='mu_new', value=4.0) + op.apply(mu=mu_new) @skipif('petsc') From 10af16fc44160591821764c19b9f641526f67b35 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 28 Jan 2025 21:39:35 +0000 Subject: [PATCH 137/151] misc: Spellings --- .github/workflows/examples-mpi.yml | 1 - .github/workflows/examples.yml | 1 - .github/workflows/flake8.yml | 1 - .github/workflows/pytest-core-mpi.yml | 1 - .github/workflows/pytest-core-nompi.yml | 1 - .github/workflows/tutorials.yml | 1 - devito/ir/equations/algorithms.py | 7 ++---- devito/ir/iet/visitors.py | 1 + devito/operator/operator.py | 3 --- devito/petsc/clusters.py | 19 +-------------- devito/petsc/iet/routines.py | 4 ++-- devito/petsc/types/array.py | 2 +- devito/types/basic.py | 31 ------------------------- docker/Dockerfile.cpu | 1 - tests/test_dimension.py | 25 +------------------- tests/test_subdomains.py | 3 +-- 16 files changed, 9 insertions(+), 93 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index bbf3105635..59cd08eaa5 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,7 +17,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 85f84f0cd5..94335a9e2d 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 92556b6be4..336cedb294 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 2eae16021b..67f841cabd 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 5cf374f222..df3a050927 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index 93d64e26ac..c792dd6acd 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,7 +10,6 @@ on: push: branches: - master - - time_loop_branch pull_request: branches: - master diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 2a52b1a41d..946eb4324b 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -169,10 +169,10 @@ def concretize_subdims(exprs, **kwargs): across `exprs`, such as the thickness symbols. """ sregistry = kwargs.get('sregistry') - mapper = kwargs.get('concretize_mapper') + mapper = {} rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions - # from IPython import embed; embed() + _concretize_subdims(exprs, mapper, rebuilt, sregistry) if not mapper: return exprs @@ -206,7 +206,6 @@ def _(v, mapper, rebuilt, sregistry): @_concretize_subdims.register(Eq) def _(expr, mapper, rebuilt, sregistry): - # from IPython import embed; embed() for d in expr.free_symbols: _concretize_subdims(d, mapper, rebuilt, sregistry) @@ -216,7 +215,6 @@ def _(expr, mapper, rebuilt, sregistry): @_concretize_subdims.register(Thickness) def _(tkn, mapper, rebuilt, sregistry): - # from IPython import embed; embed() if tkn in mapper: # Already have a substitution for this thickness return @@ -226,7 +224,6 @@ def _(tkn, mapper, rebuilt, sregistry): @_concretize_subdims.register(SubDimension) def _(d, mapper, rebuilt, sregistry): - # from IPython import embed; embed() if d in mapper: # Already have a substitution for this dimension return diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index d23f24a23d..48570fa89f 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -315,6 +315,7 @@ def _args_call(self, args): def _gen_signature(self, o, is_declaration=False): decls = self._args_decl(o.parameters) + prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),)) if o.attributes: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 73fb6f1450..fe4b5ab48c 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -263,7 +263,6 @@ def _lower(cls, expressions, **kwargs): """ # Create a symbol registry kwargs.setdefault('sregistry', SymbolRegistry()) - kwargs.setdefault('concretize_mapper', {}) expressions = as_tuple(expressions) @@ -390,8 +389,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): clusters = cls._specialize_clusters(clusters, **kwargs) - clusters = petsc_project(clusters) - # Operation count after specialization final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) try: diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 143c73f885..0a6679b2dd 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -1,5 +1,5 @@ from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr, CallbackExpr +from devito.petsc.types import LinearSolveExpr @timed_pass() @@ -25,20 +25,3 @@ def petsc_lift(clusters): else: processed.append(c) return processed - - -@timed_pass() -def petsc_project(clusters): - """ - Drop time loop for clusters which appear in PETSc callback functions. - """ - processed = [] - for c in clusters: - if isinstance(c.exprs[0].rhs, CallbackExpr): - time_dims = [d for d in c.ispace.intervals.dimensions if d.is_Time] - ispace = c.ispace.project(lambda d: d not in time_dims) - processed.append(c.rebuild(ispace=ispace)) - else: - processed.append(c) - - return processed diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index b3dda6acc9..ce082db365 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -351,7 +351,7 @@ def _make_formrhs(self, injectsolve, objs, solver_objs): # Compile formrhs `eqns` into an IET via recursive compilation irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, options={'mpi': False}, sregistry=self.sregistry) - body_formrhs = self._create_formrhs_body(injectsolve, + body_formrhs = self._create_form_rhs_body(injectsolve, List(body=irs_formrhs.uiet.body), solver_objs, objs) @@ -364,7 +364,7 @@ def _make_formrhs(self, injectsolve, objs, solver_objs): self._formrhs_callback = formrhs_callback self._efuncs[formrhs_callback.name] = formrhs_callback - def _create_formrhs_body(self, injectsolve, body, solver_objs, objs): + def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs): linsolve_expr = injectsolve.expr.rhs dmda = solver_objs['callbackdm'] diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 2b4be17da1..38ac3bb9f3 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -14,7 +14,7 @@ class PETScArray(ArrayBasic, Differentiable): """ PETScArrays are generated by the compiler only and represent a customised variant of ArrayBasic. - Differentiable enables compatability with standard Function objects, + Differentiable enables compatibility with standard Function objects, allowing for the use of the `subs` method. PETScArray objects represent vector objects within PETSc. diff --git a/devito/types/basic.py b/devito/types/basic.py index 660ceff437..14e3949fb8 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1513,37 +1513,6 @@ def __getnewargs_ex__(self): return args, kwargs -class LocalType(Basic): - """ - This is the abstract base class for local types, which are - generated by the compiler in C rather than in Python. - Notes - ----- - Subclasses should setup `_liveness`. - """ - - is_LocalType = True - - @property - def liveness(self): - return self._liveness - - @property - def _mem_internal_eager(self): - return self._liveness == 'eager' - - @property - def _mem_internal_lazy(self): - return self._liveness == 'lazy' - - """ - A modifier added to the subclass C declaration when it appears - in a function signature. For example, a subclass might define `_C_modifier = '&'` - to impose pass-by-reference semantics. - """ - _C_modifier = None - - # Extended SymPy hierarchy follows, for essentially two reasons: # - To keep track of `function` # - To override SymPy caching behaviour diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 9e51eb9814..04da0dcc8d 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -44,7 +44,6 @@ RUN cd /tmp && mkdir openmpi && \ cd openmpi && ./autogen.pl && \ mkdir build && cd build && \ ../configure --prefix=/opt/openmpi/ \ - --disable-mpi-fortran \ --enable-mca-no-build=btl-uct --enable-mpi1-compatibility && \ make -j ${nproc} && \ make install && \ diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 0132871a47..5ed0bd08a5 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -2064,8 +2064,7 @@ def test_correct_thicknesses(self): ix1 = SubDimension.right('x', x, 2) ix2 = SubDimension.middle('x', x, 2, 2) - rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry(), - concretize_mapper={}) + rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry()) assert rebuilt[0].is_left assert rebuilt[1].is_right @@ -2106,25 +2105,3 @@ def test_condition_concretization(self): assert len({c.condition for c in conditionals}) == 2 for c in conditionals: assert c.condition in expected_conditionals - - def test_repeat_concretization(self): - """ - Ensure that SubDimensions are consistently concretized to the same object - across multiple calls to the function. This is necessary when using - `rcompile` on equations with SubDimensions. - """ - - grid = Grid((2, 2)) - - x = Dimension('x') - ix = SubDimension.middle('ix', x, 2, 2) - - u = Function(name='u', grid=grid) - eq = Eq(u, ix.ltkn) - - kwargs = {'sregistry': SymbolRegistry(), 'concretize_mapper': {}} - - exprs1 = concretize_subdims([eq], **kwargs) - exprs2 = concretize_subdims([eq], **kwargs) - - assert exprs1 == exprs2 diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index db08f60de6..784dc1c3ad 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -37,8 +37,7 @@ def define(self, dimensions): with timed_region('x'): # _lower_exprs expects a SymbolRegistry, so create one expr = Operator._lower_exprs([eq0], options={}, - sregistry=SymbolRegistry(), - concretize_mapper={})[0] + sregistry=SymbolRegistry())[0] assert str(expr.rhs) == 'ix*f[ix + 1, iy + 1] + iy' def test_multiple_middle(self): From ef7ee34440d0e377bf709119890ab5504aff7b88 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 17:35:19 +0000 Subject: [PATCH 138/151] flake8 --- devito/petsc/iet/routines.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index ce082db365..81547ac64b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -352,8 +352,8 @@ def _make_formrhs(self, injectsolve, objs, solver_objs): irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, options={'mpi': False}, sregistry=self.sregistry) body_formrhs = self._create_form_rhs_body(injectsolve, - List(body=irs_formrhs.uiet.body), - solver_objs, objs) + List(body=irs_formrhs.uiet.body), + solver_objs, objs) formrhs_callback = PETScCallable( self.sregistry.make_name(prefix='FormRHS_'), body_formrhs, retval=objs['err'], From b989d0672d01223181afa33cb9d6465758899ed0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 18:51:51 +0000 Subject: [PATCH 139/151] Clean up --- .github/workflows/examples-mpi.yml | 1 - .github/workflows/examples.yml | 1 - .github/workflows/flake8.yml | 1 - .github/workflows/pytest-core-mpi.yml | 1 - .github/workflows/pytest-core-nompi.yml | 1 - .github/workflows/pytest-petsc.yml | 1 - .github/workflows/tutorials.yml | 1 - devito/symbolics/inspection.py | 1 + tests/test_symbolics.py | 1 + 9 files changed, 2 insertions(+), 7 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index 59cd08eaa5..3bbebabb1c 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -20,7 +20,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: build: diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 94335a9e2d..026331c881 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: tutorials: diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 336cedb294..95a94f6e79 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: flake8: diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 67f841cabd..4f2853e46e 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: test-mpi-basic: diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index df3a050927..54b5b05b1f 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: pytest: diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 20bfc645d1..6f0a729e95 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: pytest: diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index c792dd6acd..d27af78fcf 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: tutorials: diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 64eeeabb9d..123a8c46e4 100644 --- a/devito/symbolics/inspection.py +++ b/devito/symbolics/inspection.py @@ -295,6 +295,7 @@ def sympy_dtype(expr, base=None): """ Infer the dtype of the expression. """ + # TODO: Edit/fix/update according to PR #2513 dtypes = {base} - {None} for i in expr.args: dtype = getattr(i, 'dtype', None) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 4ae6fc2cdf..febec6a25e 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -305,6 +305,7 @@ def test_extended_sympy_arithmetic(): # noncommutative o = Object(name='o', dtype=c_void_p) bar = FieldFromPointer('bar', o) + # TODO: Edit/fix/update according to PR #2513 assert ccode(-1 + bar) == 'o->bar - 1' From 15f78ea3fdc1f41b8f2c65161710bc976e269f53 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 10 Mar 2025 14:05:58 +0000 Subject: [PATCH 140/151] compiler: Initialize and Finalize PETSc once and command line args --- devito/ir/equations/equation.py | 8 ++-- devito/ir/iet/algorithms.py | 4 +- devito/petsc/iet/nodes.py | 16 ++------ devito/petsc/iet/passes.py | 71 +++++++++++++++++++++------------ devito/petsc/iet/routines.py | 4 +- devito/petsc/iet/utils.py | 6 +-- devito/petsc/initialize.py | 42 +++++++++++++++++++ devito/petsc/solve.py | 4 +- devito/petsc/types/macros.py | 5 +++ devito/petsc/types/object.py | 9 ++++- devito/petsc/types/types.py | 22 +++++++++- devito/petsc/utils.py | 2 +- devito/types/equation.py | 2 +- examples/petsc/init_test.py | 8 ++++ examples/petsc/makefile | 5 +++ examples/petsc/petsc_test.py | 29 ++++++++++++++ examples/petsc/test_init.c | 26 ++++++++++++ tests/test_petsc.py | 25 ++++++++---- 18 files changed, 225 insertions(+), 63 deletions(-) create mode 100644 devito/petsc/initialize.py create mode 100644 devito/petsc/types/macros.py create mode 100644 examples/petsc/init_test.py create mode 100644 examples/petsc/makefile create mode 100644 examples/petsc/petsc_test.py create mode 100644 examples/petsc/test_init.c diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 67ed9269a4..6b2261b51c 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -11,10 +11,10 @@ from devito.tools import Pickable, Tag, frozendict from devito.types import (Eq, Inc, ReduceMax, ReduceMin, relational_min) -from devito.types.equation import InjectSolveEq +from devito.types.equation import PetscEq __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', - 'identity_mapper', 'OpInjectSolve'] + 'identity_mapper', 'OpPetsc'] class IREq(sympy.Eq, Pickable): @@ -105,7 +105,7 @@ def detect(cls, expr): Inc: OpInc, ReduceMax: OpMax, ReduceMin: OpMin, - InjectSolveEq: OpInjectSolve + PetscEq: OpPetsc } try: return reduction_mapper[type(expr)] @@ -122,7 +122,7 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') -OpInjectSolve = Operation('solve') +OpPetsc = Operation('solve') identity_mapper = { diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 9d2db185db..52f48e28b1 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -3,7 +3,7 @@ from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, Section, HaloSpot, ExpressionBundle) from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr +from devito.petsc.types import MetaData from devito.petsc.iet.utils import petsc_iet_mapper __all__ = ['iet_build'] @@ -26,7 +26,7 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) - elif isinstance(e.rhs, LinearSolveExpr): + elif isinstance(e.rhs, MetaData): exprs.append(petsc_iet_mapper[e.operation](e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index 2137495487..6157768b16 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -1,20 +1,12 @@ from devito.ir.iet import Expression, Callback, FixedArgsCallable, Call -from devito.ir.equations import OpInjectSolve +from devito.ir.equations import OpPetsc -class LinearSolverExpression(Expression): +class PetscMetaData(Expression): """ - Base class for general expressions required by a - matrix-free linear solve of the form Ax=b. + Base class for general expressions required to run a PETSc solver. """ - pass - - -class InjectSolveDummy(LinearSolverExpression): - """ - Placeholder expression to run the iterative solver. - """ - def __init__(self, expr, pragmas=None, operation=OpInjectSolve): + def __init__(self, expr, pragmas=None, operation=OpPetsc): super().__init__(expr, pragmas=pragmas, operation=operation) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 561f9fd0f4..9a48def8de 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,10 +1,15 @@ import cgen as c +import numpy as np from devito.passes.iet.engine import iet_pass -from devito.ir.iet import Transformer, MapNodes, Iteration, BlankLine +from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, + FindNodes, Call, CallableBody) from devito.symbolics import Byref, Macro -from devito.petsc.types import (PetscMPIInt, PetscErrorCode) -from devito.petsc.iet.nodes import InjectSolveDummy +from devito.types.basic import DataSymbol +from devito.petsc.types import (PetscMPIInt, PetscErrorCode, Initialize, + Finalize, ArgvSymbol) +from devito.petsc.types.macros import petsc_func_begin_user +from devito.petsc.iet.nodes import PetscMetaData from devito.petsc.utils import core_metadata from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup, Solver, TimeDependent, NonTimeDependent) @@ -14,14 +19,23 @@ @iet_pass def lower_petsc(iet, **kwargs): # Check if PETScSolve was used - injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy, + injectsolve_mapper = MapNodes(Iteration, PetscMetaData, 'groupby').visit(iet) if not injectsolve_mapper: return iet, {} + metadata = core_metadata() + + data = FindNodes(PetscMetaData).visit(iet) + + if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)): + return initialize(iet), metadata + + if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): + return finalize(iet), metadata + targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()] - init = init_petsc(**kwargs) # Assumption is that all targets have the same grid so can use any target here objs = build_core_objects(targets[-1], **kwargs) @@ -47,26 +61,37 @@ def lower_petsc(iet, **kwargs): iet = Transformer(subs).visit(iet) body = core + tuple(setup) + (BlankLine,) + iet.body.body - body = iet.body._rebuild( - init=init, body=body, - frees=(petsc_call('PetscFinalize', []),) - ) + body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) - metadata = core_metadata() metadata.update({'efuncs': tuple(efuncs.values())}) return iet, metadata -def init_petsc(**kwargs): - # Initialize PETSc -> for now, assuming all solver options have to be - # specified via the parameters dict in PETScSolve - # TODO: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options - initialize = petsc_call('PetscInitialize', [Null, Null, Null, Null]) +def initialize(iet): + # should be int because the correct type for argc is a C int + # and not a int32 + argc = DataSymbol(name='argc', dtype=np.int32) + argv = ArgvSymbol(name='argv') + Help = Macro('help') - return petsc_func_begin_user, initialize + help_string = c.Line(r'static char help[] = "This is help text.\n";') + + init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help]) + init_body = CallableBody( + body=(petsc_func_begin_user, help_string, init_body), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + return iet._rebuild(body=init_body) + + +def finalize(iet): + finalize_body = petsc_call('PetscFinalize', []) + finalize_body = CallableBody( + body=(petsc_func_begin_user, finalize_body), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + return iet._rebuild(body=finalize_body) def make_core_petsc_calls(objs, **kwargs): @@ -76,10 +101,7 @@ def make_core_petsc_calls(objs, **kwargs): def build_core_objects(target, **kwargs): - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' + communicator = 'PETSC_COMM_WORLD' return { 'size': PetscMPIInt(name='size'), @@ -128,9 +150,6 @@ def __init__(self, injectsolve, objs, iters, **kwargs): ) +# Move these to types folder Null = Macro('NULL') void = 'void' - - -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 81547ac64b..18572a6cbd 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -14,7 +14,7 @@ from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, - MatVecCallback, InjectSolveDummy) + MatVecCallback, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC, @@ -768,7 +768,7 @@ def _spatial_loop_nest(self, iters, injectsolve): spatial_body = [] for tree in retrieve_iteration_tree(iters[0]): root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if injectsolve in FindNodes(InjectSolveDummy).visit(root): + if injectsolve in FindNodes(PetscMetaData).visit(root): spatial_body.append(root) return spatial_body diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index adcf709eab..4f825ad196 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,5 +1,5 @@ -from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall -from devito.ir.equations import OpInjectSolve +from devito.petsc.iet.nodes import PetscMetaData, PETScCall +from devito.ir.equations import OpPetsc def petsc_call(specific_call, call_args): @@ -19,4 +19,4 @@ def petsc_struct(name, fields, pname, liveness='lazy'): # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. -petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} +petsc_iet_mapper = {OpPetsc: PetscMetaData} diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py new file mode 100644 index 0000000000..9126414658 --- /dev/null +++ b/devito/petsc/initialize.py @@ -0,0 +1,42 @@ +import os +import sys +from ctypes import POINTER, cast, c_char +import atexit + +from devito import Operator +from devito.types import Symbol +from devito.types.equation import PetscEq +from devito.petsc.types import Initialize, Finalize + +global _petsc_initialized +_petsc_initialized = False + + +def PetscInitialize(): + global _petsc_initialized + if not _petsc_initialized: + dummy = Symbol(name='d') + # TODO: Potentially just use cgen + the compiler machinery in Devito + # to generate these "dummy_ops" instead of using the Operator class. + # This would prevent circular imports when initializing during import + # from the PETSc module. + op_init = Operator( + [PetscEq(dummy, Initialize(dummy))], + name='kernel_init', opt='noop' + ) + op_finalize = Operator( + [PetscEq(dummy, Finalize(dummy))], + name='kernel_finalize', opt='noop' + ) + + # `argv_bytes` must be a list so the memory address persists + # `os.fsencode` should be preferred over `string().encode('utf-8')` + # in case there is some system specific encoding in use + argv_bytes = list(map(os.fsencode, sys.argv)) + argv_pointer = (POINTER(c_char)*len(sys.argv))( + *map(lambda s: cast(s, POINTER(c_char)), argv_bytes) + ) + op_init.apply(argc=len(sys.argv), argv=argv_pointer) + + atexit.register(op_finalize.apply) + _petsc_initialized = True diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 4f16ded1f3..b00247ae2e 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -5,7 +5,7 @@ from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq, Symbol, SteppingDimension, TimeFunction -from devito.types.equation import InjectSolveEq +from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions from devito.tools import as_tuple @@ -65,7 +65,7 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): ) # Placeholder equation for inserting calls to the solver and generating # correct time loop etc - inject_solve = InjectSolveEq(target, LinearSolveExpr( + inject_solve = PetscEq(target, LinearSolveExpr( expr=tuple(funcs), target=target, solver_parameters=solver_parameters, diff --git a/devito/petsc/types/macros.py b/devito/petsc/types/macros.py new file mode 100644 index 0000000000..4355535e64 --- /dev/null +++ b/devito/petsc/types/macros.py @@ -0,0 +1,5 @@ +import cgen as c + + +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 1bcfb3a6cf..ae9248d2a2 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,7 +1,8 @@ -from ctypes import POINTER +from ctypes import POINTER, c_char from devito.tools import CustomDtype, dtype_to_cstr from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension +from devito.types.basic import DataSymbol from devito.symbolics import Byref from devito.petsc.iet.utils import petsc_call @@ -209,3 +210,9 @@ class StartPtr(LocalObject): def __init__(self, name, dtype): super().__init__(name=name) self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + + +class ArgvSymbol(DataSymbol): + @property + def _C_ctype(self): + return POINTER(POINTER(c_char)) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 1a4a778c9e..c3b43b84af 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -3,7 +3,27 @@ from devito.tools import Reconstructable, sympy_mutex -class LinearSolveExpr(sympy.Function, Reconstructable): +class MetaData(sympy.Function, Reconstructable): + def __new__(cls, expr, **kwargs): + with sympy_mutex: + obj = sympy.Function.__new__(cls, expr) + obj._expr = expr + return obj + + @property + def expr(self): + return self._expr + + +class Initialize(MetaData): + pass + + +class Finalize(MetaData): + pass + + +class LinearSolveExpr(MetaData): """ A symbolic expression passed through the Operator, containing the metadata needed to execute a linear solver. Linear problems are handled with diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index d898db23cb..3dbb2c0d1a 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -45,7 +45,7 @@ def core_metadata(): lib_dir = os.path.join(petsc_dir, f'{petsc_arch}', 'lib') return { - 'includes': ('petscksp.h', 'petscsnes.h', 'petscdmda.h'), + 'includes': ('petscsnes.h', 'petscdmda.h'), 'include_dirs': include_dirs, 'libs': ('petsc'), 'lib_dirs': lib_dir, diff --git a/devito/types/equation.py b/devito/types/equation.py index 947d5072d0..0d33e3debf 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -221,5 +221,5 @@ class ReduceMin(Reduction): pass -class InjectSolveEq(Eq): +class PetscEq(Eq): pass diff --git a/examples/petsc/init_test.py b/examples/petsc/init_test.py new file mode 100644 index 0000000000..462865831f --- /dev/null +++ b/examples/petsc/init_test.py @@ -0,0 +1,8 @@ +import os +from devito.petsc.initialize import PetscInitialize +from devito import configuration +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() +print("helloworld") diff --git a/examples/petsc/makefile b/examples/petsc/makefile new file mode 100644 index 0000000000..ca5d2a9f57 --- /dev/null +++ b/examples/petsc/makefile @@ -0,0 +1,5 @@ +-include ${PETSC_DIR}/petscdir.mk +include ${PETSC_DIR}/lib/petsc/conf/variables +include ${PETSC_DIR}/lib/petsc/conf/rules + +all: test diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py new file mode 100644 index 0000000000..69b63c99ec --- /dev/null +++ b/examples/petsc/petsc_test.py @@ -0,0 +1,29 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, configuration) +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], u) + +op = Operator(petsc) + +op.apply() diff --git a/examples/petsc/test_init.c b/examples/petsc/test_init.c new file mode 100644 index 0000000000..8d92ce176e --- /dev/null +++ b/examples/petsc/test_init.c @@ -0,0 +1,26 @@ +#include + +extern PetscErrorCode PetscInit(); +extern PetscErrorCode PetscFinal(); + +int main(int argc, char **argv) +{ + PetscInit(argc, argv); + PetscPrintf(PETSC_COMM_WORLD, "Hello World!\n"); + return PetscFinalize(); +} + +PetscErrorCode PetscInit(int argc, char **argv) +{ + static char help[] = "Magic help string\n"; + PetscFunctionBeginUser; + PetscCall(PetscInitialize(&argc, &argv, NULL, help)); + PetscFunctionReturn(0); +} + +PetscErrorCode PetscFinal() +{ + PetscFunctionBeginUser; + PetscCall(PetscFinalize()); + PetscFunctionReturn(0); +} diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 13f0d064e5..dafc22211c 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -3,7 +3,8 @@ import pytest from conftest import skipif -from devito import Grid, Function, TimeFunction, Eq, Operator, switchconfig +from devito import (Grid, Function, TimeFunction, Eq, Operator, switchconfig, + configuration) from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, FindNodes, retrieve_iteration_tree) from devito.types import Constant, CCompositeObject @@ -13,6 +14,16 @@ LinearSolveExpr) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil from devito.petsc.iet.nodes import Expression +from devito.petsc.initialize import PetscInitialize + + +@skipif('petsc') +def test_petsc_initialization(): + # TODO: Temporary workaround until PETSc is automatically + # initialized + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() @skipif('petsc') @@ -254,14 +265,14 @@ def test_dmda_create(): op2 = Operator(petsc2, opt='noop') op3 = Operator(petsc3, opt='noop') - assert 'PetscCall(DMDACreate1d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + assert 'PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ '2,1,2,NULL,&(da_0)));' in str(op1) - assert 'PetscCall(DMDACreate2d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_0)));' \ in str(op2) - assert 'PetscCall(DMDACreate3d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + assert 'PetscCall(DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ ',1,1,1,1,6,NULL,NULL,NULL,&(da_0)));' in str(op3) @@ -604,8 +615,7 @@ def test_petsc_struct(): @skipif('petsc') -@pytest.mark.parallel(mode=[2, 4, 8]) -def test_apply(mode): +def test_apply(): grid = Grid(shape=(13, 13), dtype=np.float64) @@ -618,8 +628,7 @@ def test_apply(mode): petsc = PETScSolve(eqn, pn) # Build the op - with switchconfig(openmp=False, mpi=True): - op = Operator(petsc) + op = Operator(petsc) # Check the Operator runs without errors. Not verifying output for # now. Need to consolidate BC implementation From c68dd4f56dc73a397983de274133aff1be21e2e4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 19 Mar 2025 18:39:31 +0000 Subject: [PATCH 141/151] Coupled (#41) * compiler/dsl: Add machinery to support coupled solvers --- .github/workflows/pytest-petsc.yml | 2 +- devito/ir/equations/algorithms.py | 3 +- devito/ir/iet/nodes.py | 7 +- devito/ir/iet/visitors.py | 22 +- devito/operator/operator.py | 2 + devito/passes/iet/definitions.py | 5 +- devito/petsc/clusters.py | 2 +- devito/petsc/iet/nodes.py | 2 +- devito/petsc/iet/passes.py | 189 +++- devito/petsc/iet/routines.py | 1338 ++++++++++++++++++++-------- devito/petsc/iet/utils.py | 5 +- devito/petsc/solve.py | 222 +++-- devito/petsc/types/object.py | 208 +++-- devito/petsc/types/types.py | 216 ++++- devito/symbolics/extended_sympy.py | 45 +- devito/tools/dtypes_lowering.py | 9 +- devito/types/array.py | 15 + devito/types/object.py | 27 +- tests/test_petsc.py | 343 +++++-- 19 files changed, 2005 insertions(+), 657 deletions(-) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 6f0a729e95..91483e2fc4 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -64,7 +64,7 @@ jobs: - name: Test with pytest run: | - ${{ env.RUN_CMD }} pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 946eb4324b..f699b353f0 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -170,7 +170,8 @@ def concretize_subdims(exprs, **kwargs): """ sregistry = kwargs.get('sregistry') - mapper = {} + # Update based on changes in #2509 + mapper = kwargs.get('concretize_mapper', {}) rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions _concretize_subdims(exprs, mapper, rebuilt, sregistry) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index fa7716aaf8..dd6db491c0 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -20,7 +20,7 @@ ctypes_to_cstr) from devito.types.basic import (AbstractFunction, AbstractSymbol, Basic, Indexed, Symbol) -from devito.types.object import AbstractObject, LocalObject +from devito.types.object import AbstractObject, LocalObject, LocalCompositeObject __all__ = ['Node', 'MultiTraversable', 'Block', 'Expression', 'Callable', 'Call', 'ExprStmt', 'Conditional', 'Iteration', 'List', 'Section', @@ -1051,7 +1051,7 @@ class Dereference(ExprStmt, Node): The following cases are supported: * `pointer` is a PointerArray or TempFunction, and `pointee` is an Array. - * `pointer` is an ArrayObject or CCompositeObject representing a pointer + * `pointer` is an ArrayObject or LocalCompositeObject representing a pointer to a C struct, and `pointee` is a field in `pointer`. """ @@ -1077,7 +1077,8 @@ def expr_symbols(self): ret.extend(flatten(i.free_symbols for i in self.pointee.symbolic_shape[1:])) ret.extend(self.pointer.free_symbols) else: - assert issubclass(self.pointer._C_ctype, ctypes._Pointer) + assert (isinstance(self.pointer, LocalCompositeObject) or + issubclass(self.pointer._C_ctype, ctypes._Pointer)) ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) return tuple(filter_ordered(ret)) diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 48570fa89f..25daeb6e56 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -24,7 +24,7 @@ c_restrict_void_p, sorted_priority) from devito.types.basic import AbstractFunction, Basic from devito.types import (ArrayObject, CompositeObject, Dimension, Pointer, - IndexedData, DeviceMap) + IndexedData, DeviceMap, LocalCompositeObject) __all__ = ['FindApplications', 'FindNodes', 'FindSections', 'FindSymbols', @@ -190,7 +190,7 @@ def __init__(self, *args, compiler=None, **kwargs): def _gen_struct_decl(self, obj, masked=()): """ - Convert ctypes.Struct -> cgen.Structure. + Convert ctypes.Struct and LocalCompositeObject -> cgen.Structure. """ ctype = obj._C_ctype try: @@ -201,7 +201,16 @@ def _gen_struct_decl(self, obj, masked=()): return None except TypeError: # E.g., `ctype` is of type `dtypes_lowering.CustomDtype` - return None + if isinstance(obj, LocalCompositeObject): + # TODO: Potentially re-evaluate: Setting ctype to obj allows + # _gen_struct_decl to generate a cgen.Structure from a + # LocalCompositeObject, where obj._C_ctype is a CustomDtype. + # LocalCompositeObject has a __fields__ property, + # which allows the subsequent code in this function to function + # correctly. + ctype = obj + else: + return None try: return obj._C_typedecl @@ -678,8 +687,11 @@ def _operator_typedecls(self, o, mode='all'): for i in o._func_table.values(): if not i.local: continue - typedecls.extend([self._gen_struct_decl(j) for j in i.root.parameters - if xfilter(j)]) + typedecls.extend([ + self._gen_struct_decl(j) + for j in FindSymbols().visit(i.root) + if xfilter(j) + ]) typedecls = filter_sorted(typedecls, key=lambda i: i.tpname) return typedecls diff --git a/devito/operator/operator.py b/devito/operator/operator.py index fe4b5ab48c..3e6f17cf94 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -263,6 +263,8 @@ def _lower(cls, expressions, **kwargs): """ # Create a symbol registry kwargs.setdefault('sregistry', SymbolRegistry()) + # TODO: To be updated based on changes in #2509 + kwargs.setdefault('concretize_mapper', {}) expressions = as_tuple(expressions) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 59de8f95c7..ef5878dacb 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -256,9 +256,10 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage): """ Allocate an Array of Objects in the low latency memory. """ + frees = getattr(obj, '_C_free', None) decl = Definition(obj) - storage.update(obj, site, allocs=decl) + storage.update(obj, site, allocs=decl, frees=frees) def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage): """ @@ -335,7 +336,7 @@ def _inject_definitions(self, iet, storage): frees.extend(as_list(cbody.frees) + flatten(v.frees)) frees = sorted(frees, key=lambda x: min( (obj._C_free_priority for obj in FindSymbols().visit(x) - if obj.is_LocalObject), default=float('inf') + if obj.is_LocalType), default=float('inf') )) # maps/unmaps diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 0a6679b2dd..e035ccbefc 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -20,7 +20,7 @@ def petsc_lift(clusters): processed = [] for c in clusters: if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.space_dimensions) + ispace = c.ispace.lift(c.exprs[0].rhs.fielddata.space_dimensions) processed.append(c.rebuild(ispace=ispace)) else: processed.append(c) diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index 6157768b16..abb5da3acd 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -14,7 +14,7 @@ class PETScCallable(FixedArgsCallable): pass -class MatVecCallback(Callback): +class MatShellSetOp(Callback): @property def callback_form(self): param_types_str = ', '.join([str(t) for t in self.param_types]) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 9a48def8de..710f25a611 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,18 +1,27 @@ import cgen as c import numpy as np +from functools import cached_property from devito.passes.iet.engine import iet_pass from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, - FindNodes, Call, CallableBody) -from devito.symbolics import Byref, Macro + DummyExpr, CallableBody, List, Call, Callable, + FindNodes) +from devito.symbolics import Byref, Macro, FieldFromPointer +from devito.types import Symbol, Scalar from devito.types.basic import DataSymbol -from devito.petsc.types import (PetscMPIInt, PetscErrorCode, Initialize, - Finalize, ArgvSymbol) +from devito.tools import frozendict +from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, + PointerIS, Mat, LocalVec, GlobalVec, CallbackMat, SNES, + DummyArg, PetscInt, PointerDM, PointerMat, MatReuse, + CallbackPointerIS, CallbackPointerDM, JacobianStruct, + SubMatrixStruct, Initialize, Finalize, ArgvSymbol) from devito.petsc.types.macros import petsc_func_begin_user from devito.petsc.iet.nodes import PetscMetaData from devito.petsc.utils import core_metadata -from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup, - Solver, TimeDependent, NonTimeDependent) +from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder, + CoupledObjectBuilder, BaseSetup, CoupledSetup, + Solver, CoupledSolver, TimeDependent, + NonTimeDependent) from devito.petsc.iet.utils import petsc_call, petsc_call_mpi @@ -26,7 +35,6 @@ def lower_petsc(iet, **kwargs): return iet, {} metadata = core_metadata() - data = FindNodes(PetscMetaData).visit(iet) if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)): @@ -35,10 +43,10 @@ def lower_petsc(iet, **kwargs): if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): return finalize(iet), metadata - targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()] - - # Assumption is that all targets have the same grid so can use any target here - objs = build_core_objects(targets[-1], **kwargs) + 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: + raise ValueError("All PETScSolves must use the same Grid, but multiple found.") # Create core PETSc calls (not specific to each PETScSolve) core = make_core_petsc_calls(objs, **kwargs) @@ -54,17 +62,18 @@ def lower_petsc(iet, **kwargs): setup.extend(builder.solversetup.calls) # Transform the spatial iteration loop with the calls to execute the solver - subs.update(builder.solve.mapper) + subs.update({builder.solve.spatial_body: builder.solve.calls}) efuncs.update(builder.cbbuilder.efuncs) + populate_matrix_context(efuncs, objs) + iet = Transformer(subs).visit(iet) body = core + tuple(setup) + (BlankLine,) + iet.body.body body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) metadata.update({'efuncs': tuple(efuncs.values())}) - return iet, metadata @@ -100,56 +109,140 @@ def make_core_petsc_calls(objs, **kwargs): return call_mpi, BlankLine -def build_core_objects(target, **kwargs): - communicator = 'PETSC_COMM_WORLD' - - return { - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'err': PetscErrorCode(name='err'), - 'grid': target.grid - } - - class Builder: """ This class is designed to support future extensions, enabling different combinations of solver types, preconditioning methods, and other functionalities as needed. - The class will be extended to accommodate different solver types by returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ def __init__(self, injectsolve, objs, iters, **kwargs): + self.injectsolve = injectsolve + self.objs = objs + self.iters = iters + self.kwargs = kwargs + self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) + self.args = { + 'injectsolve': self.injectsolve, + 'objs': self.objs, + 'iters': self.iters, + **self.kwargs + } + self.args['solver_objs'] = self.objbuilder.solver_objs + self.args['timedep'] = self.timedep + self.args['cbbuilder'] = self.cbbuilder + + @cached_property + def objbuilder(self): + return ( + CoupledObjectBuilder(**self.args) + if self.coupled else + BaseObjectBuilder(**self.args) + ) - # Determine the time dependency class - time_mapper = injectsolve.expr.rhs.time_mapper - timedep = TimeDependent if time_mapper else NonTimeDependent - self.timedep = timedep(injectsolve, iters, **kwargs) + @cached_property + def timedep(self): + time_mapper = self.injectsolve.expr.rhs.time_mapper + timedep_class = TimeDependent if time_mapper else NonTimeDependent + return timedep_class(**self.args) - # Objects - self.objbuilder = BaseObjectBuilder(injectsolve, **kwargs) - self.solver_objs = self.objbuilder.solver_objs + @cached_property + def cbbuilder(self): + return CCBBuilder(**self.args) if self.coupled else CBBuilder(**self.args) - # Callbacks - self.cbbuilder = CallbackBuilder( - injectsolve, objs, self.solver_objs, timedep=self.timedep, - **kwargs - ) + @cached_property + def solversetup(self): + return CoupledSetup(**self.args) if self.coupled else BaseSetup(**self.args) - # Solver setup - self.solversetup = BaseSetup( - self.solver_objs, objs, injectsolve, self.cbbuilder - ) + @cached_property + def solve(self): + return CoupledSolver(**self.args) if self.coupled else Solver(**self.args) - # Execute the solver - self.solve = Solver( - self.solver_objs, objs, injectsolve, iters, - self.cbbuilder, timedep=self.timedep - ) + +def populate_matrix_context(efuncs, objs): + if not objs['dummyefunc'] in efuncs.values(): + return + + subdms_expr = DummyExpr( + FieldFromPointer(objs['Subdms']._C_symbol, objs['ljacctx']), + objs['Subdms']._C_symbol + ) + fields_expr = DummyExpr( + FieldFromPointer(objs['Fields']._C_symbol, objs['ljacctx']), + objs['Fields']._C_symbol + ) + body = CallableBody( + List(body=[subdms_expr, fields_expr]), + init=(objs['begin_user'],), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + ) + name = 'PopulateMatContext' + efuncs[name] = Callable( + name, body, objs['err'], + parameters=[objs['ljacctx'], objs['Subdms'], objs['Fields']] + ) -# Move these to types folder +# TODO: Devito MPI + PETSc testing +# if kwargs['options']['mpi'] -> communicator = grid.distributor._obj_comm +communicator = 'PETSC_COMM_WORLD' +subdms = PointerDM(name='subdms') +fields = PointerIS(name='fields') +submats = PointerMat(name='submats') +rows = PointerIS(name='rows') +cols = PointerIS(name='cols') + + +# A static dict containing shared symbols and objects that are not +# unique to each PETScSolve. +# Many of these objects are used as arguments in callback functions to make +# the C code cleaner and more modular. This is also a step toward leveraging +# Devito's `reuse_efuncs` functionality, allowing reuse of efuncs when +# they are semantically identical. +objs = frozendict({ + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'err': PetscErrorCode(name='err'), + 'block': CallbackMat('block'), + 'submat_arr': PointerMat(name='submat_arr'), + 'subblockrows': PetscInt('subblockrows'), + 'subblockcols': PetscInt('subblockcols'), + 'rowidx': PetscInt('rowidx'), + 'colidx': PetscInt('colidx'), + 'J': Mat('J'), + 'X': GlobalVec('X'), + 'xloc': LocalVec('xloc'), + 'Y': GlobalVec('Y'), + 'yloc': LocalVec('yloc'), + 'F': GlobalVec('F'), + 'floc': LocalVec('floc'), + 'B': GlobalVec('B'), + 'nfields': PetscInt('nfields'), + 'irow': PointerIS(name='irow'), + 'icol': PointerIS(name='icol'), + 'nsubmats': Scalar('nsubmats', dtype=np.int32), + 'matreuse': MatReuse('scall'), + 'snes': SNES('snes'), + 'rows': rows, + 'cols': cols, + 'Subdms': subdms, + 'LocalSubdms': CallbackPointerDM(name='subdms'), + 'Fields': fields, + 'LocalFields': CallbackPointerIS(name='fields'), + 'Submats': submats, + 'ljacctx': JacobianStruct( + fields=[subdms, fields, submats], modifier=' *' + ), + 'subctx': SubMatrixStruct(fields=[rows, cols]), + 'Null': Macro('NULL'), + 'dummyctx': Symbol('lctx'), + 'dummyptr': DummyArg('dummy'), + 'dummyefunc': Symbol('dummyefunc'), + 'dof': PetscInt('dof'), + 'begin_user': c.Line('PetscFunctionBeginUser;'), +}) + +# Move to macros file? Null = Macro('NULL') -void = 'void' diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 18572a6cbd..a70e987e71 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1,50 +1,54 @@ from collections import OrderedDict - -import cgen as c +from functools import cached_property from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, - retrieve_iteration_tree, filter_iterations) -from devito.symbolics import (Byref, FieldFromPointer, Macro, cast_mapper, - FieldFromComposite) + retrieve_iteration_tree, filter_iterations, Iteration) +from devito.symbolics import (Byref, FieldFromPointer, cast_mapper, VOIDP, + FieldFromComposite, IntDiv, Deref, Mod) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction -from devito.types import Temp, Symbol +from devito.types import Temp, Dimension from devito.tools import filter_ordered from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, - MatVecCallback, PetscMetaData) + MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.petsc.utils import solver_mapper -from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC, - SNES, DummyArg, PetscInt, StartPtr) +from devito.petsc.types import (DM, Mat, LocalVec, GlobalVec, KSP, PC, SNES, + PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, + DMCast, JacobianStructCast, JacobianStruct, + SubMatrixStruct, CallbackDM) -class CallbackBuilder: +class CBBuilder: """ Build IET routines to generate PETSc callback functions. """ - def __init__(self, injectsolve, objs, solver_objs, - rcompile=None, sregistry=None, timedep=None, **kwargs): + def __init__(self, **kwargs): - self.rcompile = rcompile - self.sregistry = sregistry - self.timedep = timedep - self.solver_objs = solver_objs + self.rcompile = kwargs.get('rcompile', None) + self.sregistry = kwargs.get('sregistry', None) + self.concretize_mapper = kwargs.get('concretize_mapper', {}) + self.timedep = kwargs.get('timedep') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.injectsolve = kwargs.get('injectsolve') self._efuncs = OrderedDict() self._struct_params = [] - self._matvec_callback = None - self._formfunc_callback = None - self._formrhs_callback = None - self._struct_callback = None + self._main_matvec_callback = None + self._main_formfunc_callback = None + self._user_struct_callback = None + # TODO: Test pickling. The mutability of these lists + # could cause issues when pickling? + self._matvecs = [] + self._formfuncs = [] + self._formrhs = [] - self._make_core(injectsolve, objs, solver_objs) - self._main_struct(solver_objs) - self._make_struct_callback(solver_objs, objs) - self._local_struct(solver_objs) + self._make_core() self._efuncs = self._uxreplace_efuncs() @property @@ -60,85 +64,104 @@ def filtered_struct_params(self): return filter_ordered(self.struct_params) @property - def matvec_callback(self): - return self._matvec_callback + def main_matvec_callback(self): + """ + This is the matvec callback associated with the whole Jacobian i.e + is set in the main kernel via + `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));` + """ + return self._matvecs[0] + + @property + def main_formfunc_callback(self): + return self._formfuncs[0] + + @property + def matvecs(self): + return self._matvecs @property - def formfunc_callback(self): - return self._formfunc_callback + def formfuncs(self): + return self._formfuncs @property - def formrhs_callback(self): - return self._formrhs_callback + def formrhs(self): + return self._formrhs @property - def struct_callback(self): - return self._struct_callback + def user_struct_callback(self): + return self._user_struct_callback - def _make_core(self, injectsolve, objs, solver_objs): - self._make_matvec(injectsolve, objs, solver_objs) - self._make_formfunc(injectsolve, objs, solver_objs) - self._make_formrhs(injectsolve, objs, solver_objs) + def _make_core(self): + fielddata = self.injectsolve.expr.rhs.fielddata + self._make_matvec(fielddata, fielddata.matvecs) + self._make_formfunc(fielddata) + self._make_formrhs(fielddata) + self._make_user_struct_callback() - def _make_matvec(self, injectsolve, objs, solver_objs): + def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation - irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, - options={'mpi': False}, sregistry=self.sregistry) - body_matvec = self._create_matvec_body(injectsolve, - List(body=irs_matvec.uiet.body), - solver_objs, objs) - - matvec_callback = PETScCallable( - self.sregistry.make_name(prefix='MyMatShellMult_'), body_matvec, + irs_matvec, _ = self.rcompile(matvecs, + options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper) + body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), + fielddata) + + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix=prefix), + body_matvec, retval=objs['err'], - parameters=( - solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'] - ) + parameters=(objs['J'], objs['X'], objs['Y']) ) - self._matvec_callback = matvec_callback - self._efuncs[matvec_callback.name] = matvec_callback + self._matvecs.append(cb) + self._efuncs[cb.name] = cb - def _create_matvec_body(self, injectsolve, body, solver_objs, objs): - linsolve_expr = injectsolve.expr.rhs + def _create_matvec_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs - dmda = solver_objs['callbackdm'] + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + xlocal = objs['xloc'] + ylocal = objs['yloc'] + y_matvec = fielddata.arrays['y'] + x_matvec = fielddata.arrays['x'] body = self.timedep.uxreplace_time(body) - fields = self._dummy_fields(body, solver_objs) + fields = self._dummy_fields(body) - y_matvec = linsolve_expr.arrays['y_matvec'] - x_matvec = linsolve_expr.arrays['x_matvec'] - - mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) + mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(dmda)]) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMGetLocalVector', [dmda, Byref(xlocal)] ) global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], - 'INSERT_VALUES', solver_objs['X_local']] + 'DMGlobalToLocalBegin', [dmda, objs['X'], + insert_vals, xlocal] ) global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + dmda, objs['X'], insert_vals, xlocal ]) dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + 'DMGetLocalVector', [dmda, Byref(ylocal)] ) vec_get_array_y = petsc_call( - 'VecGetArray', [solver_objs['Y_local'], Byref(y_matvec._C_symbol)] + 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] ) vec_get_array_x = petsc_call( - 'VecGetArray', [solver_objs['X_local'], Byref(x_matvec._C_symbol)] + 'VecGetArray', [xlocal, Byref(x_matvec._C_symbol)] ) dm_get_local_info = petsc_call( @@ -146,27 +169,27 @@ def _create_matvec_body(self, injectsolve, body, solver_objs, objs): ) vec_restore_array_y = petsc_call( - 'VecRestoreArray', [solver_objs['Y_local'], Byref(y_matvec._C_symbol)] + 'VecRestoreArray', [ylocal, Byref(y_matvec._C_symbol)] ) vec_restore_array_x = petsc_call( - 'VecRestoreArray', [solver_objs['X_local'], Byref(x_matvec._C_symbol)] + 'VecRestoreArray', [xlocal, Byref(x_matvec._C_symbol)] ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + dmda, ylocal, insert_vals, objs['Y'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + dmda, ylocal, insert_vals, objs['Y'] ]) dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMRestoreLocalVector', [dmda, Byref(xlocal)] ) dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['Y_local'])] + 'DMRestoreLocalVector', [dmda, Byref(ylocal)] ) # TODO: Some of the calls are placed in the `stacks` argument of the @@ -199,84 +222,88 @@ def _create_matvec_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, dummyctx) for i in + dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] matvec_body = CallableBody( List(body=body), - init=(petsc_func_begin_user,), + init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} matvec_body = Uxreplace(subs).visit(matvec_body) self._struct_params.extend(fields) - return matvec_body - def _make_formfunc(self, injectsolve, objs, solver_objs): + def _make_formfunc(self, fielddata): + formfuncs = fielddata.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( - injectsolve.expr.rhs.formfuncs, - options={'mpi': False}, sregistry=self.sregistry + formfuncs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper ) - body_formfunc = self._create_formfunc_body(injectsolve, - List(body=irs_formfunc.uiet.body), - solver_objs, objs) - - formfunc_callback = PETScCallable( - self.sregistry.make_name(prefix='FormFunction_'), body_formfunc, + body_formfunc = self._create_formfunc_body( + List(body=irs_formfunc.uiet.body), fielddata + ) + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix='FormFunction'), + body_formfunc, retval=objs['err'], - parameters=(solver_objs['snes'], solver_objs['X_global'], - solver_objs['F_global'], dummyptr) + parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._formfunc_callback = formfunc_callback - self._efuncs[formfunc_callback.name] = formfunc_callback + self._formfuncs.append(cb) + self._efuncs[cb.name] = cb - def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): - linsolve_expr = injectsolve.expr.rhs + def _create_formfunc_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs - dmda = solver_objs['callbackdm'] + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] body = self.timedep.uxreplace_time(body) - fields = self._dummy_fields(body, solver_objs) + fields = self._dummy_fields(body) + self._struct_params.extend(fields) - f_formfunc = linsolve_expr.arrays['f_formfunc'] - x_formfunc = linsolve_expr.arrays['x_formfunc'] + f_formfunc = fielddata.arrays['f'] + x_formfunc = fielddata.arrays['x'] - snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] ) global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], - 'INSERT_VALUES', solver_objs['X_local']] + 'DMGlobalToLocalBegin', [dmda, objs['X'], + insert_vals, objs['xloc']] ) global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + dmda, objs['X'], insert_vals, objs['xloc'] ]) dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['F_local'])] + 'DMGetLocalVector', [dmda, Byref(objs['floc'])] ) vec_get_array_y = petsc_call( - 'VecGetArray', [solver_objs['F_local'], Byref(f_formfunc._C_symbol)] + 'VecGetArray', [objs['floc'], Byref(f_formfunc._C_symbol)] ) vec_get_array_x = petsc_call( - 'VecGetArray', [solver_objs['X_local'], Byref(x_formfunc._C_symbol)] + 'VecGetArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] ) dm_get_local_info = petsc_call( @@ -284,27 +311,27 @@ def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) vec_restore_array_y = petsc_call( - 'VecRestoreArray', [solver_objs['F_local'], Byref(f_formfunc._C_symbol)] + 'VecRestoreArray', [objs['floc'], Byref(f_formfunc._C_symbol)] ) vec_restore_array_x = petsc_call( - 'VecRestoreArray', [solver_objs['X_local'], Byref(x_formfunc._C_symbol)] + 'VecRestoreArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] + dmda, objs['floc'], insert_vals, objs['F'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] + dmda, objs['floc'], insert_vals, objs['F'] ]) dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] ) dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['F_local'])] + 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] ) body = body._rebuild( @@ -318,7 +345,7 @@ def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) stacks = ( - snes_get_dm, + dm_cast, dm_get_app_context, dm_get_local_xvec, global_to_local_begin, @@ -330,51 +357,68 @@ def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, dummyctx) for i in + dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] formfunc_body = CallableBody( List(body=body), - init=(petsc_func_begin_user,), + init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} - formfunc_body = Uxreplace(subs).visit(formfunc_body) + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - self._struct_params.extend(fields) + return Uxreplace(subs).visit(formfunc_body) - return formfunc_body + def _make_formrhs(self, fielddata): + formrhs = fielddata.formrhs + sobjs = self.solver_objs - def _make_formrhs(self, injectsolve, objs, solver_objs): # Compile formrhs `eqns` into an IET via recursive compilation - irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, - options={'mpi': False}, sregistry=self.sregistry) - body_formrhs = self._create_form_rhs_body(injectsolve, - List(body=irs_formrhs.uiet.body), - solver_objs, objs) - - formrhs_callback = PETScCallable( - self.sregistry.make_name(prefix='FormRHS_'), body_formrhs, retval=objs['err'], - parameters=( - solver_objs['snes'], solver_objs['b_local'] - ) + irs_formrhs, _ = self.rcompile( + formrhs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper ) - self._formrhs_callback = formrhs_callback - self._efuncs[formrhs_callback.name] = formrhs_callback + body_formrhs = self._create_form_rhs_body( + List(body=irs_formrhs.uiet.body), fielddata + ) + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix='FormRHS'), + body_formrhs, + retval=objs['err'], + parameters=(sobjs['callbackdm'], objs['B']) + ) + self._formrhs.append(cb) + self._efuncs[cb.name] = cb + + def _create_form_rhs_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs - def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs): - linsolve_expr = injectsolve.expr.rhs + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] - dmda = solver_objs['callbackdm'] + dm_get_local = petsc_call( + 'DMGetLocalVector', [dmda, Byref(sobjs['blocal'])] + ) + + dm_global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['B'], + insert_vals, sobjs['blocal']] + ) - snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + dm_global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, objs['B'], insert_vals, + sobjs['blocal'] + ]) - b_arr = linsolve_expr.arrays['b_tmp'] + b_arr = fielddata.arrays['b'] vec_get_array = petsc_call( - 'VecGetArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] ) dm_get_local_info = petsc_call( @@ -383,87 +427,85 @@ def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs): body = self.timedep.uxreplace_time(body) - fields = self._dummy_fields(body, solver_objs) + fields = self._dummy_fields(body) + self._struct_params.extend(fields) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, sobjs['blocal'], insert_vals, + objs['B'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, sobjs['blocal'], insert_vals, + objs['B'] + ]) + vec_restore_array = petsc_call( - 'VecRestoreArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + 'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] ) - body = body._rebuild(body=body.body + (vec_restore_array,)) + body = body._rebuild(body=body.body + ( + dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array + )) stacks = ( - snes_get_dm, - dm_get_app_context, + dm_get_local, + dm_global_to_local_begin, + dm_global_to_local_end, vec_get_array, + dm_get_app_context, dm_get_local_info ) # Dereference function data in struct - dereference_funcs = [Dereference(i, dummyctx) for i in + dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] formrhs_body = CallableBody( List(body=[body]), - init=(petsc_func_begin_user,), + init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields if not isinstance(i.function, AbstractFunction)} - formrhs_body = Uxreplace(subs).visit(formrhs_body) - self._struct_params.extend(fields) - - return formrhs_body - - def _local_struct(self, solver_objs): - """ - This is the struct used within callback functions, - usually accessed via DMGetApplicationContext. - """ - solver_objs['localctx'] = petsc_struct( - dummyctx.name, - self.filtered_struct_params, - solver_objs['Jac'].name+'_ctx', - liveness='eager' - ) + return Uxreplace(subs).visit(formrhs_body) - def _main_struct(self, solver_objs): + def _make_user_struct_callback(self): """ This is the struct initialised inside the main kernel and attached to the DM via DMSetApplicationContext. + # TODO: this could be common between all PETScSolves instead? """ - solver_objs['mainctx'] = petsc_struct( + mainctx = self.solver_objs['userctx'] = petsc_struct( self.sregistry.make_name(prefix='ctx'), self.filtered_struct_params, - solver_objs['Jac'].name+'_ctx' + self.sregistry.make_name(prefix='UserCtx'), ) - - def _make_struct_callback(self, solver_objs, objs): - mainctx = solver_objs['mainctx'] body = [ DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) for i in mainctx.callback_fields ] struct_callback_body = CallableBody( - List(body=body), init=(petsc_func_begin_user,), - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + List(body=body), init=(self.objs['begin_user'],), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - struct_callback = Callable( - self.sregistry.make_name(prefix='PopulateMatContext_'), - struct_callback_body, objs['err'], + cb = Callable( + self.sregistry.make_name(prefix='PopulateUserContext'), + struct_callback_body, self.objs['err'], parameters=[mainctx] ) - self._efuncs[struct_callback.name] = struct_callback - self._struct_callback = struct_callback + self._efuncs[cb.name] = cb + self._user_struct_callback = cb - def _dummy_fields(self, iet, solver_objs): + def _dummy_fields(self, iet): # Place all context data required by the shell routines into a struct fields = [f.function for f in FindSymbols('basics').visit(iet)] fields = [f for f in fields if not isinstance(f.function, (PETScArray, Temp))] @@ -473,84 +515,385 @@ def _dummy_fields(self, iet, solver_objs): return fields def _uxreplace_efuncs(self): + sobjs = self.solver_objs + luserctx = petsc_struct( + sobjs['userctx'].name, + self.filtered_struct_params, + sobjs['userctx'].pname, + modifier=' *' + ) mapper = {} - visitor = Uxreplace({dummyctx: self.solver_objs['localctx']}) + visitor = Uxreplace({self.objs['dummyctx']: luserctx}) for k, v in self._efuncs.items(): mapper.update({k: visitor.visit(v)}) return mapper +class CCBBuilder(CBBuilder): + def __init__(self, **kwargs): + self._submatrices_callback = None + super().__init__(**kwargs) + + @property + def submatrices_callback(self): + return self._submatrices_callback + + @property + def submatrices(self): + return self.injectsolve.expr.rhs.fielddata.submatrices + + @property + def main_matvec_callback(self): + """ + This is the matvec callback associated with the whole Jacobian i.e + is set in the main kernel via + `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));` + """ + return self._main_matvec_callback + + @property + def main_formfunc_callback(self): + return self._main_formfunc_callback + + def _make_core(self): + injectsolve = self.injectsolve + targets = injectsolve.expr.rhs.fielddata.targets + all_fielddata = injectsolve.expr.rhs.fielddata + + for t in targets: + data = all_fielddata.get_field_data(t) + self._make_formfunc(data) + self._make_formrhs(data) + + row_matvecs = all_fielddata.submatrices.submatrices[t] + for submat, mtvs in row_matvecs.items(): + if mtvs['matvecs']: + self._make_matvec(data, mtvs['matvecs'], prefix=f'{submat}_MatMult') + + self._make_user_struct_callback() + self._make_whole_matvec() + self._make_whole_formfunc() + self._create_submatrices() + self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] + + def _make_whole_matvec(self): + objs = self.objs + body = self._whole_matvec_body() + + cb = PETScCallable( + self.sregistry.make_name(prefix='WholeMatMult'), + List(body=body), + retval=objs['err'], + parameters=(objs['J'], objs['X'], objs['Y']) + ) + self._main_matvec_callback = cb + self._efuncs[cb.name] = cb + + def _whole_matvec_body(self): + objs = self.objs + sobjs = self.solver_objs + + jctx = objs['ljacctx'] + ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) + + nonzero_submats = self.submatrices.nonzero_submatrix_keys + + calls = () + for sm in nonzero_submats: + idx = self.submatrices.submat_to_index[sm] + ctx = sobjs[f'{sm}ctx'] + X = sobjs[f'{sm}X'] + Y = sobjs[f'{sm}Y'] + rows = objs['rows'].base + cols = objs['cols'].base + sm_indexed = objs['Submats'].indexed[idx] + + calls += ( + DummyExpr(sobjs[sm], FieldFromPointer(sm_indexed, jctx)), + petsc_call('MatShellGetContext', [sobjs[sm], Byref(ctx)]), + petsc_call( + 'VecGetSubVector', + [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] + ), + petsc_call( + 'VecGetSubVector', + [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] + ), + petsc_call('MatMult', [sobjs[sm], X, Y]), + petsc_call( + 'VecRestoreSubVector', + [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] + ), + petsc_call( + 'VecRestoreSubVector', + [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] + ), + ) + return CallableBody( + List(body=(ctx_main, BlankLine) + calls), + init=(objs['begin_user'],), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + def _make_whole_formfunc(self): + objs = self.objs + body = self._whole_formfunc_body() + + cb = PETScCallable( + self.sregistry.make_name(prefix='WholeFormFunc'), + List(body=body), + retval=objs['err'], + parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + ) + self._main_formfunc_callback = cb + self._efuncs[cb.name] = cb + + def _whole_formfunc_body(self): + objs = self.objs + sobjs = self.solver_objs + + ljacctx = objs['ljacctx'] + struct_cast = DummyExpr(ljacctx, JacobianStructCast(objs['dummyptr'])) + X = objs['X'] + F = objs['F'] + + targets = self.injectsolve.expr.rhs.fielddata.targets + + deref_subdms = Dereference(objs['LocalSubdms'], ljacctx) + deref_fields = Dereference(objs['LocalFields'], ljacctx) + + calls = () + for i, t in enumerate(targets): + field_ptr = FieldFromPointer(objs['LocalFields'].indexed[i], ljacctx) + x_name = f'Xglobal{t.name}' + f_name = f'Fglobal{t.name}' + + calls += ( + petsc_call('VecGetSubVector', [X, field_ptr, Byref(sobjs[x_name])]), + petsc_call('VecGetSubVector', [F, field_ptr, Byref(sobjs[f_name])]), + petsc_call(self.formfuncs[i].name, [objs['snes'], sobjs[x_name], + sobjs[f_name], VOIDP(objs['LocalSubdms'].indexed[i])]), + petsc_call('VecRestoreSubVector', [X, field_ptr, Byref(sobjs[x_name])]), + petsc_call('VecRestoreSubVector', [F, field_ptr, Byref(sobjs[f_name])]), + ) + return CallableBody( + List(body=calls + (BlankLine,)), + init=(objs['begin_user'],), + stacks=(struct_cast, deref_subdms, deref_fields), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + def _create_submatrices(self): + body = self._submat_callback_body() + objs = self.objs + params = ( + objs['J'], + objs['nfields'], + objs['irow'], + objs['icol'], + objs['matreuse'], + objs['Submats'], + ) + cb = PETScCallable( + self.sregistry.make_name(prefix='MatCreateSubMatrices'), + body, + retval=objs['err'], + parameters=params + ) + self._submatrices_callback = cb + self._efuncs[cb.name] = cb + + def _submat_callback_body(self): + objs = self.objs + sobjs = self.solver_objs + + n_submats = DummyExpr( + objs['nsubmats'], Mul(objs['nfields'], objs['nfields']) + ) + + malloc_submats = petsc_call('PetscCalloc1', [objs['nsubmats'], objs['Submats']]) + + mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(sobjs['callbackdm'])]) + + dm_get_app = petsc_call( + 'DMGetApplicationContext', [sobjs['callbackdm'], Byref(objs['dummyctx'])] + ) + + get_ctx = petsc_call('MatShellGetContext', [objs['J'], Byref(objs['ljacctx'])]) + + Null = objs['Null'] + dm_get_info = petsc_call( + 'DMDAGetInfo', [ + sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']), + Null, Null, Null, Null, Byref(objs['dof']), Null, Null, Null, Null, Null + ] + ) + subblock_rows = DummyExpr(objs['subblockrows'], Mul(sobjs['M'], sobjs['N'])) + subblock_cols = DummyExpr(objs['subblockcols'], Mul(sobjs['M'], sobjs['N'])) + + ptr = DummyExpr(objs['submat_arr']._C_symbol, Deref(objs['Submats']), init=True) + + mat_create = petsc_call('MatCreate', [self.objs['comm'], Byref(objs['block'])]) + + mat_set_sizes = petsc_call( + 'MatSetSizes', [ + objs['block'], 'PETSC_DECIDE', 'PETSC_DECIDE', + objs['subblockrows'], objs['subblockcols'] + ] + ) + + mat_set_type = petsc_call('MatSetType', [objs['block'], 'MATSHELL']) + + malloc = petsc_call('PetscMalloc1', [1, Byref(objs['subctx'])]) + i = Dimension(name='i') + + row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof'])) + col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof'])) + + deref_subdm = Dereference(objs['Subdms'], objs['ljacctx']) + + set_rows = DummyExpr( + FieldFromPointer(objs['rows'].base, objs['subctx']), + Byref(objs['irow'].indexed[objs['rowidx']]) + ) + set_cols = DummyExpr( + FieldFromPointer(objs['cols'].base, objs['subctx']), + Byref(objs['icol'].indexed[objs['colidx']]) + ) + dm_set_ctx = petsc_call( + 'DMSetApplicationContext', [ + objs['Subdms'].indexed[objs['rowidx']], objs['dummyctx'] + ] + ) + matset_dm = petsc_call('MatSetDM', [ + objs['block'], objs['Subdms'].indexed[objs['rowidx']] + ]) + + set_ctx = petsc_call('MatShellSetContext', [objs['block'], objs['subctx']]) + + mat_setup = petsc_call('MatSetUp', [objs['block']]) + + assign_block = DummyExpr(objs['submat_arr'].indexed[i], objs['block']) + + iter_body = ( + mat_create, + mat_set_sizes, + mat_set_type, + malloc, + row_idx, + col_idx, + set_rows, + set_cols, + dm_set_ctx, + matset_dm, + set_ctx, + mat_setup, + assign_block + ) + + upper_bound = objs['nsubmats'] - 1 + iteration = Iteration(List(body=iter_body), i, upper_bound) + + nonzero_submats = self.submatrices.nonzero_submatrix_keys + matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} + + matmult_op = [ + petsc_call( + 'MatShellSetOperation', + [ + objs['submat_arr'].indexed[self.submatrices.submat_to_index[sb]], + 'MATOP_MULT', + MatShellSetOp(matvec_lookup[sb].name, void, void), + ], + ) + for sb in nonzero_submats if sb in matvec_lookup + ] + + body = [ + n_submats, + malloc_submats, + mat_get_dm, + dm_get_app, + dm_get_info, + subblock_rows, + subblock_cols, + ptr, + BlankLine, + iteration, + ] + matmult_op + + return CallableBody( + List(body=tuple(body)), + init=(objs['begin_user'],), + stacks=(get_ctx, deref_subdm), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + class BaseObjectBuilder: """ A base class for constructing objects needed for a PETSc solver. Designed to be extended by subclasses, which can override the `_extend_build` method to support specific use cases. """ - - def __init__(self, injectsolve, sregistry=None, **kwargs): - self.sregistry = sregistry - self.solver_objs = self._build(injectsolve) - - def _build(self, injectsolve): + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.objs = kwargs.get('objs') + self.sregistry = kwargs.get('sregistry') + self.fielddata = self.injectsolve.expr.rhs.fielddata + self.solver_objs = self._build() + + def _build(self): """ + # TODO: update docs Constructs the core dictionary of solver objects and allows subclasses to extend or modify it via `_extend_build`. - Returns: dict: A dictionary containing the following objects: - 'Jac' (Mat): A matrix representing the jacobian. - - 'x_global' (GlobalVec): The global solution vector. - - 'x_local' (LocalVec): The local solution vector. - - 'b_global': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. - - 'b_local': (LocalVec) Local RHS vector `b`, where `F(x) = b`. + - 'xglobal' (GlobalVec): The global solution vector. + - 'xlocal' (LocalVec): The local solution vector. + - 'bglobal': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. + - 'blocal': (LocalVec) Local RHS vector `b`, where `F(x) = b`. - 'ksp': (KSP) Krylov solver object that manages the linear solver. - 'pc': (PC) Preconditioner object. - 'snes': (SNES) Nonlinear solver object. - - 'F_global': (GlobalVec) Global residual vector `F`, where `F(x) = b`. - - 'F_local': (LocalVec) Local residual vector `F`, where `F(x) = b`. - - 'Y_global': (GlobalVector) The output vector populated by the - matrix-free `MyMatShellMult` callback function. - - 'Y_local': (LocalVector) The output vector populated by the matrix-free - `MyMatShellMult` callback function. - - 'X_global': (GlobalVec) Current guess for the solution, - required by the FormFunction callback. - - 'X_local': (LocalVec) Current guess for the solution, - required by the FormFunction callback. - 'localsize' (PetscInt): The local length of the solution vector. - - 'start_ptr' (StartPtr): A pointer to the beginning of the solution array - that will be updated at each time step. - 'dmda' (DM): The DMDA object associated with this solve, linked to the SNES object via `SNESSetDM`. - 'callbackdm' (CallbackDM): The DM object accessed within callback functions via `SNESGetDM`. """ - target = injectsolve.expr.rhs.target sreg = self.sregistry + targets = self.fielddata.targets base_dict = { - 'Jac': Mat(sreg.make_name(prefix='J_')), - 'x_global': GlobalVec(sreg.make_name(prefix='x_global_')), - 'x_local': LocalVec(sreg.make_name(prefix='x_local_'), liveness='eager'), - 'b_global': GlobalVec(sreg.make_name(prefix='b_global_')), - 'b_local': LocalVec(sreg.make_name(prefix='b_local_')), - 'ksp': KSP(sreg.make_name(prefix='ksp_')), - 'pc': PC(sreg.make_name(prefix='pc_')), - 'snes': SNES(sreg.make_name(prefix='snes_')), - 'F_global': GlobalVec(sreg.make_name(prefix='F_global_')), - 'F_local': LocalVec(sreg.make_name(prefix='F_local_'), liveness='eager'), - 'Y_global': GlobalVec(sreg.make_name(prefix='Y_global_')), - 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), - 'X_global': GlobalVec(sreg.make_name(prefix='X_global_')), - 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), - 'localsize': PetscInt(sreg.make_name(prefix='localsize_')), - 'start_ptr': StartPtr(sreg.make_name(prefix='start_ptr_'), target.dtype), - 'dmda': DM(sreg.make_name(prefix='da_'), liveness='eager', - stencil_width=target.space_order), - 'callbackdm': CallbackDM(sreg.make_name(prefix='dm_'), - liveness='eager', stencil_width=target.space_order), + 'Jac': Mat(sreg.make_name(prefix='J')), + 'xglobal': GlobalVec(sreg.make_name(prefix='xglobal')), + 'xlocal': LocalVec(sreg.make_name(prefix='xlocal')), + 'bglobal': GlobalVec(sreg.make_name(prefix='bglobal')), + 'blocal': LocalVec(sreg.make_name(prefix='blocal')), + 'ksp': KSP(sreg.make_name(prefix='ksp')), + 'pc': PC(sreg.make_name(prefix='pc')), + 'snes': SNES(sreg.make_name(prefix='snes')), + 'localsize': PetscInt(sreg.make_name(prefix='localsize')), + 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), + 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } - return self._extend_build(base_dict, injectsolve) + self._target_dependent(base_dict) + return self._extend_build(base_dict) - def _extend_build(self, base_dict, injectsolve): + def _target_dependent(self, base_dict): + """ + '_ptr' (StartPtr): A pointer to the beginning of the solution array + that will be updated at each time step. + """ + sreg = self.sregistry + target = self.fielddata.target + base_dict[f'{target.name}_ptr'] = StartPtr( + sreg.make_name(prefix=f'{target.name}_ptr'), target.dtype + ) + + def _extend_build(self, base_dict): """ Subclasses can override this method to extend or modify the base dictionary of solver objects. @@ -558,83 +901,173 @@ def _extend_build(self, base_dict, injectsolve): return base_dict +class CoupledObjectBuilder(BaseObjectBuilder): + def _extend_build(self, base_dict): + injectsolve = self.injectsolve + sreg = self.sregistry + objs = self.objs + targets = self.fielddata.targets + + base_dict['fields'] = PointerIS( + name=sreg.make_name(prefix='fields'), nindices=len(targets) + ) + base_dict['subdms'] = PointerDM( + name=sreg.make_name(prefix='subdms'), nindices=len(targets) + ) + base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields')) + + space_dims = len(self.fielddata.grid.dimensions) + + dim_labels = ["M", "N", "P"] + base_dict.update({ + dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) + }) + + submatrices = injectsolve.expr.rhs.fielddata.submatrices + submatrix_keys = submatrices.submatrix_keys + + base_dict['jacctx'] = JacobianStruct( + name=sreg.make_name(prefix=objs['ljacctx'].name), + fields=objs['ljacctx'].fields, + ) + + for key in submatrix_keys: + base_dict[key] = Mat(name=key) + base_dict[f'{key}ctx'] = SubMatrixStruct( + name=f'{key}ctx', + fields=objs['subctx'].fields, + ) + base_dict[f'{key}X'] = LocalVec(f'{key}X') + base_dict[f'{key}Y'] = LocalVec(f'{key}Y') + base_dict[f'{key}F'] = LocalVec(f'{key}F') + + return base_dict + + def _target_dependent(self, base_dict): + sreg = self.sregistry + targets = self.fielddata.targets + for t in targets: + name = t.name + base_dict[f'{name}_ptr'] = StartPtr( + sreg.make_name(prefix=f'{name}_ptr'), t.dtype + ) + base_dict[f'xlocal{name}'] = LocalVec( + sreg.make_name(prefix=f'xlocal{name}'), liveness='eager' + ) + base_dict[f'Fglobal{name}'] = LocalVec( + sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager' + ) + base_dict[f'Xglobal{name}'] = LocalVec( + sreg.make_name(prefix=f'Xglobal{name}') + ) + base_dict[f'xglobal{name}'] = GlobalVec( + sreg.make_name(prefix=f'xglobal{name}') + ) + base_dict[f'blocal{name}'] = LocalVec( + sreg.make_name(prefix=f'blocal{name}'), liveness='eager' + ) + base_dict[f'bglobal{name}'] = GlobalVec( + sreg.make_name(prefix=f'bglobal{name}') + ) + base_dict[f'da{name}'] = DM( + sreg.make_name(prefix=f'da{name}'), liveness='eager' + ) + base_dict[f'scatter{name}'] = VecScatter( + sreg.make_name(prefix=f'scatter{name}') + ) + + class BaseSetup: - def __init__(self, solver_objs, objs, injectsolve, cbbuilder): - self.calls = self._setup(solver_objs, objs, injectsolve, cbbuilder) + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.cbbuilder = kwargs.get('cbbuilder') + self.fielddata = self.injectsolve.expr.rhs.fielddata + self.calls = self._setup() - def _setup(self, solver_objs, objs, injectsolve, cbbuilder): - dmda = solver_objs['dmda'] + @property + def snes_ctx(self): + """ + The [optional] context for private data for the function evaluation routine. + https://petsc.org/main/manualpages/SNES/SNESSetFunction/ + """ + return VOIDP(self.solver_objs['dmda']) - solver_params = injectsolve.expr.rhs.solver_parameters + def _setup(self): + objs = self.objs + sobjs = self.solver_objs - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) + dmda = sobjs['dmda'] - snes_set_dm = petsc_call('SNESSetDM', [solver_objs['snes'], dmda]) + solver_params = self.injectsolve.expr.rhs.solver_parameters - create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(solver_objs['Jac'])]) + snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) - # NOTE: Assuming all solves are linear for now. - snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) + + # NOTE: Assuming all solves are linear for now + snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY']) snes_set_jac = petsc_call( - 'SNESSetJacobian', [solver_objs['snes'], solver_objs['Jac'], - solver_objs['Jac'], 'MatMFFDComputeJacobian', Null] + 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], + sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']] ) global_x = petsc_call('DMCreateGlobalVector', - [dmda, Byref(solver_objs['x_global'])]) + [dmda, Byref(sobjs['xglobal'])]) global_b = petsc_call('DMCreateGlobalVector', - [dmda, Byref(solver_objs['b_global'])]) - - local_b = petsc_call('DMCreateLocalVector', - [dmda, Byref(solver_objs['b_local'])]) + [dmda, Byref(sobjs['bglobal'])]) snes_get_ksp = petsc_call('SNESGetKSP', - [solver_objs['snes'], Byref(solver_objs['ksp'])]) + [sobjs['snes'], Byref(sobjs['ksp'])]) ksp_set_tols = petsc_call( - 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], + 'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], solver_params['ksp_divtol'], solver_params['ksp_max_it']] ) ksp_set_type = petsc_call( - 'KSPSetType', [solver_objs['ksp'], solver_mapper[solver_params['ksp_type']]] + 'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]] ) ksp_get_pc = petsc_call( - 'KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])] + 'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])] ) # Even though the default will be jacobi, set to PCNONE for now - pc_set_type = petsc_call('PCSetType', [solver_objs['pc'], 'PCNONE']) + pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE']) - ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) + ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']]) + matvec = self.cbbuilder.main_matvec_callback matvec_operation = petsc_call( 'MatShellSetOperation', - [solver_objs['Jac'], 'MATOP_MULT', - MatVecCallback(cbbuilder.matvec_callback.name, void, void)] + [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] ) - + formfunc = self.cbbuilder.main_formfunc_callback formfunc_operation = petsc_call( 'SNESSetFunction', - [solver_objs['snes'], Null, - FormFunctionCallback(cbbuilder.formfunc_callback.name, void, void), Null] + [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), + self.snes_ctx] ) - dmda_calls = self._create_dmda_calls(dmda, objs) + dmda_calls = self._create_dmda_calls(dmda) - mainctx = solver_objs['mainctx'] + mainctx = sobjs['userctx'] call_struct_callback = petsc_call( - cbbuilder.struct_callback.name, [Byref(mainctx)] + self.cbbuilder.user_struct_callback.name, [Byref(mainctx)] ) - calls_set_app_ctx = [ - petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) - ] - calls = [call_struct_callback] + calls_set_app_ctx + [BlankLine] + + # TODO: maybe don't need to explictly set this + mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) + + calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) base_setup = dmda_calls + ( snes_create, @@ -644,7 +1077,6 @@ def _setup(self, solver_objs, objs, injectsolve, cbbuilder): snes_set_type, global_x, global_b, - local_b, snes_get_ksp, ksp_set_tols, ksp_set_type, @@ -653,26 +1085,29 @@ def _setup(self, solver_objs, objs, injectsolve, cbbuilder): ksp_set_from_ops, matvec_operation, formfunc_operation, - ) + tuple(calls) - - extended_setup = self._extend_setup(solver_objs, objs, injectsolve, cbbuilder) - return base_setup + tuple(extended_setup) + call_struct_callback, + mat_set_dm, + calls_set_app_ctx, + BlankLine + ) + extended_setup = self._extend_setup() + return base_setup + extended_setup - def _extend_setup(self, solver_objs, objs, injectsolve, cbbuilder): + def _extend_setup(self): """ Hook for subclasses to add additional setup calls. """ - return [] + return () - def _create_dmda_calls(self, dmda, objs): - dmda_create = self._create_dmda(dmda, objs) + def _create_dmda_calls(self, dmda): + dmda_create = self._create_dmda(dmda) dm_setup = petsc_call('DMSetUp', [dmda]) dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) return dmda_create, dm_setup, dm_mat_type - def _create_dmda(self, dmda, objs): - grid = objs['grid'] - + def _create_dmda(self, dmda): + objs = self.objs + grid = self.fielddata.grid nspace_dims = len(grid.dimensions) # MPI communicator @@ -692,64 +1127,128 @@ def _create_dmda(self, dmda, objs): args.extend(list(grid.distributor.topology)[::-1]) # Number of degrees of freedom per node - args.append(1) + args.append(dmda.dofs) # "Stencil width" -> size of overlap - args.append(dmda.stencil_width) - args.extend([Null]*nspace_dims) + stencil_width = self.fielddata.space_order + args.append(stencil_width) + args.extend([objs['Null']]*nspace_dims) # The distributed array object args.append(Byref(dmda)) # The PETSc call used to create the DMDA - dmda = petsc_call('DMDACreate%sd' % nspace_dims, args) + dmda = petsc_call(f'DMDACreate{nspace_dims}d', args) return dmda -class Solver: - def __init__(self, solver_objs, objs, injectsolve, iters, cbbuilder, - timedep=None, **kwargs): - self.timedep = timedep - self.calls = self._execute_solve(solver_objs, objs, injectsolve, iters, cbbuilder) - self.spatial_body = self._spatial_loop_nest(iters, injectsolve) +class CoupledSetup(BaseSetup): + @property + def snes_ctx(self): + return Byref(self.solver_objs['jacctx']) + + def _extend_setup(self): + objs = self.objs + sobjs = self.solver_objs + + dmda = sobjs['dmda'] + create_field_decomp = petsc_call( + 'DMCreateFieldDecomposition', + [dmda, Byref(sobjs['nfields']), objs['Null'], Byref(sobjs['fields']), + Byref(sobjs['subdms'])] + ) + submat_cb = self.cbbuilder.submatrices_callback + matop_create_submats_op = petsc_call( + 'MatShellSetOperation', + [sobjs['Jac'], 'MATOP_CREATE_SUBMATRICES', + MatShellSetOp(submat_cb.name, void, void)] + ) + + call_coupled_struct_callback = petsc_call( + 'PopulateMatContext', + [Byref(sobjs['jacctx']), sobjs['subdms'], sobjs['fields']] + ) + + shell_set_ctx = petsc_call( + 'MatShellSetContext', [sobjs['Jac'], Byref(sobjs['jacctx']._C_symbol)] + ) - space_iter, = self.spatial_body - self.mapper = {space_iter: self.calls} + create_submats = petsc_call( + 'MatCreateSubMatrices', + [sobjs['Jac'], sobjs['nfields'], sobjs['fields'], + sobjs['fields'], 'MAT_INITIAL_MATRIX', + Byref(FieldFromComposite(objs['Submats'].base, sobjs['jacctx']))] + ) + + targets = self.fielddata.targets + + deref_dms = [ + DummyExpr(sobjs[f'da{t.name}'], sobjs['subdms'].indexed[i]) + for i, t in enumerate(targets) + ] - def _execute_solve(self, solver_objs, objs, injectsolve, iters, cbbuilder): + xglobals = [petsc_call( + 'DMCreateGlobalVector', + [sobjs[f'da{t.name}'], Byref(sobjs[f'xglobal{t.name}'])] + ) for t in targets] + + bglobals = [petsc_call( + 'DMCreateGlobalVector', + [sobjs[f'da{t.name}'], Byref(sobjs[f'bglobal{t.name}'])] + ) for t in targets] + + return ( + create_field_decomp, + matop_create_submats_op, + call_coupled_struct_callback, + shell_set_ctx, + create_submats + ) + tuple(deref_dms) + tuple(xglobals) + tuple(bglobals) + + +class Solver: + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.iters = kwargs.get('iters') + self.cbbuilder = kwargs.get('cbbuilder') + self.timedep = kwargs.get('timedep') + # TODO: Should/could _execute_solve be a cached_property? + self.calls = self._execute_solve() + + def _execute_solve(self): """ Assigns the required time iterators to the struct and executes the necessary calls to execute the SNES solver. """ - struct_assignment = self.timedep.assign_time_iters(solver_objs['mainctx']) + sobjs = self.solver_objs + target = self.injectsolve.expr.rhs.fielddata.target + + struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - rhs_callback = cbbuilder.formrhs_callback + rhs_callback = self.cbbuilder.formrhs[0] - dmda = solver_objs['dmda'] + dmda = sobjs['dmda'] - rhs_call = petsc_call(rhs_callback.name, list(rhs_callback.parameters)) + rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']]) local_x = petsc_call('DMCreateLocalVector', - [dmda, Byref(solver_objs['x_local'])]) + [dmda, Byref(sobjs['xlocal'])]) - vec_replace_array = self.timedep.replace_array(solver_objs) + vec_replace_array = self.timedep.replace_array(target) dm_local_to_global_x = petsc_call( - 'DMLocalToGlobal', [dmda, solver_objs['x_local'], 'INSERT_VALUES', - solver_objs['x_global']] - ) - - dm_local_to_global_b = petsc_call( - 'DMLocalToGlobal', [dmda, solver_objs['b_local'], 'INSERT_VALUES', - solver_objs['b_global']] + 'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_vals, + sobjs['xglobal']] ) snes_solve = petsc_call('SNESSolve', [ - solver_objs['snes'], solver_objs['b_global'], solver_objs['x_global']] + sobjs['snes'], sobjs['bglobal'], sobjs['xglobal']] ) dm_global_to_local_x = petsc_call('DMGlobalToLocal', [ - dmda, solver_objs['x_global'], 'INSERT_VALUES', solver_objs['x_local']] + dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']] ) run_solver_calls = (struct_assignment,) + ( @@ -757,37 +1256,120 @@ def _execute_solve(self, solver_objs, objs, injectsolve, iters, cbbuilder): local_x ) + vec_replace_array + ( dm_local_to_global_x, - dm_local_to_global_b, snes_solve, dm_global_to_local_x, BlankLine, ) return List(body=run_solver_calls) - def _spatial_loop_nest(self, iters, injectsolve): + @cached_property + def spatial_body(self): spatial_body = [] - for tree in retrieve_iteration_tree(iters[0]): + # TODO: remove the iters[0] + for tree in retrieve_iteration_tree(self.iters[0]): root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if injectsolve in FindNodes(PetscMetaData).visit(root): + if self.injectsolve in FindNodes(PetscMetaData).visit(root): spatial_body.append(root) + spatial_body, = spatial_body return spatial_body -class NonTimeDependent: - def __init__(self, injectsolve, iters, **kwargs): - self.injectsolve = injectsolve - self.iters = iters - self.kwargs = kwargs - self.origin_to_moddim = self._origin_to_moddim_mapper(iters) - self.time_idx_to_symb = injectsolve.expr.rhs.time_mapper +class CoupledSolver(Solver): + def _execute_solve(self): + """ + Assigns the required time iterators to the struct and executes + the necessary calls to execute the SNES solver. + """ + sobjs = self.solver_objs + + struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) + + rhs_callbacks = self.cbbuilder.formrhs + + xglob = sobjs['xglobal'] + bglob = sobjs['bglobal'] + + targets = self.injectsolve.expr.rhs.fielddata.targets + + # TODO: optimise the ccode generated here + pre_solve = () + post_solve = () + + for i, (c, t) in enumerate(zip(rhs_callbacks, targets)): + name = t.name + dm = sobjs[f'da{name}'] + target_xloc = sobjs[f'xlocal{name}'] + target_xglob = sobjs[f'xglobal{name}'] + target_bglob = sobjs[f'bglobal{name}'] + field = sobjs['fields'].indexed[i] + s = sobjs[f'scatter{name}'] + + pre_solve += ( + petsc_call(c.name, [dm, target_bglob]), + petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), + self.timedep.replace_array(t), + petsc_call( + 'DMLocalToGlobal', + [dm, target_xloc, insert_vals, target_xglob] + ), + petsc_call( + 'VecScatterCreate', + [xglob, field, target_xglob, self.objs['Null'], Byref(s)] + ), + petsc_call( + 'VecScatterBegin', + [s, target_xglob, xglob, insert_vals, sreverse] + ), + petsc_call( + 'VecScatterEnd', + [s, target_xglob, xglob, insert_vals, sreverse] + ), + petsc_call( + 'VecScatterBegin', + [s, target_bglob, bglob, insert_vals, sreverse] + ), + petsc_call( + 'VecScatterEnd', + [s, target_bglob, bglob, insert_vals, sreverse] + ), + ) - @property - def is_target_time(self): - return False + post_solve += ( + petsc_call( + 'VecScatterBegin', + [s, xglob, target_xglob, insert_vals, sforward] + ), + petsc_call( + 'VecScatterEnd', + [s, xglob, target_xglob, insert_vals, sforward] + ), + petsc_call( + 'DMGlobalToLocal', + [dm, target_xglob, insert_vals, target_xloc] + ) + ) - @property - def target(self): - return self.injectsolve.expr.rhs.target + snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], bglob, xglob]),) + + return List( + body=( + (struct_assignment,) + + pre_solve + + snes_solve + + post_solve + + (BlankLine,) + ) + ) + + +class NonTimeDependent: + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.iters = kwargs.get('iters') + self.sobjs = kwargs.get('solver_objs') + self.kwargs = kwargs + self.origin_to_moddim = self._origin_to_moddim_mapper(self.iters) + self.time_idx_to_symb = self.injectsolve.expr.rhs.time_mapper def _origin_to_moddim_mapper(self, iters): return {} @@ -795,7 +1377,7 @@ def _origin_to_moddim_mapper(self, iters): def uxreplace_time(self, body): return body - def replace_array(self, solver_objs): + def replace_array(self, target): """ VecReplaceArray() is a PETSc function that allows replacing the array of a `Vec` with a user provided array. @@ -806,19 +1388,19 @@ def replace_array(self, solver_objs): Examples -------- - >>> self.target + >>> target f1(x, y) - >>> call = replace_array(solver_objs) + >>> call = replace_array(target) >>> print(call) - PetscCall(VecReplaceArray(x_local_0,f1_vec->data)); + PetscCall(VecReplaceArray(xlocal0,f1_vec->data)); """ + sobjs = self.sobjs + field_from_ptr = FieldFromPointer( - self.target.function._C_field_data, self.target.function._C_symbol + target.function._C_field_data, target.function._C_symbol ) - vec_replace_array = (petsc_call( - 'VecReplaceArray', [solver_objs['x_local'], field_from_ptr] - ),) - return vec_replace_array + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + return (petsc_call('VecReplaceArray', [xlocal, field_from_ptr]),) def assign_time_iters(self, struct): return [] @@ -851,25 +1433,11 @@ class TimeDependent(NonTimeDependent): - Modulo dimensions are updated in the matrix context struct at each time step and can be accessed in the callback functions where needed. """ - @property - def is_target_time(self): - return any(i.is_Time for i in self.target.dimensions) - @property def time_spacing(self): - return self.target.grid.stepping_dim.spacing + return self.injectsolve.expr.rhs.grid.stepping_dim.spacing - @property - def target_time(self): - target_time = [ - i for i, d in zip(self.target.indices, self.target.dimensions) - if d.is_Time - ] - assert len(target_time) == 1 - target_time = target_time.pop() - return target_time - - @property + @cached_property def symb_to_moddim(self): """ Maps temporary `Symbol` objects created during `PETScSolve` to their @@ -881,6 +1449,18 @@ def symb_to_moddim(self): } return {symb: self.origin_to_moddim[mapper[symb]] for symb in mapper} + def is_target_time(self, target): + return any(i.is_Time for i in target.dimensions) + + def target_time(self, target): + target_time = [ + i for i, d in zip(target.indices, target.dimensions) + if d.is_Time + ] + assert len(target_time) == 1 + target_time = target_time.pop() + return target_time + def uxreplace_time(self, body): return Uxreplace(self.symb_to_moddim).visit(body) @@ -911,7 +1491,7 @@ def _origin_to_moddim_mapper(self, iters): mapper[d] = d return mapper - def replace_array(self, solver_objs): + def replace_array(self, target): """ In the case that the actual target is time-dependent e.g a `TimeFunction`, a pointer to the first element in the array that will be updated during @@ -919,51 +1499,45 @@ def replace_array(self, solver_objs): Examples -------- - >>> self.target + >>> target f1(time + dt, x, y) - >>> calls = replace_array(solver_objs) + >>> calls = replace_array(target) >>> print(List(body=calls)) - PetscCall(VecGetSize(x_local_0,&(localsize_0))); - float * start_ptr_0 = (time + 1)*localsize_0 + (float*)(f1_vec->data); - PetscCall(VecReplaceArray(x_local_0,start_ptr_0)); + PetscCall(VecGetSize(xlocal0,&(localsize0))); + float * f1_ptr0 = (time + 1)*localsize0 + (float*)(f1_vec->data); + PetscCall(VecReplaceArray(xlocal0,f1_ptr0)); - >>> self.target + >>> target f1(t + dt, x, y) - >>> calls = replace_array(solver_objs) + >>> calls = replace_array(target) >>> print(List(body=calls)) - PetscCall(VecGetSize(x_local_0,&(localsize_0))); - float * start_ptr_0 = t1*localsize_0 + (float*)(f1_vec->data); + PetscCall(VecGetSize(xlocal0,&(localsize0))); + float * f1_ptr0 = t1*localsize0 + (float*)(f1_vec->data); + PetscCall(VecReplaceArray(xlocal0,f1_ptr0)); """ - if self.is_target_time: - mapper = {self.time_spacing: 1, -self.time_spacing: -1} - target_time = self.target_time.xreplace(mapper) - - try: - target_time = self.origin_to_moddim[target_time] - except KeyError: - pass - - start_ptr = solver_objs['start_ptr'] - - vec_get_size = petsc_call( - 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] - ) - - field_from_ptr = FieldFromPointer( - self.target.function._C_field_data, self.target.function._C_symbol - ) + sobjs = self.sobjs - expr = DummyExpr( - start_ptr, cast_mapper[(self.target.dtype, '*')](field_from_ptr) + - Mul(target_time, solver_objs['localsize']), init=True - ) + if self.is_target_time(target): + mapper = {self.time_spacing: 1, -self.time_spacing: -1} - vec_replace_array = petsc_call( - 'VecReplaceArray', [solver_objs['x_local'], start_ptr] + target_time = self.target_time(target).xreplace(mapper) + target_time = self.origin_to_moddim.get(target_time, target_time) + + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + start_ptr = sobjs[f'{target.name}_ptr'] + + return ( + petsc_call('VecGetSize', [xlocal, Byref(sobjs['localsize'])]), + DummyExpr( + start_ptr, + cast_mapper[(target.dtype, '*')]( + FieldFromPointer(target._C_field_data, target._C_symbol) + ) + Mul(target_time, sobjs['localsize']), + init=True + ), + petsc_call('VecReplaceArray', [xlocal, start_ptr]) ) - return (vec_get_size, expr, vec_replace_array) - else: - return super().replace_array(solver_objs) + return super().replace_array(target) def assign_time_iters(self, struct): """ @@ -993,11 +1567,7 @@ def assign_time_iters(self, struct): return time_iter_assignments -Null = Macro('NULL') void = 'void' -dummyctx = Symbol('lctx') -dummyptr = DummyArg('dummy') - - -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') +insert_vals = 'INSERT_VALUES' +sreverse = 'SCATTER_REVERSE' +sforward = 'SCATTER_FORWARD' diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 4f825ad196..99da0468ad 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -10,11 +10,12 @@ def petsc_call_mpi(specific_call, call_args): return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)]) -def petsc_struct(name, fields, pname, liveness='lazy'): +def petsc_struct(name, fields, pname, liveness='lazy', modifier=None): # TODO: Fix this circular import from devito.petsc.types.object import PETScStruct return PETScStruct(name=name, pname=pname, - fields=fields, liveness=liveness) + fields=fields, liveness=liveness, + modifier=modifier) # Mapping special Eq operations to their corresponding IET Expression subclass types. diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index b00247ae2e..971fc8678b 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -8,76 +8,154 @@ from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions -from devito.tools import as_tuple -from devito.petsc.types import LinearSolveExpr, PETScArray, DMDALocalInfo +from devito.tools import as_tuple, filter_ordered +from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, + FieldData, MultipleFieldData, SubMatrices) -__all__ = ['PETScSolve'] +__all__ = ['PETScSolve', 'EssentialBC'] -def PETScSolve(eqns, target, solver_parameters=None, **kwargs): - prefixes = ['y_matvec', 'x_matvec', 'f_formfunc', 'x_formfunc', 'b_tmp'] - - localinfo = DMDALocalInfo(name='info', liveness='eager') - - arrays = { - p: PETScArray(name='%s_%s' % (p, target.name), - target=target, - liveness='eager', - localinfo=localinfo) - for p in prefixes - } - - matvecs = [] - formfuncs = [] - formrhs = [] +def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): + if target is not None: + return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() + else: + return InjectSolveNested(solver_parameters, target_eqns).build_eq() + + +class InjectSolve: + def __init__(self, solver_parameters=None, target_eqns=None): + self.solver_params = solver_parameters + self.time_mapper = None + self.target_eqns = target_eqns + + def build_eq(self): + target, funcs, fielddata = self.linear_solve_args() + # Placeholder equation for inserting calls to the solver + linear_solve = LinearSolveExpr( + funcs, + self.solver_params, + fielddata=fielddata, + time_mapper=self.time_mapper, + localinfo=localinfo + ) + return [PetscEq(target, linear_solve)] + + def linear_solve_args(self): + target, eqns = next(iter(self.target_eqns.items())) + eqns = as_tuple(eqns) + + funcs = get_funcs(eqns) + self.time_mapper = generate_time_mapper(funcs) + arrays = self.generate_arrays(target) + + return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) + + def generate_field_data(self, eqns, target, arrays): + formfuncs, formrhs = zip( + *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + ) + matvecs = [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + + return FieldData( + target=target, + matvecs=matvecs, + formfuncs=formfuncs, + formrhs=formrhs, + arrays=arrays + ) + + 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) - eqns = as_tuple(eqns) + return tuple(expr.subs(self.time_mapper) for expr in (formfunc, formrhs)) - for eq in eqns: + def build_matvec_eqns(self, eq, target, arrays): b, F_target, targets = separate_eqn(eq, target) - - # TODO: Current assumption is that problem is linear and user has not provided - # a jacobian. Hence, we can use F_target to form the jac-vec product - matvecs.append(Eq( - arrays['y_matvec'], - F_target.subs(targets_to_arrays(arrays['x_matvec'], targets)), - subdomain=eq.subdomain - )) - - formfuncs.append(Eq( - arrays['f_formfunc'], - F_target.subs(targets_to_arrays(arrays['x_formfunc'], targets)), - subdomain=eq.subdomain - )) - - formrhs.append(Eq( - arrays['b_tmp'], - b, - subdomain=eq.subdomain - )) - - funcs = retrieve_functions(eqns) - time_mapper = generate_time_mapper(funcs) - - matvecs, formfuncs, formrhs = ( - [eq.xreplace(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs) - ) - # Placeholder equation for inserting calls to the solver and generating - # correct time loop etc - inject_solve = PetscEq(target, LinearSolveExpr( - expr=tuple(funcs), - target=target, - solver_parameters=solver_parameters, - matvecs=matvecs, - formfuncs=formfuncs, - formrhs=formrhs, - arrays=arrays, - time_mapper=time_mapper, - localinfo=localinfo - )) - - return [inject_solve] + if not F_target: + return None + matvec = self.make_matvec(eq, F_target, arrays, targets) + return matvec.subs(self.time_mapper) + + def make_matvec(self, eq, F_target, arrays, targets): + rhs = arrays['x'] if isinstance(eq, EssentialBC) else F_target.subs( + targets_to_arrays(arrays['x'], targets) + ) + return Eq(arrays['y'], rhs, subdomain=eq.subdomain) + + def make_formfunc(self, eq, F_target, arrays, targets): + rhs = 0. if isinstance(eq, EssentialBC) else F_target.subs( + targets_to_arrays(arrays['x'], targets) + ) + return Eq(arrays['f'], rhs, subdomain=eq.subdomain) + + def make_rhs(self, eq, b, arrays): + rhs = 0. if isinstance(eq, EssentialBC) else b + return Eq(arrays['b'], rhs, subdomain=eq.subdomain) + + def generate_arrays(self, target): + return { + p: PETScArray(name=f'{p}_{target.name}', + target=target, + liveness='eager', + localinfo=localinfo) + for p in prefixes + } + + +class InjectSolveNested(InjectSolve): + def linear_solve_args(self): + combined_eqns = [] + for eqns in self.target_eqns.values(): + combined_eqns.extend(eqns) + funcs = get_funcs(combined_eqns) + self.time_mapper = generate_time_mapper(funcs) + + targets = list(self.target_eqns.keys()) + jacobian = SubMatrices(targets) + + all_data = MultipleFieldData(jacobian) + + for target, eqns in self.target_eqns.items(): + eqns = as_tuple(eqns) + arrays = self.generate_arrays(target) + + self.update_jacobian(eqns, target, jacobian, arrays) + + fielddata = self.generate_field_data( + eqns, target, arrays + ) + all_data.add_field_data(fielddata) + + return target, tuple(funcs), all_data + + def update_jacobian(self, eqns, target, jacobian, arrays): + for submat, mtvs in jacobian.submatrices[target].items(): + matvecs = [ + self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) + for eq in eqns + ] + # Set submatrix only if there's at least one non-zero matvec + if any(m is not None for m in matvecs): + jacobian.set_submatrix(target, submat, matvecs) + + def generate_field_data(self, eqns, target, arrays): + formfuncs, formrhs = zip( + *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + ) + + return FieldData( + target=target, + formfuncs=formfuncs, + formrhs=formrhs, + arrays=arrays + ) + + +class EssentialBC(Eq): + pass def separate_eqn(eqn, target): @@ -113,15 +191,12 @@ def targets_to_arrays(array, targets): """ Map each target in `targets` to a corresponding array generated from `array`, matching the spatial indices of the target. - Example: -------- >>> array vec_u(x, y) - >>> targets {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} - >>> targets_to_arrays(array, targets) {u(t + dt, x - h_x, y): vec_u(x - h_x, y), u(t + dt, x + h_x, y): vec_u(x + h_x, y), @@ -218,7 +293,6 @@ def generate_time_mapper(funcs): (and subsequently dropped and replaced with calls to run the solver). Therefore, the appropriate time loop will always be correctly generated inside the main kernel. - Examples -------- >>> funcs = [ @@ -229,7 +303,6 @@ def generate_time_mapper(funcs): >>> ] >>> generate_time_mapper(funcs) {t + dt: tau0, t: tau1} - """ time_indices = list({ i if isinstance(d, SteppingDimension) else d @@ -239,3 +312,16 @@ def generate_time_mapper(funcs): }) tau_symbs = [Symbol('tau%d' % i) for i in range(len(time_indices))] return dict(zip(time_indices, tau_symbs)) + + +def get_funcs(eqns): + funcs = [ + func + for eq in eqns + for func in retrieve_functions(eval_time_derivatives(eq.lhs - eq.rhs)) + ] + return filter_ordered(funcs) + + +localinfo = DMDALocalInfo(name='info', liveness='eager') +prefixes = ['y', 'x', 'f', 'b'] diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index ae9248d2a2..9acf7def46 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,13 +1,21 @@ from ctypes import POINTER, c_char - -from devito.tools import CustomDtype, dtype_to_cstr -from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension +from devito.tools import CustomDtype, dtype_to_cstr, as_tuple, CustomIntType +from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, + TimeDimension, ArrayObject, CustomDimension) +from devito.symbolics import Byref, Cast from devito.types.basic import DataSymbol -from devito.symbolics import Byref - from devito.petsc.iet.utils import petsc_call +class CallbackDM(LocalObject): + """ + PETSc Data Management object (DM). This is the DM instance + accessed within the callback functions via `SNESGetDM` and + is not destroyed during callback execution. + """ + dtype = CustomDtype('DM') + + class DM(LocalObject): """ PETSc Data Management object (DM). This is the primary DM instance @@ -16,43 +24,38 @@ class DM(LocalObject): """ dtype = CustomDtype('DM') - def __init__(self, *args, stencil_width=None, **kwargs): + def __init__(self, *args, dofs=1, **kwargs): super().__init__(*args, **kwargs) - self._stencil_width = stencil_width + self._dofs = dofs @property - def stencil_width(self): - return self._stencil_width + def dofs(self): + return self._dofs @property def _C_free(self): return petsc_call('DMDestroy', [Byref(self.function)]) + # TODO: This is growing out of hand so switch to an enumeration or something? @property def _C_free_priority(self): - return 3 + return 4 -class CallbackDM(LocalObject): - """ - PETSc Data Management object (DM). This is the DM instance - accessed within the callback functions via `SNESGetDM`. - """ - dtype = CustomDtype('DM') +class DMCast(Cast): + _base_typ = 'DM' - def __init__(self, *args, stencil_width=None, **kwargs): - super().__init__(*args, **kwargs) - self._stencil_width = stencil_width - @property - def stencil_width(self): - return self._stencil_width +class CallbackMat(LocalObject): + """ + PETSc Matrix object (Mat) used within callback functions. + These instances are not destroyed during callback execution; + instead, they are managed and destroyed in the main kernel. + """ + dtype = CustomDtype('Mat') class Mat(LocalObject): - """ - PETSc Matrix object (Mat). - """ dtype = CustomDtype('Mat') @property @@ -61,7 +64,7 @@ def _C_free(self): @property def _C_free_priority(self): - return 1 + return 2 class LocalVec(LocalObject): @@ -73,21 +76,26 @@ class LocalVec(LocalObject): dtype = CustomDtype('Vec') -class GlobalVec(LocalObject): +class CallbackGlobalVec(LocalVec): + """ + PETSc global vector object (Vec). For example, used for coupled + solves inside the `WholeFormFunc` callback. + """ + + +class GlobalVec(LocalVec): """ PETSc global vector object (Vec). A global vector is a parallel vector that has no duplicate values between MPI ranks. A global vector has no ghost locations. """ - dtype = CustomDtype('Vec') - @property def _C_free(self): return petsc_call('VecDestroy', [Byref(self.function)]) @property def _C_free_priority(self): - return 0 + return 1 class PetscMPIInt(LocalObject): @@ -103,7 +111,7 @@ class PetscInt(LocalObject): PETSc datatype used to represent `int` parameters to PETSc functions. """ - dtype = CustomDtype('PetscInt') + dtype = CustomIntType('PetscInt') class KSP(LocalObject): @@ -114,19 +122,21 @@ class KSP(LocalObject): dtype = CustomDtype('KSP') -class SNES(LocalObject): +class CallbackSNES(LocalObject): """ PETSc SNES : Non-Linear Systems Solvers. """ dtype = CustomDtype('SNES') + +class SNES(CallbackSNES): @property def _C_free(self): return petsc_call('SNESDestroy', [Byref(self.function)]) @property def _C_free_priority(self): - return 2 + return 3 class PC(LocalObject): @@ -168,18 +178,25 @@ class DummyArg(LocalObject): dtype = CustomDtype('void', modifier='*') -class PETScStruct(CCompositeObject): +class MatReuse(LocalObject): + dtype = CustomDtype('MatReuse') - __rargs__ = ('name', 'pname', 'fields') - def __init__(self, name, pname, fields, liveness='lazy'): - pfields = [(i._C_name, i._C_ctype) for i in fields] - super().__init__(name, pname, pfields, liveness) - self._fields = fields +class VecScatter(LocalObject): + dtype = CustomDtype('VecScatter') - @property - def fields(self): - return self._fields + +class StartPtr(LocalObject): + def __init__(self, name, dtype): + super().__init__(name=name) + self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + + +class SingleIS(LocalObject): + dtype = CustomDtype('IS') + + +class PETScStruct(LocalCompositeObject): @property def time_dim_fields(self): @@ -198,18 +215,109 @@ def callback_fields(self): """ return [f for f in self.fields if f not in self.time_dim_fields] + _C_modifier = ' *' + + +class JacobianStruct(PETScStruct): + def __init__(self, name='jctx', pname='JacobianCtx', fields=None, + modifier='', liveness='lazy'): + super().__init__(name, pname, fields, modifier, liveness) + _C_modifier = None + + +class SubMatrixStruct(PETScStruct): + def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, + modifier=' *', liveness='lazy'): + super().__init__(name, pname, fields, modifier, liveness) + _C_modifier = None + + +class JacobianStructCast(Cast): + _base_typ = 'struct JacobianCtx *' + + +class PETScArrayObject(ArrayObject): + _data_alignment = False + + def __init_finalize__(self, *args, **kwargs): + self._nindices = kwargs.pop('nindices', 1) + super().__init_finalize__(*args, **kwargs) + + @classmethod + def __indices_setup__(cls, **kwargs): + try: + return as_tuple(kwargs['dimensions']), as_tuple(kwargs['dimensions']) + except KeyError: + nindices = kwargs.get('nindices', 1) + dim = CustomDimension(name='d', symbolic_size=nindices) + return (dim,), (dim,) + @property - def _C_ctype(self): - return POINTER(self.dtype) if self.liveness == \ - 'eager' else self.dtype + def dim(self): + assert len(self.dimensions) == 1 + return self.dimensions[0] - _C_modifier = ' *' + @property + def nindices(self): + return self._nindices + @property + def _C_name(self): + return self.name -class StartPtr(LocalObject): - def __init__(self, name, dtype): - super().__init__(name=name) - self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + @property + def _mem_stack(self): + return False + + @property + def _C_free_priority(self): + return 0 + + +class CallbackPointerIS(PETScArrayObject): + """ + Index set object used for efficient indexing into vectors and matrices. + https://petsc.org/release/manualpages/IS/IS/ + """ + @property + def dtype(self): + return CustomDtype('IS', modifier=' *') + + +class PointerIS(CallbackPointerIS): + @property + def _C_free(self): + destroy_calls = [ + petsc_call('ISDestroy', [Byref(self.indexify().subs({self.dim: i}))]) + for i in range(self._nindices) + ] + destroy_calls.append(petsc_call('PetscFree', [self.function])) + return destroy_calls + + +class CallbackPointerDM(PETScArrayObject): + @property + def dtype(self): + return CustomDtype('DM', modifier=' *') + + +class PointerDM(CallbackPointerDM): + @property + def _C_free(self): + destroy_calls = [ + petsc_call('DMDestroy', [Byref(self.indexify().subs({self.dim: i}))]) + for i in range(self._nindices) + ] + destroy_calls.append(petsc_call('PetscFree', [self.function])) + return destroy_calls + + +class PointerMat(PETScArrayObject): + _C_modifier = ' *' + + @property + def dtype(self): + return CustomDtype('Mat', modifier=' *') class ArgvSymbol(DataSymbol): diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index c3b43b84af..c7c3397181 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -29,20 +29,16 @@ class LinearSolveExpr(MetaData): needed to execute a linear solver. Linear problems are handled with `SNESSetType(snes, KSPONLY)`, enabling a unified interface for both linear and nonlinear solvers. - # TODO: extend this defaults: - 'ksp_type': String with the name of the PETSc Krylov method. Default is 'gmres' (Generalized Minimal Residual Method). https://petsc.org/main/manualpages/KSP/KSPType/ - - 'pc_type': String with the name of the PETSc preconditioner. Default is 'jacobi' (i.e diagonal scaling preconditioning). https://petsc.org/main/manualpages/PC/PCType/ - KSP tolerances: https://petsc.org/release/manualpages/KSP/KSPSetTolerances/ - - 'ksp_rtol': Relative convergence tolerance. Default is 1e-5. - 'ksp_atol': Absolute convergence for tolerance. Default @@ -55,8 +51,7 @@ class LinearSolveExpr(MetaData): """ __rargs__ = ('expr',) - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays', 'time_mapper', + __rkwargs__ = ('solver_parameters', 'fielddata', 'time_mapper', 'localinfo') defaults = { @@ -68,9 +63,8 @@ class LinearSolveExpr(MetaData): 'ksp_max_it': 1e4 # Maximum iterations } - def __new__(cls, expr, target=None, solver_parameters=None, - matvecs=None, formfuncs=None, formrhs=None, - arrays=None, time_mapper=None, localinfo=None, **kwargs): + def __new__(cls, expr, solver_parameters=None, + fielddata=None, time_mapper=None, localinfo=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -82,12 +76,8 @@ def __new__(cls, expr, target=None, solver_parameters=None, obj = sympy.Function.__new__(cls, expr) obj._expr = expr - obj._target = target obj._solver_parameters = solver_parameters - obj._matvecs = matvecs - obj._formfuncs = formfuncs - obj._formrhs = formrhs - obj._arrays = arrays + obj._fielddata = fielddata if fielddata else FieldData() obj._time_mapper = time_mapper obj._localinfo = localinfo return obj @@ -105,21 +95,52 @@ def __hash__(self): def __eq__(self, other): return (isinstance(other, LinearSolveExpr) and - self.expr == other.expr and - self.target == other.target) + self.expr == other.expr) @property def expr(self): return self._expr @property - def target(self): - return self._target + def fielddata(self): + return self._fielddata @property def solver_parameters(self): return self._solver_parameters + @property + def time_mapper(self): + return self._time_mapper + + @property + def localinfo(self): + return self._localinfo + + @property + def grid(self): + return self.fielddata.grid + + @classmethod + def eval(cls, *args): + return None + + func = Reconstructable._rebuild + + +class FieldData: + def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, + arrays=None, **kwargs): + self._target = kwargs.get('target', target) + self._matvecs = matvecs + self._formfuncs = formfuncs + self._formrhs = formrhs + self._arrays = arrays + + @property + def target(self): + return self._target + @property def matvecs(self): return self._matvecs @@ -137,15 +158,160 @@ def arrays(self): return self._arrays @property - def time_mapper(self): - return self._time_mapper + def space_dimensions(self): + return self.target.space_dimensions @property - def localinfo(self): - return self._localinfo + def grid(self): + return self.target.grid - @classmethod - def eval(cls, *args): + @property + def space_order(self): + return self.target.space_order + + @property + def targets(self): + return (self.target,) + + +class MultipleFieldData(FieldData): + def __init__(self, submatrices=None): + self.field_data_list = [] + self._submatrices = submatrices + + def add_field_data(self, field_data): + self.field_data_list.append(field_data) + + def get_field_data(self, target): + for field_data in self.field_data_list: + if field_data.target == target: + return field_data + raise ValueError(f"FieldData with target {target} not found.") + pass + + @property + def target(self): return None - func = Reconstructable._rebuild + @property + def targets(self): + return tuple(field_data.target for field_data in self.field_data_list) + + @property + def space_dimensions(self): + space_dims = {field_data.space_dimensions for field_data in self.field_data_list} + if len(space_dims) > 1: + # TODO: This may not actually have to be the case, but enforcing it for now + raise ValueError( + "All targets within a PETScSolve have to have the same space dimensions." + ) + return space_dims.pop() + + @property + def grid(self): + grids = [t.grid for t in self.targets] + if len(set(grids)) > 1: + raise ValueError( + "All targets within a PETScSolve have to have the same grid." + ) + return grids.pop() + + @property + def space_order(self): + # NOTE: since we use DMDA to create vecs for the coupled solves, + # all fields must have the same space order + # ... re think this? limitation. For now, just force the + # space order to be the same. + # This isn't a problem for segregated solves. + space_orders = [t.space_order for t in self.targets] + if len(set(space_orders)) > 1: + raise ValueError( + "All targets within a PETScSolve have to have the same space order." + ) + return space_orders.pop() + + @property + def submatrices(self): + return self._submatrices + + +class SubMatrices: + def __init__(self, targets): + self.targets = targets + self.submatrices = self._initialize_submatrices() + + def _initialize_submatrices(self): + """ + Create a dict of submatrices for each target with metadata. + """ + submatrices = {} + num_targets = len(self.targets) + + for i, target in enumerate(self.targets): + submatrices[target] = {} + for j in range(num_targets): + key = f'J{i}{j}' + submatrices[target][key] = { + 'matvecs': None, + 'derivative_wrt': self.targets[j], + 'index': i * num_targets + j + } + + return submatrices + + @property + def submatrix_keys(self): + """ + Return a list of all submatrix keys (e.g., ['J00', 'J01', 'J10', 'J11']). + """ + return [key for submats in self.submatrices.values() for key in submats.keys()] + + @property + def nonzero_submatrix_keys(self): + """ + Returns a list of submats where 'matvecs' is not None. + """ + return [ + key + for submats in self.submatrices.values() + for key, value in submats.items() + if value['matvecs'] is not None + ] + + @property + def submat_to_index(self): + """ + Returns a dict mapping submatrix keys to their index. + """ + return { + key: value['index'] + for submats in self.submatrices.values() + for key, value in submats.items() + } + + def set_submatrix(self, field, key, matvecs): + """ + Set a specific submatrix for a field. + + Parameters + ---------- + field : Function + The target field that the submatrix operates on. + key: str + The identifier for the submatrix (e.g., 'J00', 'J01'). + matvecs: list of Eq + The matrix-vector equations forming the submatrix. + """ + if field in self.submatrices and key in self.submatrices[field]: + self.submatrices[field][key]["matvecs"] = matvecs + else: + raise KeyError(f'Invalid field ({field}) or submatrix key ({key})') + + def get_submatrix(self, field, key): + """ + Retrieve a specific submatrix. + """ + return self.submatrices.get(field, {}).get(key, None) + + def __repr__(self): + return str(self.submatrices) diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index 9db65d8bb2..72413bcaaa 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -10,7 +10,7 @@ from devito.finite_differences.elementary import Min, Max from devito.tools import (Pickable, Bunch, as_tuple, is_integer, float2, # noqa float3, float4, double2, double3, double4, int2, int3, - int4) + int4, CustomIntType) from devito.types import Symbol from devito.types.basic import Basic @@ -19,8 +19,8 @@ 'ListInitializer', 'Byref', 'IndexedPointer', 'Cast', 'DefFunction', 'MathFunction', 'InlineIf', 'ReservedWord', 'Keyword', 'String', 'Macro', 'Class', 'MacroArgument', 'CustomType', 'Deref', 'Namespace', - 'Rvalue', 'INT', 'FLOAT', 'DOUBLE', 'VOID', 'Null', 'SizeOf', 'rfunc', - 'cast_mapper', 'BasicWrapperMixin', 'ValueLimit', 'limits_mapper'] + 'Rvalue', 'INT', 'FLOAT', 'DOUBLE', 'VOID', 'VOIDP', 'Null', 'SizeOf', 'rfunc', + 'cast_mapper', 'BasicWrapperMixin', 'ValueLimit', 'limits_mapper', 'Mod'] class CondEq(sympy.Eq): @@ -90,9 +90,16 @@ def __new__(cls, lhs, rhs, params=None): # Perhaps it's a symbolic RHS -- but we wanna be sure it's of type int if not hasattr(rhs, 'dtype'): raise ValueError("Symbolic RHS `%s` lacks dtype" % rhs) - if not issubclass(rhs.dtype, np.integer): - raise ValueError("Symbolic RHS `%s` must be of type `int`, found " - "`%s` instead" % (rhs, rhs.dtype)) + + # TODO: Move into a utility function? + is_int_type = isinstance(rhs.dtype, type) and \ + issubclass(rhs.dtype, np.integer) + is_custom_int_type = isinstance(rhs.dtype, CustomIntType) + assert is_int_type or is_custom_int_type, ( + f"Symbolic RHS `{rhs}` must be of type `int`, " + f"found `{rhs.dtype}` instead" + ) + rhs = sympify(rhs) obj = sympy.Expr.__new__(cls, lhs, rhs) @@ -115,6 +122,26 @@ def __mul__(self, other): return super().__mul__(other) +class Mod(sympy.Expr): + # TODO: Add tests + is_Atom = True + is_commutative = True + + def __new__(cls, lhs, rhs, params=None): + rhs = sympify(rhs) + + obj = sympy.Expr.__new__(cls, lhs, rhs) + + obj.lhs = lhs + obj.rhs = rhs + return obj + + def __str__(self): + return "Mod(%s, %s)" % (self.lhs, self.rhs) + + __repr__ = __str__ + + class BasicWrapperMixin: """ @@ -167,7 +194,7 @@ def __new__(cls, call, pointer, params=None, **kwargs): pointer = Symbol(pointer) if isinstance(call, str): call = Symbol(call) - elif not isinstance(call, Basic): + elif not isinstance(call.base, Basic): raise ValueError("`call` must be a `devito.Basic` or a type " "with compatible interface") _params = [] @@ -823,6 +850,10 @@ class VOID(Cast): _base_typ = 'void' +class VOIDP(CastStar): + base = VOID + + class CHARP(CastStar): base = CHAR diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index b5b564a4d7..7c3309d404 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -13,7 +13,7 @@ 'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype', 'dtype_to_cstr', 'dtype_to_ctype', 'dtype_to_mpitype', 'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', - 'is_external_ctype', 'infer_dtype', 'CustomDtype'] + 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'CustomIntType'] # *** Custom np.dtypes @@ -123,6 +123,11 @@ def __repr__(self): __str__ = __repr__ +# TODO: Consider if this should be an instance instead of a subclass? +class CustomIntType(CustomDtype): + pass + + # *** np.dtypes lowering @@ -278,6 +283,8 @@ def is_external_ctype(ctype, includes): True if `ctype` is known to be declared in one of the given `includes` files, False otherwise. """ + if isinstance(ctype, CustomDtype): + return False # Get the base type while issubclass(ctype, ctypes._Pointer): ctype = ctype._type_ diff --git a/devito/types/array.py b/devito/types/array.py index 44d7fabf9d..d2be10eb67 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -60,6 +60,21 @@ def shape_allocated(self): def is_const(self): return self._is_const + @property + def _C_free(self): + """ + A symbolic destructor for the Array, injected in the generated code. + + Notes + ----- + To be overridden by subclasses, ignored otherwise. + """ + return None + + @property + def _C_free_priority(self): + return 0 + class Array(ArrayBasic): diff --git a/devito/types/object.py b/devito/types/object.py index aa738bd19b..df8f94c171 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -1,14 +1,14 @@ from ctypes import byref import sympy -from devito.tools import Pickable, as_tuple, sympy_mutex +from devito.tools import Pickable, as_tuple, sympy_mutex, CustomDtype from devito.types.args import ArgProvider from devito.types.caching import Uncached from devito.types.basic import Basic, LocalType from devito.types.utils import CtypesFactory -__all__ = ['Object', 'LocalObject', 'CompositeObject', 'CCompositeObject'] +__all__ = ['Object', 'LocalObject', 'CompositeObject', 'LocalCompositeObject'] class AbstractObject(Basic, sympy.Basic, Pickable): @@ -241,19 +241,30 @@ def _mem_global(self): return self._is_global -class CCompositeObject(CompositeObject, LocalType): +class LocalCompositeObject(CompositeObject, LocalType): """ Object with composite type (e.g., a C struct) defined in C. """ - __rargs__ = ('name', 'pname', 'pfields') + __rargs__ = ('name', 'pname', 'fields') - def __init__(self, name, pname, pfields, liveness='lazy'): - super().__init__(name, pname, pfields) + def __init__(self, name, pname, fields, modifier=None, liveness='lazy'): + dtype = CustomDtype(f"struct {pname}", modifier=modifier) + Object.__init__(self, name, dtype, None) + self._pname = pname assert liveness in ['eager', 'lazy'] self._liveness = liveness + self._fields = fields @property - def dtype(self): - return self._dtype._type_ + def fields(self): + return self._fields + + @property + def _fields_(self): + return [(i._C_name, i._C_ctype) for i in self.fields] + + @property + def __name__(self): + return self.pname diff --git a/tests/test_petsc.py b/tests/test_petsc.py index dafc22211c..1f0f8bc91e 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -4,14 +4,14 @@ from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, switchconfig, - configuration) + configuration, norm) from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, FindNodes, retrieve_iteration_tree) -from devito.types import Constant, CCompositeObject +from devito.types import Constant, LocalCompositeObject from devito.passes.iet.languages.C import CDataManager from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, - LinearSolveExpr) + LinearSolveExpr, FieldData, MultipleFieldData) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -126,20 +126,21 @@ def test_petsc_solve(): callable_roots = [meta_call.root for meta_call in op._func_table.values()] - matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_0'] + matvec_callback = [root for root in callable_roots if root.name == 'MatMult0'] - formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS_0'] + formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS0'] action_expr = FindNodes(Expression).visit(matvec_callback[0]) rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) - assert str(action_expr[-1].expr.rhs) == \ - 'x_matvec_f[x + 1, y + 2]/lctx->h_x**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_x**2' + \ - ' + x_matvec_f[x + 3, y + 2]/lctx->h_x**2' + \ - ' + x_matvec_f[x + 2, y + 1]/lctx->h_y**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_y**2' + \ - ' + x_matvec_f[x + 2, y + 3]/lctx->h_y**2' + assert str(action_expr[-1].expr.rhs) == ( + 'x_f[x + 1, y + 2]/ctx0->h_x**2' + ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_x**2' + ' + x_f[x + 3, y + 2]/ctx0->h_x**2' + ' + x_f[x + 2, y + 1]/ctx0->h_y**2' + ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_y**2' + ' + x_f[x + 2, y + 3]/ctx0->h_y**2' + ) assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' @@ -213,12 +214,12 @@ def test_petsc_cast(): cb2 = [meta_call.root for meta_call in op2._func_table.values()] cb3 = [meta_call.root for meta_call in op3._func_table.values()] - assert 'float (*restrict x_matvec_f1) = ' + \ - '(float (*)) x_matvec_f1_vec;' in str(cb1[0]) - assert 'float (*restrict x_matvec_f2)[info.gxm] = ' + \ - '(float (*)[info.gxm]) x_matvec_f2_vec;' in str(cb2[0]) - assert 'float (*restrict x_matvec_f3)[info.gym][info.gxm] = ' + \ - '(float (*)[info.gym][info.gxm]) x_matvec_f3_vec;' in str(cb3[0]) + assert 'float (*restrict x_f1) = ' + \ + '(float (*)) x_f1_vec;' in str(cb1[0]) + assert 'float (*restrict x_f2)[info.gxm] = ' + \ + '(float (*)[info.gxm]) x_f2_vec;' in str(cb2[0]) + assert 'float (*restrict x_f3)[info.gym][info.gxm] = ' + \ + '(float (*)[info.gym][info.gxm]) x_f3_vec;' in str(cb3[0]) @skipif('petsc') @@ -233,8 +234,9 @@ def test_LinearSolveExpr(): linsolveexpr = LinearSolveExpr(eqn.rhs, target=f) - # Check the target - assert linsolveexpr.target == f + # TODO: maybe expand this test now to check the fielddata etc + linsolveexpr = LinearSolveExpr(eqn.rhs) + # Check the solver parameters assert linsolveexpr.solver_parameters == \ {'ksp_type': 'gmres', 'pc_type': 'jacobi', 'ksp_rtol': 1e-05, @@ -266,15 +268,15 @@ def test_dmda_create(): op3 = Operator(petsc3, opt='noop') assert 'PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - '2,1,2,NULL,&(da_0)));' in str(op1) + '2,1,2,NULL,&(da0)));' in str(op1) assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_0)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da0)));' \ in str(op2) assert 'PetscCall(DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ - ',1,1,1,1,6,NULL,NULL,NULL,&(da_0)));' in str(op3) + ',1,1,1,1,6,NULL,NULL,NULL,&(da0)));' in str(op3) @skipif('petsc') @@ -301,9 +303,9 @@ def test_cinterface_petsc_struct(): assert 'include "%s.h"' % name in ccode - # The public `struct MatContext` only appears in the header file - assert 'struct J_0_ctx\n{' not in ccode - assert 'struct J_0_ctx\n{' in hcode + # The public `struct UserCtx` only appears in the header file + assert 'struct UserCtx0\n{' not in ccode + assert 'struct UserCtx0\n{' in hcode @skipif('petsc') @@ -571,14 +573,14 @@ def test_callback_arguments(): with switchconfig(openmp=False): op = Operator(petsc1) - mv = op._func_table['MyMatShellMult_0'].root - ff = op._func_table['FormFunction_0'].root + mv = op._func_table['MatMult0'].root + ff = op._func_table['FormFunction0'].root assert len(mv.parameters) == 3 assert len(ff.parameters) == 4 - assert str(mv.parameters) == '(J_0, X_global_0, Y_global_0)' - assert str(ff.parameters) == '(snes_0, X_global_0, F_global_0, dummy)' + assert str(mv.parameters) == '(J, X, Y)' + assert str(ff.parameters) == '(snes, X, F, dummy)' @skipif('petsc') @@ -611,7 +613,7 @@ def test_petsc_struct(): assert mu2 in op.parameters # Check PETSc struct not in op.parameters - assert all(not isinstance(i, CCompositeObject) for i in op.parameters) + assert all(not isinstance(i, LocalCompositeObject) for i in op.parameters) @skipif('petsc') @@ -630,8 +632,7 @@ def test_apply(): # Build the op op = Operator(petsc) - # Check the Operator runs without errors. Not verifying output for - # now. Need to consolidate BC implementation + # Check the Operator runs without errors op.apply() # Verify that users can override `mu` @@ -656,11 +657,11 @@ def test_petsc_frees(): frees = op.body.frees # Check the frees appear in the following order - assert str(frees[0]) == 'PetscCall(VecDestroy(&(b_global_0)));' - assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_0)));' - assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_0)));' - assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_0)));' - assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_0)));' + assert str(frees[0]) == 'PetscCall(VecDestroy(&(bglobal0)));' + assert str(frees[1]) == 'PetscCall(VecDestroy(&(xglobal0)));' + assert str(frees[2]) == 'PetscCall(MatDestroy(&(J0)));' + assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes0)));' + assert str(frees[4]) == 'PetscCall(DMDestroy(&(da0)));' @skipif('petsc') @@ -679,8 +680,8 @@ def test_calls_to_callbacks(): ccode = str(op.ccode) - assert '(void (*)(void))MyMatShellMult_0' in ccode - assert 'PetscCall(SNESSetFunction(snes_0,NULL,FormFunction_0,NULL));' in ccode + assert '(void (*)(void))MatMult0' in ccode + assert 'PetscCall(SNESSetFunction(snes0,NULL,FormFunction0,(void*)(da0)));' in ccode @skipif('petsc') @@ -702,7 +703,7 @@ def test_start_ptr(): op1 = Operator(petsc1) # Verify the case with modulo time stepping - assert 'float * start_ptr_0 = t1*localsize_0 + (float*)(u1_vec->data);' in str(op1) + assert 'float * u1_ptr0 = t1*localsize0 + (float*)(u1_vec->data);' in str(op1) # Verify the case with no modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, dtype=np.float32, save=5) @@ -712,7 +713,7 @@ def test_start_ptr(): with switchconfig(openmp=False): op2 = Operator(petsc2) - assert 'float * start_ptr_0 = (time + 1)*localsize_0 + ' + \ + assert 'float * u2_ptr0 = (time + 1)*localsize0 + ' + \ '(float*)(u2_vec->data);' in str(op2) @@ -735,12 +736,12 @@ def test_time_loop(): with switchconfig(openmp=False): op1 = Operator(petsc1) body1 = str(op1.body) - rhs1 = str(op1._func_table['FormRHS_0'].root.ccode) + rhs1 = str(op1._func_table['FormRHS0'].root.ccode) assert 'ctx0.t0 = t0' in body1 assert 'ctx0.t1 = t1' not in body1 - assert 'lctx->t0' in rhs1 - assert 'lctx->t1' not in rhs1 + assert 'ctx0->t0' in rhs1 + assert 'ctx0->t1' not in rhs1 # Non-modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) @@ -750,10 +751,10 @@ def test_time_loop(): with switchconfig(openmp=False): op2 = Operator(petsc2) body2 = str(op2.body) - rhs2 = str(op2._func_table['FormRHS_0'].root.ccode) + rhs2 = str(op2._func_table['FormRHS0'].root.ccode) assert 'ctx0.time = time' in body2 - assert 'lctx->time' in rhs2 + assert 'ctx0->time' in rhs2 # Modulo time stepping with more than one time step # used in one of the callback functions @@ -762,12 +763,12 @@ def test_time_loop(): with switchconfig(openmp=False): op3 = Operator(petsc3) body3 = str(op3.body) - rhs3 = str(op3._func_table['FormRHS_0'].root.ccode) + rhs3 = str(op3._func_table['FormRHS0'].root.ccode) assert 'ctx0.t0 = t0' in body3 assert 'ctx0.t1 = t1' in body3 - assert 'lctx->t0' in rhs3 - assert 'lctx->t1' in rhs3 + assert 'ctx0->t0' in rhs3 + assert 'ctx0->t1' in rhs3 # Multiple petsc solves within the same time loop v2 = Function(name='v2', grid=grid, space_order=2) @@ -781,3 +782,245 @@ def test_time_loop(): assert 'ctx0.t0 = t0' in body4 assert body4.count('ctx0.t0 = t0') == 1 + + +class TestCoupledLinear: + # The coupled interface can be used even for uncoupled problems, meaning + # the equations will be solved within a single matrix system. + # These tests use simple problems to validate functionality, but they help + # ensure correctness in code generation. + # TODO: Add more comprehensive tests for fully coupled problems. + # TODO: Add subdomain tests, time loop, multiple coupled etc. + + @skipif('petsc') + def test_coupled_vs_non_coupled(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + f.data[:] = 5. + h.data[:] = 5. + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + # Non-coupled + petsc1 = PETScSolve(eq1, target=e) + petsc2 = PETScSolve(eq2, target=g) + + with switchconfig(openmp=False): + op1 = Operator(petsc1 + petsc2, opt='noop') + op1.apply() + + enorm1 = norm(e) + gnorm1 = norm(g) + + # Reset + e.data[:] = 0 + g.data[:] = 0 + + # Coupled + # TODO: Need more friendly API for coupled - just + # using a dict for now + petsc3 = PETScSolve({e: [eq1], g: [eq2]}) + with switchconfig(openmp=False): + op2 = Operator(petsc3, opt='noop') + op2.apply() + + 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) + + callbacks1 = [meta_call.root for meta_call in op1._func_table.values()] + callbacks2 = [meta_call.root for meta_call in op2._func_table.values()] + + # Solving for multiple fields within the same matrix system requires + # additional machinery and more callback functions + assert len(callbacks1) == 8 + assert len(callbacks2) == 11 + + # Check fielddata type + fielddata1 = petsc1[0].rhs.fielddata + fielddata2 = petsc2[0].rhs.fielddata + fielddata3 = petsc3[0].rhs.fielddata + + assert isinstance(fielddata1, FieldData) + assert isinstance(fielddata2, FieldData) + assert isinstance(fielddata3, MultipleFieldData) + + @skipif('petsc') + def test_coupled_structs(self): + grid = Grid(shape=(11, 11)) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e + 5, f) + eq2 = Eq(g + 10, h) + + petsc = PETScSolve({f: [eq1], h: [eq2]}) + + name = "foo" + with switchconfig(openmp=False): + op = Operator(petsc, name=name) + + # Trigger the generation of a .c and a .h files + ccode, hcode = op.cinterface(force=True) + + dirname = op._compiler.get_jit_dir() + assert os.path.isfile(os.path.join(dirname, f"{name}.c")) + assert os.path.isfile(os.path.join(dirname, f"{name}.h")) + + ccode = str(ccode) + hcode = str(hcode) + + assert f'include "{name}.h"' in ccode + + # The public `struct JacobianCtx` only appears in the header file + assert 'struct JacobianCtx\n{' not in ccode + assert 'struct JacobianCtx\n{' in hcode + + # The public `struct SubMatrixCtx` only appears in the header file + assert 'struct SubMatrixCtx\n{' not in ccode + assert 'struct SubMatrixCtx\n{' in hcode + + # The public `struct UserCtx0` only appears in the header file + assert 'struct UserCtx0\n{' not in ccode + assert 'struct UserCtx0\n{' in hcode + + @skipif('petsc') + def test_coupled_frees(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, h) + eq2 = Eq(f.laplace, h) + eq3 = Eq(g.laplace, h) + + petsc1 = PETScSolve({e: [eq1], f: [eq2]}) + petsc2 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + + with switchconfig(openmp=False): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + + frees1 = op1.body.frees + frees2 = op2.body.frees + + # Check solver with two fields + # IS destroys + assert str(frees1[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' + assert str(frees1[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' + assert str(frees1[2]) == 'PetscCall(PetscFree(fields0));' + # Sub DM destroys + assert str(frees1[3]) == 'PetscCall(DMDestroy(&(subdms0[0])));' + assert str(frees1[4]) == 'PetscCall(DMDestroy(&(subdms0[1])));' + assert str(frees1[5]) == 'PetscCall(PetscFree(subdms0));' + + # Check solver with three fields + # IS destroys + assert str(frees2[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' + assert str(frees2[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' + assert str(frees2[2]) == 'PetscCall(ISDestroy(&(fields0[2])));' + assert str(frees2[3]) == 'PetscCall(PetscFree(fields0));' + # Sub DM destroys + assert str(frees2[4]) == 'PetscCall(DMDestroy(&(subdms0[0])));' + assert str(frees2[5]) == 'PetscCall(DMDestroy(&(subdms0[1])));' + assert str(frees2[6]) == 'PetscCall(DMDestroy(&(subdms0[2])));' + assert str(frees2[7]) == 'PetscCall(PetscFree(subdms0));' + + @skipif('petsc') + def test_dmda_dofs(self): + grid = Grid(shape=(11, 11)) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, h) + eq2 = Eq(f.laplace, h) + eq3 = Eq(g.laplace, h) + + petsc1 = PETScSolve({e: [eq1]}) + petsc2 = PETScSolve({e: [eq1], f: [eq2]}) + petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + + with switchconfig(openmp=False): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + op3 = Operator(petsc3, opt='noop') + + # Check the number of dofs in the DMDA for each field + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,1,2,NULL,NULL,&(da0)));' \ + in str(op1) + + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,2,2,NULL,NULL,&(da0)));' \ + in str(op2) + + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,3,2,NULL,NULL,&(da0)));' \ + in str(op3) + + @skipif('petsc') + def test_submatrices(self): + grid = Grid(shape=(11, 11)) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + submatrices = petsc[0].rhs.fielddata.submatrices + + j00 = submatrices.get_submatrix(e, 'J00') + j01 = submatrices.get_submatrix(e, 'J01') + j10 = submatrices.get_submatrix(g, 'J10') + j11 = submatrices.get_submatrix(g, 'J11') + + # Check the number of submatrices + assert len(submatrices.submatrix_keys) == 4 + assert str(submatrices.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" + + # Technically a non-coupled problem, so the only non-zero submatrices + # should be the diagonal ones i.e J00 and J11 + assert submatrices.nonzero_submatrix_keys == ['J00', 'J11'] + assert submatrices.get_submatrix(e, 'J01')['matvecs'] is None + assert submatrices.get_submatrix(g, 'J10')['matvecs'] is None + + j00 = submatrices.get_submatrix(e, 'J00') + j11 = submatrices.get_submatrix(g, 'J11') + + assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ + + ' Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2)))' + + assert str(j11['matvecs'][0]) == 'Eq(y_g(x, y),' \ + + ' Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2)))' + + # Check the derivative wrt fields + assert j00['derivative_wrt'] == e + assert j01['derivative_wrt'] == g + assert j10['derivative_wrt'] == e + assert j11['derivative_wrt'] == g + + # TODO: + # @skipif('petsc') + # def test_create_submats(self): + + # add tests for all new callbacks + # def test_create_whole_matmult(): From ccb68db14764408276c1ded4dbfc4fdb9606a9c7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 19 Mar 2025 17:26:17 +0000 Subject: [PATCH 142/151] edit petsc utils --- devito/petsc/utils.py | 41 +++++++++++++++++------------------------ 1 file changed, 17 insertions(+), 24 deletions(-) diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 3dbb2c0d1a..bb1c123160 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -2,6 +2,9 @@ from devito.tools import memoized_func +class PetscOSError(OSError): + pass + solver_mapper = { 'gmres': 'KSPGMRES', @@ -12,42 +15,32 @@ @memoized_func def get_petsc_dir(): - # *** First try: via commonly used environment variables - for i in ['PETSC_DIR']: - petsc_dir = os.environ.get(i) - if petsc_dir: - return petsc_dir - # TODO: Raise error if PETSC_DIR is not set - return None + petsc_dir = os.environ.get('PETSC_DIR') + if petsc_dir is None: + raise PetscOSError("PETSC_DIR environment variable not set") + return petsc_dir @memoized_func def get_petsc_arch(): - # *** First try: via commonly used environment variables - for i in ['PETSC_ARCH']: - petsc_arch = os.environ.get(i) - if petsc_arch: - return petsc_arch - # TODO: Raise error if PETSC_ARCH is not set - return None + # Note: users don't have to set PETSC_ARCH + return os.environ.get('PETSC_ARCH') def core_metadata(): petsc_dir = get_petsc_dir() petsc_arch = get_petsc_arch() - # Include directories - global_include = os.path.join(petsc_dir, 'include') - config_specific_include = os.path.join(petsc_dir, f'{petsc_arch}', 'include') - include_dirs = (global_include, config_specific_include) - - # Lib directories - lib_dir = os.path.join(petsc_dir, f'{petsc_arch}', 'lib') + petsc_include = (os.path.join(petsc_dir, 'include'),) + petsc_lib = (os.path.join(petsc_dir, 'lib'),) + if petsc_arch: + petsc_include += (os.path.join(petsc_dir, petsc_arch, 'include'),) + petsc_lib += (os.path.join(petsc_dir, petsc_arch, 'lib'),) return { 'includes': ('petscsnes.h', 'petscdmda.h'), - 'include_dirs': include_dirs, + 'include_dirs': petsc_include, 'libs': ('petsc'), - 'lib_dirs': lib_dir, - 'ldflags': ('-Wl,-rpath,%s' % lib_dir) + 'lib_dirs': petsc_lib, + 'ldflags': tuple([f"-Wl,-rpath, {lib}" for lib in petsc_lib]) } From 70e0b0f2949a5349d93efe0494ac4f823e421e0d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 19 Mar 2025 17:54:01 +0000 Subject: [PATCH 143/151] compiler: petsc memory allocator --- devito/petsc/allocators.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 devito/petsc/allocators.py diff --git a/devito/petsc/allocators.py b/devito/petsc/allocators.py new file mode 100644 index 0000000000..a081465ae6 --- /dev/null +++ b/devito/petsc/allocators.py @@ -0,0 +1,30 @@ +from devito.data.allocators import MemoryAllocator +from pathlib import Path +from devito.petsc.utils import core_metadata + + +class PetscMemoryAllocator(MemoryAllocator): + """ + """ + @classmethod + def initialize(cls): + 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): + c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) + c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) + ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) + + if ret == 0: + return c_pointer, (c_pointer, ) + else: + return None, None + + def free(self, c_pointer): + self.lib.PetscFree(c_pointer) From cddee0d87150e84103e1c5dae49a2c38dd5f7eba Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 19 Mar 2025 18:02:39 +0000 Subject: [PATCH 144/151] edit utils --- devito/petsc/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index bb1c123160..98a2ed5284 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -27,6 +27,7 @@ def get_petsc_arch(): return os.environ.get('PETSC_ARCH') +@memoized_func def core_metadata(): petsc_dir = get_petsc_dir() petsc_arch = get_petsc_arch() From 8a273f46b0915c7c39056c970295554f57885e27 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 19 Mar 2025 18:13:04 +0000 Subject: [PATCH 145/151] update utils --- devito/petsc/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 98a2ed5284..c1f33b31b0 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -43,5 +43,5 @@ def core_metadata(): 'include_dirs': petsc_include, 'libs': ('petsc'), 'lib_dirs': petsc_lib, - 'ldflags': tuple([f"-Wl,-rpath, {lib}" for lib in petsc_lib]) + 'ldflags': tuple([f"-Wl,-rpath,{lib}" for lib in petsc_lib]) } From c29ac830db7d94b16a39834342421de2c52b3baf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 19 Mar 2025 18:31:24 +0000 Subject: [PATCH 146/151] petsc malloc --- devito/petsc/allocators.py | 2 ++ examples/petsc/memory_allocator.py | 33 ++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) create mode 100644 examples/petsc/memory_allocator.py diff --git a/devito/petsc/allocators.py b/devito/petsc/allocators.py index a081465ae6..793918af53 100644 --- a/devito/petsc/allocators.py +++ b/devito/petsc/allocators.py @@ -28,3 +28,5 @@ def _alloc_C_libcall(self, size, ctype): def free(self, c_pointer): self.lib.PetscFree(c_pointer) + +PETSC_ALLOC = PetscMemoryAllocator() diff --git a/examples/petsc/memory_allocator.py b/examples/petsc/memory_allocator.py new file mode 100644 index 0000000000..9435f13f67 --- /dev/null +++ b/examples/petsc/memory_allocator.py @@ -0,0 +1,33 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, configuration) +from devito.petsc import PETScSolve +from devito.petsc.initialize import PetscInitialize +from devito.petsc.allocators import PetscMemoryAllocator +import devito.data.allocators +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +devito.data.allocators.ALLOC_NUMA_ANY = PetscMemoryAllocator() + +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) +u.data._allocator = PetscMemoryAllocator() +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) + +op = Operator(petsc) + +op.apply() \ No newline at end of file From 08417551b82b4c0c64c392096ea00b401e220939 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 19 Mar 2025 19:57:50 +0000 Subject: [PATCH 147/151] fix mem issue --- devito/data/allocators.py | 30 +++++++++++++++++++- devito/petsc/allocators.py | 44 +++++++++++++++--------------- devito/petsc/iet/passes.py | 16 +++++------ devito/petsc/iet/routines.py | 35 ++++++++++++++---------- devito/petsc/types/object.py | 20 ++------------ examples/petsc/memory_allocator.py | 44 ++++++++++++++++++++---------- tests/test_petsc.py | 8 ++++-- 7 files changed, 117 insertions(+), 80 deletions(-) diff --git a/devito/data/allocators.py b/devito/data/allocators.py index aff28ef108..770d86bdaa 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -15,7 +15,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: @@ -141,6 +141,33 @@ def _alloc_C_libcall(self, size, ctype): return +class PetscMemoryAllocator(MemoryAllocator): + """ + """ + @classmethod + def initialize(cls): + 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): + c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) + c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) + ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) + + if ret == 0: + return c_pointer, (c_pointer, ) + else: + return None, None + + def free(self, c_pointer): + self.lib.PetscFree(c_pointer) + + class PosixAllocator(MemoryAllocator): """ @@ -396,6 +423,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/allocators.py b/devito/petsc/allocators.py index 793918af53..0285c5afd2 100644 --- a/devito/petsc/allocators.py +++ b/devito/petsc/allocators.py @@ -3,30 +3,30 @@ from devito.petsc.utils import core_metadata -class PetscMemoryAllocator(MemoryAllocator): - """ - """ - @classmethod - def initialize(cls): - metadata = core_metadata() - lib_dir = Path(metadata['lib_dirs'][-1]) +# class PetscMemoryAllocator(MemoryAllocator): +# """ +# """ +# @classmethod +# def initialize(cls): +# metadata = core_metadata() +# lib_dir = Path(metadata['lib_dirs'][-1]) - try: - cls.lib = ctypes.CDLL(lib_dir/'libpetsc.so') - except OSError: - cls.lib = None +# try: +# cls.lib = ctypes.CDLL(lib_dir/'libpetsc.so') +# except OSError: +# cls.lib = None - def _alloc_C_libcall(self, size, ctype): - c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) - c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) - ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) +# def _alloc_C_libcall(self, size, ctype): +# c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) +# c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) +# ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) - if ret == 0: - return c_pointer, (c_pointer, ) - else: - return None, None +# if ret == 0: +# return c_pointer, (c_pointer, ) +# else: +# return None, None - def free(self, c_pointer): - self.lib.PetscFree(c_pointer) +# def free(self, c_pointer): +# self.lib.PetscFree(c_pointer) -PETSC_ALLOC = PetscMemoryAllocator() +# PETSC_ALLOC = PetscMemoryAllocator() diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 710f25a611..f191ad2052 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -11,7 +11,7 @@ from devito.types.basic import DataSymbol from devito.tools import frozendict from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, - PointerIS, Mat, LocalVec, GlobalVec, CallbackMat, SNES, + PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES, DummyArg, PetscInt, PointerDM, PointerMat, MatReuse, CallbackPointerIS, CallbackPointerDM, JacobianStruct, SubMatrixStruct, Initialize, Finalize, ArgvSymbol) @@ -212,13 +212,13 @@ def populate_matrix_context(efuncs, objs): 'rowidx': PetscInt('rowidx'), 'colidx': PetscInt('colidx'), 'J': Mat('J'), - 'X': GlobalVec('X'), - 'xloc': LocalVec('xloc'), - 'Y': GlobalVec('Y'), - 'yloc': LocalVec('yloc'), - 'F': GlobalVec('F'), - 'floc': LocalVec('floc'), - 'B': GlobalVec('B'), + 'X': Vec('X'), + 'xloc': CallbackVec('xloc'), + 'Y': Vec('Y'), + 'yloc': CallbackVec('yloc'), + 'F': Vec('F'), + 'floc': CallbackVec('floc'), + 'B': Vec('B'), 'nfields': PetscInt('nfields'), 'irow': PointerIS(name='irow'), 'icol': PointerIS(name='icol'), diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index a70e987e71..bf127214a1 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -16,7 +16,7 @@ MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.petsc.utils import solver_mapper -from devito.petsc.types import (DM, Mat, LocalVec, GlobalVec, KSP, PC, SNES, +from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStructCast, JacobianStruct, SubMatrixStruct, CallbackDM) @@ -448,8 +448,13 @@ def _create_form_rhs_body(self, body, fielddata): 'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] ) + dm_restore_local_bvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(sobjs['blocal'])] + ) + body = body._rebuild(body=body.body + ( - dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array + dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array, + dm_restore_local_bvec )) stacks = ( @@ -868,10 +873,10 @@ def _build(self): targets = self.fielddata.targets base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), - 'xglobal': GlobalVec(sreg.make_name(prefix='xglobal')), - 'xlocal': LocalVec(sreg.make_name(prefix='xlocal')), - 'bglobal': GlobalVec(sreg.make_name(prefix='bglobal')), - 'blocal': LocalVec(sreg.make_name(prefix='blocal')), + 'xglobal': Vec(sreg.make_name(prefix='xglobal')), + 'xlocal': CallbackVec(sreg.make_name(prefix='xlocal')), + 'bglobal': Vec(sreg.make_name(prefix='bglobal')), + 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), 'ksp': KSP(sreg.make_name(prefix='ksp')), 'pc': PC(sreg.make_name(prefix='pc')), 'snes': SNES(sreg.make_name(prefix='snes')), @@ -937,9 +942,9 @@ def _extend_build(self, base_dict): name=f'{key}ctx', fields=objs['subctx'].fields, ) - base_dict[f'{key}X'] = LocalVec(f'{key}X') - base_dict[f'{key}Y'] = LocalVec(f'{key}Y') - base_dict[f'{key}F'] = LocalVec(f'{key}F') + base_dict[f'{key}X'] = CallbackVec(f'{key}X') + base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') + base_dict[f'{key}F'] = CallbackVec(f'{key}F') return base_dict @@ -951,22 +956,22 @@ def _target_dependent(self, base_dict): base_dict[f'{name}_ptr'] = StartPtr( sreg.make_name(prefix=f'{name}_ptr'), t.dtype ) - base_dict[f'xlocal{name}'] = LocalVec( + base_dict[f'xlocal{name}'] = CallbackVec( sreg.make_name(prefix=f'xlocal{name}'), liveness='eager' ) - base_dict[f'Fglobal{name}'] = LocalVec( + base_dict[f'Fglobal{name}'] = CallbackVec( sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager' ) - base_dict[f'Xglobal{name}'] = LocalVec( + base_dict[f'Xglobal{name}'] = CallbackVec( sreg.make_name(prefix=f'Xglobal{name}') ) - base_dict[f'xglobal{name}'] = GlobalVec( + base_dict[f'xglobal{name}'] = Vec( sreg.make_name(prefix=f'xglobal{name}') ) - base_dict[f'blocal{name}'] = LocalVec( + base_dict[f'blocal{name}'] = CallbackVec( sreg.make_name(prefix=f'blocal{name}'), liveness='eager' ) - base_dict[f'bglobal{name}'] = GlobalVec( + base_dict[f'bglobal{name}'] = Vec( sreg.make_name(prefix=f'bglobal{name}') ) base_dict[f'da{name}'] = DM( diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 9acf7def46..7e5110ef34 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -67,28 +67,14 @@ def _C_free_priority(self): return 2 -class LocalVec(LocalObject): +class CallbackVec(LocalObject): """ - PETSc local vector object (Vec). - A local vector has ghost locations that contain values that are - owned by other MPI ranks. + PETSc vector object (Vec). """ dtype = CustomDtype('Vec') -class CallbackGlobalVec(LocalVec): - """ - PETSc global vector object (Vec). For example, used for coupled - solves inside the `WholeFormFunc` callback. - """ - - -class GlobalVec(LocalVec): - """ - PETSc global vector object (Vec). - A global vector is a parallel vector that has no duplicate values - between MPI ranks. A global vector has no ghost locations. - """ +class Vec(CallbackVec): @property def _C_free(self): return petsc_call('VecDestroy', [Byref(self.function)]) diff --git a/examples/petsc/memory_allocator.py b/examples/petsc/memory_allocator.py index 9435f13f67..996a76c710 100644 --- a/examples/petsc/memory_allocator.py +++ b/examples/petsc/memory_allocator.py @@ -1,33 +1,47 @@ import os import numpy as np -from devito import (Grid, Function, Eq, Operator, configuration) +from devito import (Grid, Function, Eq, Operator, configuration, TimeFunction, switchconfig) from devito.petsc import PETScSolve from devito.petsc.initialize import PetscInitialize -from devito.petsc.allocators import PetscMemoryAllocator -import devito.data.allocators +# from devito.petsc.allocators import PetscMemoryAllocator +# import devito.data.allocators configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' -devito.data.allocators.ALLOC_NUMA_ANY = PetscMemoryAllocator() +# devito.data.allocators.ALLOC_NUMA_ANY = PetscMemoryAllocator() PetscInitialize() -nx = 81 -ny = 81 +# nx = 81 +# ny = 81 -grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) +# grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) -u = Function(name='u', grid=grid, dtype=np.float64, space_order=2) -u.data._allocator = PetscMemoryAllocator() -v = Function(name='v', grid=grid, dtype=np.float64, space_order=2) +# u = Function(name='u', grid=grid, dtype=np.float64, space_order=2) +# # u.data._allocator = PetscMemoryAllocator() +# v = Function(name='v', grid=grid, dtype=np.float64, space_order=2) -v.data[:] = 5.0 +# v.data[:] = 5.0 -eq = Eq(v, u.laplace, subdomain=grid.interior) +# eq = Eq(v, u.laplace, subdomain=grid.interior) -petsc = PETScSolve([eq], u) +# petsc = PETScSolve([eq], u) -op = Operator(petsc) +# op = Operator(petsc) +# # print(op.ccode) +# op.apply() -op.apply() \ No newline at end of file + + +grid = Grid((11, 11)) + +# Modulo time stepping +u1 = TimeFunction(name='u1', grid=grid, space_order=2) +v1 = Function(name='v1', grid=grid, space_order=2) +eq1 = Eq(v1.laplace, u1) +petsc1 = PETScSolve(eq1, v1) +with switchconfig(openmp=False): + op1 = Operator(petsc1) + print(op1.ccode) +op1.apply(time_M=5) \ No newline at end of file diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 1f0f8bc91e..5f3f177122 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -9,7 +9,7 @@ FindNodes, retrieve_iteration_tree) from devito.types import Constant, LocalCompositeObject from devito.passes.iet.languages.C import CDataManager -from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, +from devito.petsc.types import (DM, Mat, CallbackVec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, LinearSolveExpr, FieldData, MultipleFieldData) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil @@ -33,7 +33,7 @@ def test_petsc_local_object(): """ lo0 = DM('da', stencil_width=1) lo1 = Mat('A') - lo2 = LocalVec('x') + lo2 = CallbackVec('x') lo3 = PetscMPIInt('size') lo4 = KSP('ksp') lo5 = PC('pc') @@ -735,6 +735,7 @@ def test_time_loop(): petsc1 = PETScSolve(eq1, v1) with switchconfig(openmp=False): op1 = Operator(petsc1) + op1.apply(time_M=5) body1 = str(op1.body) rhs1 = str(op1._func_table['FormRHS0'].root.ccode) @@ -750,6 +751,7 @@ def test_time_loop(): petsc2 = PETScSolve(eq2, v2) with switchconfig(openmp=False): op2 = Operator(petsc2) + # op2.apply(time_M=5) body2 = str(op2.body) rhs2 = str(op2._func_table['FormRHS0'].root.ccode) @@ -762,6 +764,7 @@ def test_time_loop(): petsc3 = PETScSolve(eq3, v1) with switchconfig(openmp=False): op3 = Operator(petsc3) + # op3.apply(time_M=5) body3 = str(op3.body) rhs3 = str(op3._func_table['FormRHS0'].root.ccode) @@ -778,6 +781,7 @@ def test_time_loop(): petsc5 = PETScSolve(eq5, v2) with switchconfig(openmp=False): op4 = Operator(petsc4 + petsc5) + # op4.apply(time_M=5) body4 = str(op4.body) assert 'ctx0.t0 = t0' in body4 From cdee484b476090c0ce5596e4646120ec0660fde6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 20 Mar 2025 10:44:25 +0000 Subject: [PATCH 148/151] edit memory allocator --- devito/data/allocators.py | 54 +++++++++++++++--------------- devito/petsc/allocators.py | 32 ------------------ devito/petsc/iet/routines.py | 2 +- examples/petsc/memory_allocator.py | 43 +++++++----------------- tests/test_petsc.py | 20 +++++------ 5 files changed, 51 insertions(+), 100 deletions(-) delete mode 100644 devito/petsc/allocators.py diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 770d86bdaa..c8660ce86e 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -141,33 +141,6 @@ def _alloc_C_libcall(self, size, ctype): return -class PetscMemoryAllocator(MemoryAllocator): - """ - """ - @classmethod - def initialize(cls): - 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): - c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) - c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) - ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) - - if ret == 0: - return c_pointer, (c_pointer, ) - else: - return None, None - - def free(self, c_pointer): - self.lib.PetscFree(c_pointer) - - class PosixAllocator(MemoryAllocator): """ @@ -364,6 +337,33 @@ def put_local(self): return self._node == 'local' +class PetscMemoryAllocator(MemoryAllocator): + """ + """ + @classmethod + def initialize(cls): + 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): + c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) + c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) + ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) + + if ret == 0: + return c_pointer, (c_pointer, ) + else: + return None, None + + def free(self, c_pointer): + self.lib.PetscFree(c_pointer) + + class DataReference(MemoryAllocator): """ diff --git a/devito/petsc/allocators.py b/devito/petsc/allocators.py deleted file mode 100644 index 0285c5afd2..0000000000 --- a/devito/petsc/allocators.py +++ /dev/null @@ -1,32 +0,0 @@ -from devito.data.allocators import MemoryAllocator -from pathlib import Path -from devito.petsc.utils import core_metadata - - -# class PetscMemoryAllocator(MemoryAllocator): -# """ -# """ -# @classmethod -# def initialize(cls): -# 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): -# c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) -# c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) -# ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) - -# if ret == 0: -# return c_pointer, (c_pointer, ) -# else: -# return None, None - -# def free(self, c_pointer): -# self.lib.PetscFree(c_pointer) - -# PETSC_ALLOC = PetscMemoryAllocator() diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index bf127214a1..e7972cf2ef 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -874,7 +874,7 @@ def _build(self): base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), 'xglobal': Vec(sreg.make_name(prefix='xglobal')), - 'xlocal': CallbackVec(sreg.make_name(prefix='xlocal')), + 'xlocal': Vec(sreg.make_name(prefix='xlocal')), 'bglobal': Vec(sreg.make_name(prefix='bglobal')), 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), 'ksp': KSP(sreg.make_name(prefix='ksp')), diff --git a/examples/petsc/memory_allocator.py b/examples/petsc/memory_allocator.py index 996a76c710..1490d8dd2e 100644 --- a/examples/petsc/memory_allocator.py +++ b/examples/petsc/memory_allocator.py @@ -4,44 +4,27 @@ from devito import (Grid, Function, Eq, Operator, configuration, TimeFunction, switchconfig) from devito.petsc import PETScSolve from devito.petsc.initialize import PetscInitialize -# from devito.petsc.allocators import PetscMemoryAllocator -# import devito.data.allocators +from devito.data.allocators import ALLOC_PETSC +import devito.data.allocators configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' -# devito.data.allocators.ALLOC_NUMA_ANY = PetscMemoryAllocator() - 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) -# # u.data._allocator = PetscMemoryAllocator() -# v = Function(name='v', grid=grid, dtype=np.float64, space_order=2) - -# v.data[:] = 5.0 - -# eq = Eq(v, u.laplace, subdomain=grid.interior) +nx = 81 +ny = 81 -# petsc = PETScSolve([eq], u) +grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) -# op = Operator(petsc) -# # print(op.ccode) -# op.apply() +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) -grid = Grid((11, 11)) +petsc = PETScSolve([eq], u) -# Modulo time stepping -u1 = TimeFunction(name='u1', grid=grid, space_order=2) -v1 = Function(name='v1', grid=grid, space_order=2) -eq1 = Eq(v1.laplace, u1) -petsc1 = PETScSolve(eq1, v1) -with switchconfig(openmp=False): - op1 = Operator(petsc1) - print(op1.ccode) -op1.apply(time_M=5) \ No newline at end of file +op = Operator(petsc) +print(op.ccode) +op.apply() diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 5f3f177122..93d51d8546 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -726,16 +726,16 @@ def test_time_loop(): - Only assign/update the modulo dimensions required by any of the PETSc callback functions. """ - grid = Grid((11, 11)) + grid = Grid((11, 11), dtype=np.float64) # Modulo time stepping - u1 = TimeFunction(name='u1', grid=grid, space_order=2) - v1 = Function(name='v1', grid=grid, space_order=2) + u1 = TimeFunction(name='u1', grid=grid, space_order=2, dtype=np.float64) + v1 = Function(name='v1', grid=grid, space_order=2, dtype=np.float64) eq1 = Eq(v1.laplace, u1) petsc1 = PETScSolve(eq1, v1) with switchconfig(openmp=False): op1 = Operator(petsc1) - op1.apply(time_M=5) + op1.apply(time_M=3) body1 = str(op1.body) rhs1 = str(op1._func_table['FormRHS0'].root.ccode) @@ -745,13 +745,13 @@ def test_time_loop(): assert 'ctx0->t1' not in rhs1 # Non-modulo time stepping - u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) - v2 = Function(name='v2', grid=grid, space_order=2, save=5) + u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5, dtype=np.float64) + v2 = Function(name='v2', grid=grid, space_order=2, save=5, dtype=np.float64) eq2 = Eq(v2.laplace, u2) petsc2 = PETScSolve(eq2, v2) with switchconfig(openmp=False): op2 = Operator(petsc2) - # op2.apply(time_M=5) + op2.apply(time_M=3) body2 = str(op2.body) rhs2 = str(op2._func_table['FormRHS0'].root.ccode) @@ -764,7 +764,7 @@ def test_time_loop(): petsc3 = PETScSolve(eq3, v1) with switchconfig(openmp=False): op3 = Operator(petsc3) - # op3.apply(time_M=5) + op3.apply(time_M=3) body3 = str(op3.body) rhs3 = str(op3._func_table['FormRHS0'].root.ccode) @@ -774,14 +774,14 @@ def test_time_loop(): assert 'ctx0->t1' in rhs3 # Multiple petsc solves within the same time loop - v2 = Function(name='v2', grid=grid, space_order=2) + v2 = Function(name='v2', grid=grid, space_order=2, dtype=np.float64) eq4 = Eq(v1.laplace, u1) petsc4 = PETScSolve(eq4, v1) eq5 = Eq(v2.laplace, u1) petsc5 = PETScSolve(eq5, v2) with switchconfig(openmp=False): op4 = Operator(petsc4 + petsc5) - # op4.apply(time_M=5) + op4.apply(time_M=3) body4 = str(op4.body) assert 'ctx0.t0 = t0' in body4 From b512417ecab216fda607468669ea4cabc7dca6b3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 20 Mar 2025 11:14:21 +0000 Subject: [PATCH 149/151] fix --- conftest.py | 5 ++--- devito/petsc/iet/routines.py | 3 ++- devito/petsc/utils.py | 3 ++- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/conftest.py b/conftest.py index b804aaebc7..92bc8148b4 100644 --- a/conftest.py +++ b/conftest.py @@ -14,7 +14,6 @@ from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock, retrieve_iteration_tree) from devito.tools import as_tuple -from devito.petsc.utils import get_petsc_dir, get_petsc_arch try: from mpi4py import MPI # noqa @@ -89,8 +88,8 @@ def skipif(items, whole_module=False): skipit = "pyrevolve not installed" break if i == 'petsc': - petsc_dir = get_petsc_dir() - petsc_arch = get_petsc_arch() + petsc_dir = os.environ.get('PETSC_DIR') + petsc_arch = os.environ.get('PETSC_ARCH') if petsc_dir is None or petsc_arch is None: skipit = "PETSC_DIR or PETSC_ARCH are not set" break diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index e7972cf2ef..9f25748405 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -874,7 +874,8 @@ def _build(self): base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), 'xglobal': Vec(sreg.make_name(prefix='xglobal')), - 'xlocal': Vec(sreg.make_name(prefix='xlocal')), + # TODO: I think 'xloc' should technically be a Vec because it should be destroyed? + 'xlocal': CallbackVec(sreg.make_name(prefix='xlocal')), 'bglobal': Vec(sreg.make_name(prefix='bglobal')), 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), 'ksp': KSP(sreg.make_name(prefix='ksp')), diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index c1f33b31b0..94dd858684 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -23,7 +23,8 @@ def get_petsc_dir(): @memoized_func def get_petsc_arch(): - # Note: users don't have to set PETSC_ARCH + # Note: users don't have to explicitly set PETSC_ARCH + # if they add it to the PETSC_DIR path return os.environ.get('PETSC_ARCH') From 8046b75e1be6e715c87b8f66e8bcc4120bf31945 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 20 Mar 2025 11:25:50 +0000 Subject: [PATCH 150/151] ferfrwe --- devito/data/allocators.py | 3 ++- devito/petsc/iet/routines.py | 3 ++- devito/petsc/utils.py | 1 + examples/petsc/memory_allocator.py | 3 +-- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/devito/data/allocators.py b/devito/data/allocators.py index c8660ce86e..664338b2a2 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -6,12 +6,14 @@ import mmap import os import sys +from pathlib import Path import numpy as np from devito.logger import logger from devito.parameters import configuration from devito.tools import dtype_to_ctype, is_integer +from devito.petsc.utils import core_metadata __all__ = ['ALLOC_ALIGNED', 'ALLOC_NUMA_LOCAL', 'ALLOC_NUMA_ANY', 'ALLOC_KNL_MCDRAM', 'ALLOC_KNL_DRAM', 'ALLOC_GUARD', @@ -351,7 +353,6 @@ def initialize(cls): cls.lib = None def _alloc_C_libcall(self, size, ctype): - c_bytesize = ctypes.c_ulong(size * ctypes.sizeof(ctype)) c_pointer = ctypes.cast(ctypes.c_void_p(), ctypes.c_void_p) ret = self.lib.PetscMalloc(size, ctypes.byref(c_pointer)) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 9f25748405..4baad8a35b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -874,7 +874,8 @@ def _build(self): base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), 'xglobal': Vec(sreg.make_name(prefix='xglobal')), - # TODO: I think 'xloc' should technically be a Vec because it should be destroyed? + # TODO: I think 'xloc' should technically be a Vec because + # it should be destroyed? 'xlocal': CallbackVec(sreg.make_name(prefix='xlocal')), 'bglobal': Vec(sreg.make_name(prefix='bglobal')), 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 94dd858684..47ab4eba3f 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -2,6 +2,7 @@ from devito.tools import memoized_func + class PetscOSError(OSError): pass diff --git a/examples/petsc/memory_allocator.py b/examples/petsc/memory_allocator.py index 1490d8dd2e..90a6ac061f 100644 --- a/examples/petsc/memory_allocator.py +++ b/examples/petsc/memory_allocator.py @@ -1,11 +1,10 @@ import os import numpy as np -from devito import (Grid, Function, Eq, Operator, configuration, TimeFunction, switchconfig) +from devito import Grid, Function, Eq, Operator, configuration from devito.petsc import PETScSolve from devito.petsc.initialize import PetscInitialize from devito.data.allocators import ALLOC_PETSC -import devito.data.allocators configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' From a927ca61f2a03f2ed9c7ae1f2125f061b8b11f28 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 20 Mar 2025 11:28:09 +0000 Subject: [PATCH 151/151] circular import --- devito/data/allocators.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 664338b2a2..cdba6bd0f8 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -13,7 +13,7 @@ from devito.logger import logger from devito.parameters import configuration from devito.tools import dtype_to_ctype, is_integer -from devito.petsc.utils import core_metadata + __all__ = ['ALLOC_ALIGNED', 'ALLOC_NUMA_LOCAL', 'ALLOC_NUMA_ANY', 'ALLOC_KNL_MCDRAM', 'ALLOC_KNL_DRAM', 'ALLOC_GUARD', @@ -344,6 +344,8 @@ class PetscMemoryAllocator(MemoryAllocator): """ @classmethod def initialize(cls): + # from IPython import embed; embed() + from devito.petsc.utils import core_metadata metadata = core_metadata() lib_dir = Path(metadata['lib_dirs'][-1])