From 0b92950dc2744ec7b5dffbb6b8abe6539b3c0874 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 10:46:24 +0000 Subject: [PATCH 001/107] 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 9bcc3460f6..339a2e70ba 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 505fe2e001..a9428f08f9 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -509,6 +509,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) @@ -1416,3 +1422,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 ceba1b91190172725dea210212cff3fadd248db4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 11:46:01 +0000 Subject: [PATCH 002/107] 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 339a2e70ba..197889ed03 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 1c8ef9493a7ff1c99721c9b8b871f0e9f2bf043e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 11:04:27 +0000 Subject: [PATCH 003/107] 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 197889ed03..752de95e5b 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): """ @@ -1133,6 +1120,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 a9428f08f9..32c1964fc5 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -509,12 +509,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) @@ -619,6 +613,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) @@ -1424,12 +1421,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 1d6c21e958bd9cf9577a98fcc30f9996f86c7ab4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 5 Jan 2024 11:20:48 +0000 Subject: [PATCH 004/107] 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 752de95e5b..1c6c85a923 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1133,7 +1133,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 1b50fc729b119a39d05617da2e99cd7d5e24d9d9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sat, 6 Jan 2024 10:37:33 +0000 Subject: [PATCH 005/107] 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 1c6c85a923..fcccbeae88 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1121,7 +1121,6 @@ def defines(self): class Callback(Call): - """ Callback as a function pointer. @@ -1143,9 +1142,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 dd3b021a811c480695330de11111807d9ccbb21c Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 006/107] 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 2e2888071f..78d95d29e6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -45,6 +45,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 @@ -58,6 +59,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo \ # 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 @@ -82,7 +93,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 @@ -99,6 +114,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 @@ -113,6 +157,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 8302ea149121f4f17d7f24abdac21b2d004a571c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 27 Nov 2023 19:16:56 +0000 Subject: [PATCH 007/107] 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 e7e5c5fc061e74d32fb482395e9bf65aae224145 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 30 Nov 2023 23:58:44 +0000 Subject: [PATCH 008/107] 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 5fa3ba7dd55be26f5e7afa615c29af387886e8e0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 4 Dec 2023 14:44:33 +0000 Subject: [PATCH 009/107] 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 09dbaecc5e88ded535c4c869b4a92e442972d5ca Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 4 Dec 2023 17:47:08 +0000 Subject: [PATCH 010/107] 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 b2202c56e85be8c57c8c1e8b7df8f9354d20b003 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 7 Dec 2023 10:58:32 +0000 Subject: [PATCH 011/107] 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 2063c1d9136acbcab459ef719d44683ce184c030 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 13:37:34 +0000 Subject: [PATCH 012/107] 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 a5dbdc71f9258224f19f165177a7ab08f34f8a28 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 15:02:49 +0000 Subject: [PATCH 013/107] 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 3592bf734c83b3e1852e74fa258d8c4cd394148e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 10:18:11 +0000 Subject: [PATCH 014/107] 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 8b80d5b995d2bd1763810d2bfde2ae7f8070f811 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:03:02 +0000 Subject: [PATCH 015/107] 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 28981d3bf1..e9161e7cc7 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -222,6 +222,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: @@ -359,6 +367,8 @@ def place_definitions(self, iet, globs=None, **kwargs): 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) + 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 3d3241f27d..6bf0d21e7f 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -257,6 +257,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 d557f9661e25e7c7a57ee6959e5f5d9fdf3f9d60 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:08:40 +0000 Subject: [PATCH 016/107] 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 ffc05039b5987aa6a2f96588c29ab01fbc343af4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:20:16 +0000 Subject: [PATCH 017/107] 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 5e292adfa07eeb5908bb1cae73a05f04522e41b8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 20 Dec 2023 21:22:33 +0000 Subject: [PATCH 018/107] 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 71d109ca1975e908b9fcf8d3ee1cefd711c5bca9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 019/107] 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 78d95d29e6..2e2888071f 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -45,7 +45,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 @@ -59,16 +58,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo \ # 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 @@ -93,11 +82,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 @@ -114,35 +99,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 @@ -157,12 +113,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 17e8570471b9bad77e839277d8d98ab882b6e846 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Jan 2024 11:46:05 +0000 Subject: [PATCH 020/107] compiler: Change name to PETScArray from PETScFunction --- devito/passes/iet/definitions.py | 12 +----------- devito/types/petsc.py | 2 +- tests/test_petsc.py | 20 ++++++++++---------- 3 files changed, 12 insertions(+), 22 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index e9161e7cc7..2387225121 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -222,14 +222,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: @@ -363,12 +355,10 @@ def place_definitions(self, iet, globs=None, **kwargs): elif globs is not None: # Track, to be handled by the EntryFunction being a global obj! globs.add(i) - elif i.is_ObjectArray: + elif i.is_ObjectArray or i.is_PETScArray: 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) - 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 e5df4344eb3fa2d050d69d7df68bec9181b54798 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 11:12:00 +0000 Subject: [PATCH 021/107] compiler: Add intermediate class between PETScArray and ArrayBasic --- devito/passes/iet/definitions.py | 2 +- devito/types/basic.py | 1 - devito/types/petsc.py | 28 ++++++++++++++++++++-------- 3 files changed, 21 insertions(+), 10 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 2387225121..5c9fec4ed8 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -355,7 +355,7 @@ def place_definitions(self, iet, globs=None, **kwargs): elif globs is not None: # Track, to be handled by the EntryFunction being a global obj! globs.add(i) - elif i.is_ObjectArray or i.is_PETScArray: + elif i.is_ObjectArray or i.is_AbstractArray: 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) diff --git a/devito/types/basic.py b/devito/types/basic.py index 6bf0d21e7f..3d3241f27d 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -257,7 +257,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 c56405fb845cde7f6d487c48464900096aad6619 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 13:33:58 +0000 Subject: [PATCH 022/107] compiler: Remove PETSc specific checks in place_definitions --- devito/passes/iet/definitions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 5c9fec4ed8..204426d3bf 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -355,10 +355,12 @@ def place_definitions(self, iet, globs=None, **kwargs): elif globs is not None: # Track, to be handled by the EntryFunction being a global obj! globs.add(i) - elif i.is_ObjectArray or i.is_AbstractArray: + 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 ee12a3375cc0e254d0343668714a4872e00f1ffc Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 14:08:14 +0000 Subject: [PATCH 023/107] 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 204426d3bf..ebcb50b04e 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -222,6 +222,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: @@ -360,7 +371,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 651885a8c361c0cb5c4cc59bbcaa058d263b689f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 16:14:31 +0000 Subject: [PATCH 024/107] 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 43bbdca5197484237c22233f2c34713d4a417bad Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 4 Jan 2024 18:53:04 +0000 Subject: [PATCH 025/107] compiler: Remove elif is_ObjectArray and use alloc_object_.. inside the else statement --- devito/passes/iet/definitions.py | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index ebcb50b04e..ee55c90a45 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -222,17 +222,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: @@ -366,12 +355,10 @@ def place_definitions(self, iet, globs=None, **kwargs): elif globs is not None: # Track, to be handled by the EntryFunction being a global obj! globs.add(i) - elif i.is_ObjectArray: - self._alloc_object_array_on_low_lat_mem(iet, i, storage) elif i.is_PointerArray: self._alloc_pointed_array_on_high_bw_mem(iet, i, storage) else: - self._alloc_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 4dcc4b926c5a02f606069b6b71f0f61c13eb7572 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 8 Jan 2024 10:04:22 +0000 Subject: [PATCH 026/107] 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 f1b31a23b7a9714ed02552a5ada060bfcf5b3f79 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 8 Jan 2024 10:12:19 +0000 Subject: [PATCH 027/107] 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 9aec979a6dda9b78fb72c94fe068ed08763f52f5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 8 Jan 2024 10:14:28 +0000 Subject: [PATCH 028/107] 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 866ac5f9ce41205b99de8cdc4fe85f32a5ccf700 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/107] 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 fcccbeae88..203aa0ff9e 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 16028adf409379262dc0d37542867c148a568525 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/107] 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 203aa0ff9e..e908e65275 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 34ea7509676421477c7bc55a6edaa53cff2f5528 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 10:46:24 +0000 Subject: [PATCH 031/107] 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 e908e65275..3a93d83e48 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 32c1964fc5..347505da9b 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -509,6 +509,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 dd2a906d0e08db8f53cd8cb1710cf7182ca04af3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 19 Dec 2023 11:04:27 +0000 Subject: [PATCH 032/107] 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 3a93d83e48..e908e65275 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 347505da9b..32c1964fc5 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -509,12 +509,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 54c21de013dcecb70684d207ad200047fced1cc8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 28 Feb 2024 13:03:28 +0000 Subject: [PATCH 033/107] 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 978cd389fdd353ef5700e6514cd2a6af583fdb81 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 1 Mar 2024 15:45:07 +0000 Subject: [PATCH 034/107] 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 77a0d0b5adf14160fed885da33dbdb38a456689f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 11 Mar 2024 17:25:23 +0000 Subject: [PATCH 035/107] 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 34237f01bf..da2ed44921 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 fe08af8116b336849c7711be781b3a8f7e0fac2d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 11 Mar 2024 17:26:57 +0000 Subject: [PATCH 036/107] 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 da2ed44921..6900496f35 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 4cae443b3d09d2cba3fb768c3f42f52472aa3239 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 11 Mar 2024 19:54:12 +0000 Subject: [PATCH 037/107] 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 aa0a27e8ae3914e0b9d2dbd4059ee8f2c6add656 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 12 Mar 2024 11:59:50 +0000 Subject: [PATCH 038/107] 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 850d5c0ce07811572f304c5cf4d8fede2e4a4140 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 14 Mar 2024 17:25:22 +0000 Subject: [PATCH 039/107] 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 d74b316a21d5386018feef770e116d58a865f32a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 14 Mar 2024 17:51:22 +0000 Subject: [PATCH 040/107] 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 504ee2eac15bba75789d4992016474cff4f34a2e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 15 Mar 2024 11:36:07 +0000 Subject: [PATCH 041/107] 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 42034814ed0e9305a6d97ee830005f3cfd6e0f1e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 15 Mar 2024 15:44:21 +0000 Subject: [PATCH 042/107] 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 61c0404a3b293485f9fc211abfc34f7f65d0da14 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 15 Mar 2024 17:55:47 +0000 Subject: [PATCH 043/107] 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 cd95a1e9a4c3002517356f561b558ea08f535b1c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 14:49:37 +0000 Subject: [PATCH 044/107] 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 629ebdde4a..828e705472 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -373,11 +373,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 e908e65275..fc4dc019ea 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 24258ad671..8cf287e611 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -24,7 +24,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, @@ -473,6 +473,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 73f714f19d8ea6ab1aad32a4a41d4cb9815e845f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:08:42 +0000 Subject: [PATCH 045/107] 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 0275205fe0fa39f4e87a86248ed95743ef113462 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:24:03 +0000 Subject: [PATCH 046/107] 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 828e705472..cfdd6750bc 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -375,7 +375,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 b341fea0722418db04ba7f508223246a4d7a0921 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:27:50 +0000 Subject: [PATCH 047/107] 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 cfdd6750bc..7c52b9c5cf 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -373,6 +373,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 4c0fb93e79463da10bde24d81926b056234d0db7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 16:30:39 +0000 Subject: [PATCH 048/107] 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 2b8309cef6ee732782d0177a8299c3fe43eff95f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Mar 2024 17:04:41 +0000 Subject: [PATCH 049/107] 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 ca6a96a7bbdf5e236807f35a8baa8f19dd4277ac Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 27 Mar 2024 14:03:24 +0000 Subject: [PATCH 050/107] 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 7c52b9c5cf..3e8f4da86c 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -367,25 +367,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 fc4dc019ea..cf5d5cd081 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 de8c6c4be4ea72eea4aecc9a586fcd4c3829830c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 27 Mar 2024 14:07:15 +0000 Subject: [PATCH 051/107] 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 3e8f4da86c..95997e65b2 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -367,7 +367,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 37bb9003349f9533983b7e4f14de1e1b9f7ce576 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 09:27:24 +0000 Subject: [PATCH 052/107] 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 95997e65b2..3aadc93554 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -367,22 +367,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 fc3730530b16c022447a2984ef5b0ef578620be8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 12:55:01 +0000 Subject: [PATCH 053/107] 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 3aadc93554..78d0bd6f54 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 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 8cf287e611..7d27eea6d1 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -347,7 +347,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 672971cf4f..e33df0fad8 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -211,6 +211,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 b8efafd3e3393bd3cfb92c0fcbe2cb2ae5c4a94f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 12:57:59 +0000 Subject: [PATCH 054/107] 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 7d27eea6d1..8cf287e611 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -347,7 +347,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 4268b47a0953798756e99a0640c7f0b0040e19c9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 13:01:40 +0000 Subject: [PATCH 055/107] 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 cf5d5cd081..500c23c75b 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 c13181dbaee6e2b5dfe8853dfe77684909f3b5ea Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 13:36:15 +0000 Subject: [PATCH 056/107] 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 d875a9d702ee571dba442e67e33709b8ce5c262b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 13:51:06 +0000 Subject: [PATCH 057/107] 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 e33df0fad8..78cac05480 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -211,7 +211,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 206f84ed3b08d233e02a728f999d69aeda4de558 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 14:33:54 +0000 Subject: [PATCH 058/107] 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 ad70691e03d6b60a281a3fa87da117022e2a1c5e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 14:36:40 +0000 Subject: [PATCH 059/107] 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 d3f52dbc8c689040b75d4aa7a34a15ec15d22cee Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 14:51:59 +0000 Subject: [PATCH 060/107] 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 ea36ecff5893fd907fc437ef8e723094cc6c5732 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 28 Mar 2024 16:42:32 +0000 Subject: [PATCH 061/107] 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 500c23c75b..cb1c8abbd9 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 20353336d2ce803d634197ab3a9abae6090c7c6b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Apr 2024 13:44:34 +0100 Subject: [PATCH 062/107] 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 cb1c8abbd9..525340c49a 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 6370848c6152068c17adc8b2547c219aca3aadb1 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Apr 2024 18:42:57 +0100 Subject: [PATCH 063/107] 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 525340c49a..122793b2d7 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 78cac05480..672971cf4f 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -211,9 +211,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 6734a2df9b4f84e3eb9ed7479016106e734be2c3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 2 Apr 2024 18:44:31 +0100 Subject: [PATCH 064/107] 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 acf232762a4ae9a953abec658ca7a0da0b48013a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 17 Apr 2024 21:55:51 +0100 Subject: [PATCH 065/107] compiler: Edit PETScArrays to support devito multi-dim indexing --- devito/ir/clusters/cluster.py | 2 -- devito/passes/iet/definitions.py | 7 +++-- devito/types/basic.py | 5 ++- devito/types/petsc.py | 34 ++++++++++++++++---- tests/test_petsc.py | 54 ++++++++++++++++++++++---------- 5 files changed, 75 insertions(+), 27 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index 78d0bd6f54..629ebdde4a 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -367,8 +367,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 ee55c90a45..40a8cd6b0f 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -17,7 +17,8 @@ from devito.symbolics import (Byref, DefFunction, FieldFromPointer, IndexedPointer, SizeOf, VOID, Keyword, pow_to_mul) from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten -from devito.types import Array, CustomDimension, DeviceMap, DeviceRM, Eq, Symbol +from devito.types import (Array, CustomDimension, DeviceMap, DeviceRM, Eq, Symbol, + PETScArray) __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage'] @@ -399,7 +400,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 3d3241f27d..d4369b8f47 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 a1fa4dd489d794bdcd7dc6e3ddfa1c6bb4923103 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 17 Apr 2024 22:11:49 +0100 Subject: [PATCH 066/107] 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 b10beafcc428952e727117d55386a73cb29d60bc Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 09:52:45 +0100 Subject: [PATCH 067/107] 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 07f8e99d5e1ed94b6a7a692c3f6e69d59ab0798a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 09:57:36 +0100 Subject: [PATCH 068/107] 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 b700c430cea479112c191d4e76e462ca8584f2b2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 13:23:08 +0100 Subject: [PATCH 069/107] compiler: Add PETScIndexedData --- devito/passes/iet/definitions.py | 14 +++++++------- devito/types/petsc.py | 10 ++++++++++ 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 40a8cd6b0f..53dc7fa0cd 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -18,7 +18,8 @@ SizeOf, VOID, Keyword, pow_to_mul) from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten from devito.types import (Array, CustomDimension, DeviceMap, DeviceRM, Eq, Symbol, - PETScArray) + PointerArray, PETScArray, IndexedData) +from devito.types.dense import DiscreteFunction __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage'] @@ -386,8 +387,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 @@ -395,14 +396,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 ddc36c4e31fa35971fe282b0a5ce9fa69b237a30 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 13:27:36 +0100 Subject: [PATCH 070/107] misc: Clean --- devito/passes/iet/definitions.py | 8 +++----- devito/types/petsc.py | 5 ++--- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 53dc7fa0cd..3dec70ad42 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -18,8 +18,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, CustomDimension, DeviceMap, DeviceRM, Eq, Symbol, - PointerArray, PETScArray, IndexedData) -from devito.types.dense import DiscreteFunction + IndexedData) __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage'] @@ -387,7 +386,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` @@ -396,8 +394,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 4a7e43645f2c94ca15d45001d345fb53041c6b63 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Apr 2024 13:42:54 +0100 Subject: [PATCH 071/107] 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 3dec70ad42..2302b96501 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -386,7 +386,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 db50cc58c2dd7288ab1866eabbffe24d2a904506 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 25 Apr 2024 16:28:04 +0100 Subject: [PATCH 072/107] 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 2a4886aa85b09d63edbcf969bcc9409ac311eb06 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 26 Apr 2024 13:37:11 +0100 Subject: [PATCH 073/107] 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 9cc79dfd8707ec76976e982877178e8033221a2b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 1 May 2024 10:53:57 +0100 Subject: [PATCH 074/107] 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 2302b96501..765b8ba219 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -387,6 +387,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 bc082ae73f706e1122e99272c37a288ff871ee15 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 01:12:18 +0100 Subject: [PATCH 075/107] 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 122793b2d7..4a99334c52 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 8cf287e611..a2aff3c24d 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -471,6 +471,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 13e1f294809e0d9d4bac9d15bb214ea2ccc2a779 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 01:18:36 +0100 Subject: [PATCH 076/107] 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 4a99334c52..888a2e1143 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 a2aff3c24d..8cf287e611 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -471,7 +471,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 019d1c6b9340555170c62166bc5702c71775b571 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 10:08:12 +0100 Subject: [PATCH 077/107] 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 8cf287e611..55c5fe9935 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -471,10 +471,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 b7fc3329c2..7be1fdffc5 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 cbea467cbc06453c4e4c71c5df477fac6da906c5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 10:12:05 +0100 Subject: [PATCH 078/107] 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 888a2e1143..4266d8a62f 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 d9c6fcd8d23ec610d67091a0af8f7f2215ac5ddb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 11:14:26 +0100 Subject: [PATCH 079/107] 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 4266d8a62f..122793b2d7 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 55c5fe9935..a489c67655 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -418,6 +418,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 7be1fdffc5..b7fc3329c2 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 0d1325cb0c42f9a2050b0342e393ccb1a927e146 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 11:17:51 +0100 Subject: [PATCH 080/107] 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 a489c67655..55c5fe9935 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -418,7 +418,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 b22cd377c1a4a051ab8b62271ca27ea13f49d02c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 11:30:41 +0100 Subject: [PATCH 081/107] 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 5a869b5d6c233c3793397957b0a9d71b67a063bd Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 12:18:20 +0100 Subject: [PATCH 082/107] 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 cf8e9c1d2359951cd2d5b0ea3f61ae5dcb143db5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 3 Apr 2024 12:20:08 +0100 Subject: [PATCH 083/107] 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 20416b748e79b17b946383f071e18fae9362f3d3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 10:40:15 +0100 Subject: [PATCH 084/107] 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 1a05f0842a..03bf8dfe24 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 59f38b74a5c538ca1816bb40c17c376252d2de08 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 10:43:37 +0100 Subject: [PATCH 085/107] 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 046516882904a238a1793a7d84c196b017f65071 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 11:13:02 +0100 Subject: [PATCH 086/107] 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 c6230ddbefdb67f665a5ec1d246d7afe9d5a777a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 11:48:07 +0100 Subject: [PATCH 087/107] 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 536e1399108050cb4a4a9ef65fe13e457aa24f02 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 12:30:38 +0100 Subject: [PATCH 088/107] 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 03bf8dfe24..ec7c97eed3 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 3a5f0a232610aec666e6b915800e2caaa0b807db Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 12:33:49 +0100 Subject: [PATCH 089/107] 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 4824e40a406b87345ee4fe9e9b549ef68faf0fdf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 14:01:28 +0100 Subject: [PATCH 090/107] 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 c49033d2fb013387f587c8e255f2cad944c090af Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 15:45:44 +0100 Subject: [PATCH 091/107] 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 1f245e775ed979c5f15ca40990e48ac6b97befd1 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 2 May 2024 16:02:19 +0100 Subject: [PATCH 092/107] 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 ec7c97eed3..1a05f0842a 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 55c5fe9935..6671d1011b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -24,7 +24,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, @@ -374,6 +374,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 73d1083a502a3780cfdb09a9bc933811849d467f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 9 Jul 2024 23:02:26 +0100 Subject: [PATCH 093/107] 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 e2d7321391b475ef53e84527130b3697dcb4709b Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Tue, 19 Dec 2023 08:32:06 +0000 Subject: [PATCH 094/107] 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 2e2888071f..78d95d29e6 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -45,6 +45,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 @@ -58,6 +59,16 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo \ # 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 @@ -82,7 +93,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 @@ -99,6 +114,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 @@ -113,6 +157,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 f4d9c6a64e56d3b73177c09b61d66de9d68e4edf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 22 Dec 2023 14:46:10 +0000 Subject: [PATCH 095/107] 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 629ebdde4a..dd513bcaed 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -374,7 +374,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 122793b2d7..7026ad449b 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): """ @@ -1169,8 +1140,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 6671d1011b..aee617330b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -24,13 +24,15 @@ 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, split, timed_pass, timed_region, contains_val) from devito.types import (Buffer, Grid, Evaluable, host_layer, device_layer, disk_layer) +from devito.petsc.iet.passes import lower_petsc +from devito.petsc.clusters import petsc_lift __all__ = ['Operator'] @@ -422,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/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 765b8ba219..19aed147e4 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -401,7 +401,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 7221e985d5..d38f25ff2b 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 672971cf4f..e33df0fad8 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -211,6 +211,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 bed0bbad24..cae16941ac 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -18,7 +18,8 @@ RUN apt-get update && \ # 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 + autoconf libtool build-essential procps \ + gfortran pkgconf libopenblas-serial-dev # Install tmpi RUN curl https://raw.githubusercontent.com/Azrael3000/tmpi/master/tmpi -o /usr/local/bin/tmpi diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index 99b21c87fb..c38e8a1668 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 @@ -57,6 +75,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 78d95d29e6..2e2888071f 100644 --- a/docker/Dockerfile.intel +++ b/docker/Dockerfile.intel @@ -45,7 +45,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 @@ -59,16 +58,6 @@ RUN apt-get update -y && apt-get dist-upgrade -y && \ mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo \ # 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 @@ -93,11 +82,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 @@ -114,35 +99,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 @@ -157,12 +113,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 eb58f3c4970abb21fbd82d69695e8e9b28f435f9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 10 Jul 2024 11:31:30 +0100 Subject: [PATCH 096/107] 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 | 16 +- devito/ir/stree/algorithms.py | 1 - devito/operator/operator.py | 18 +- 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, 1822 insertions(+), 540 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 8c6068521f..2472f7268c 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 f9a86a5232..bb7a0fc2b2 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 a5d2354e33..9be714b055 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 d30f1752e1..49c66c6874 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 292ae9bec8..a5e3017b5e 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 dd513bcaed..629ebdde4a 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -374,7 +374,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 7026ad449b..8b85ad44df 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1122,7 +1122,7 @@ def defines(self): class Callback(Call): """ - Callback as a function pointer. + Base class for special callback types. Parameters ---------- @@ -1135,18 +1135,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 32c1964fc5..161f6ede19 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -314,7 +314,10 @@ def _args_call(self, args): return ret def _gen_signature(self, o, is_declaration=False): - decls = self._args_decl(o.parameters) + try: + decls = self._args_decl(o.parameters + o.unused_parameters) + except AttributeError: + decls = self._args_decl(o.parameters) prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),)) signature = c.FunctionDeclaration(c.Value(prefix, o.name), decls) if o.templates: @@ -614,7 +617,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) @@ -1423,11 +1426,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 aee617330b..6586cce1f9 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -31,8 +31,9 @@ split, timed_pass, timed_region, contains_val) from devito.types import (Buffer, Grid, Evaluable, host_layer, device_layer, disk_layer) -from devito.petsc.iet.passes import lower_petsc +from devito.petsc.iet.passes import lower_petsc, sort_frees from devito.petsc.clusters import petsc_lift +from devito.petsc.utils import derive_callback_dims, derive_struct_inputs __all__ = ['Operator'] @@ -349,7 +350,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 @@ -494,6 +495,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 @@ -512,7 +516,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'))) @@ -520,7 +529,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 c38e8a1668..1e6edbfc28 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -110,4 +110,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 f229935e5bc1cdfcab3a1784bd7943bb98c21d58 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 23 Jul 2024 21:18:45 +0100 Subject: [PATCH 097/107] 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 | 7 +- devito/operator/operator.py | 25 +- 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, 1006 insertions(+), 1527 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 2472f7268c..03efcb0c3d 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 bb7a0fc2b2..3b84bc4f61 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 9be714b055..3450c78a7f 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 49c66c6874..bdabf58f7d 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 a5e3017b5e..39ec4b5439 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 e7192ecf9b..86eeac77e0 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 8b85ad44df..ebe5f81f1a 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 @@ -750,6 +751,15 @@ def defines(self): return self.all_parameters +class FixedArgsCallable(Callable): + + """ + A Callable class that enforces a fixed function signature. + """ + + pass + + class CallableBody(MultiTraversable): """ @@ -1028,8 +1038,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 @@ -1048,13 +1058,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 161f6ede19..9359b2c3c0 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 @@ -314,10 +314,7 @@ def _args_call(self, args): return ret def _gen_signature(self, o, is_declaration=False): - try: - decls = self._args_decl(o.parameters + o.unused_parameters) - except AttributeError: - decls = self._args_decl(o.parameters) + decls = self._args_decl(o.parameters) prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),)) signature = c.FunctionDeclaration(c.Value(prefix, o.name), decls) if o.templates: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 6586cce1f9..61d297fb47 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -31,9 +31,8 @@ split, timed_pass, timed_region, contains_val) from devito.types import (Buffer, Grid, Evaluable, host_layer, device_layer, disk_layer) -from devito.petsc.iet.passes import lower_petsc, sort_frees -from devito.petsc.clusters import petsc_lift -from devito.petsc.utils import derive_callback_dims, derive_struct_inputs +from devito.petsc.iet.passes import lower_petsc +from devito.petsc.clusters import petsc_lift, petsc_project __all__ = ['Operator'] @@ -350,7 +349,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 @@ -377,8 +376,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 @@ -386,6 +383,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: @@ -425,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) @@ -495,9 +493,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 @@ -516,12 +511,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'))) @@ -529,8 +519,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 19aed147e4..557990032c 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -296,6 +296,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) @@ -401,9 +405,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 d38f25ff2b..cb36d18541 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 @@ -602,12 +602,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 9f49440709..b7f6572a26 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -79,7 +79,8 @@ def rule1(dep, candidates, loc_dims): rules = [rule0, rule1] # 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 e33df0fad8..672971cf4f 100644 --- a/devito/symbolics/printer.py +++ b/devito/symbolics/printer.py @@ -211,9 +211,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 6900496f35..34237f01bf 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 d4369b8f47..73a2e4956c 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1475,6 +1475,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 662cdd0d34..8b3ead9873 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -218,3 +218,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 0f9bb9bfb8c84e85ac173b8d9b9b0fbeb5c0dff6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 2 Oct 2024 20:22:24 +0100 Subject: [PATCH 098/107] compiler: Don't ever generate time loop for petsc callbacks --- devito/operator/operator.py | 6 +-- 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, 126 insertions(+), 137 deletions(-) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 61d297fb47..f51845ac20 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -32,7 +32,7 @@ from devito.types import (Buffer, Grid, Evaluable, host_layer, device_layer, disk_layer) from devito.petsc.iet.passes import lower_petsc -from devito.petsc.clusters import petsc_lift, petsc_project +from devito.petsc.clusters import petsc_lift __all__ = ['Operator'] @@ -376,6 +376,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 @@ -383,8 +385,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 b7f6572a26..e590de3d4e 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -79,8 +79,9 @@ def rule1(dep, candidates, loc_dims): rules = [rule0, rule1] # 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 f5b07465ed4a635d3039eb62a0b5e2970adab1d7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 4 Oct 2024 12:03:46 +0100 Subject: [PATCH 099/107] compiler: Move passes from utils into passes and create petsc_preprocess --- devito/operator/operator.py | 7 ++- 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, 104 insertions(+), 90 deletions(-) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index f51845ac20..3172b90ae4 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -32,7 +32,7 @@ from devito.types import (Buffer, Grid, Evaluable, host_layer, device_layer, disk_layer) from devito.petsc.iet.passes import lower_petsc -from devito.petsc.clusters import petsc_lift +from devito.petsc.clusters import petsc_preprocess __all__ = ['Operator'] @@ -376,9 +376,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 07a27bf32b7f29f47f680e8b9f126a8ee92e2f41 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 4 Oct 2024 14:24:08 +0100 Subject: [PATCH 100/107] 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 03efcb0c3d..8c6068521f 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 3b84bc4f61..f9a86a5232 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 3450c78a7f..a5d2354e33 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 bdabf58f7d..d30f1752e1 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 39ec4b5439..292ae9bec8 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 7ad60cc911827ce2439c3f304dcbe64108b35f6d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 14 Oct 2024 11:01:50 +0100 Subject: [PATCH 101/107] 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 73a2e4956c..d4369b8f47 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -1475,37 +1475,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 9ef17eb4c37aa0df699432ca9b199f8f42d3b867 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 31 Oct 2024 15:40:25 +0000 Subject: [PATCH 102/107] workflows: Add petsc docker image --- .github/workflows/docker-petsc.yml | 96 ++++++++++++++++++++++++++++++ docker/Dockerfile.devito | 2 +- 2 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/docker-petsc.yml diff --git a/.github/workflows/docker-petsc.yml b/.github/workflows/docker-petsc.yml new file mode 100644 index 0000000000..a96c945ce0 --- /dev/null +++ b/.github/workflows/docker-petsc.yml @@ -0,0 +1,96 @@ +name: Build PETSc Docker images + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +on: + pull_request: + branches: + - master + workflow_dispatch: + schedule: + # Run once a week + - cron: "0 13 * * 1" + +jobs: +####################################################### +############## Basic gcc CPU ########################## +####################################################### + deploy-cpu-bases: + name: "cpu-base" + runs-on: ubuntu-latest + env: + DOCKER_BUILDKIT: "1" + + steps: + - name: Checkout devito + uses: actions/checkout@v4 + + - name: Check event name + run: echo ${{ github.event_name }} + + - name: Set up QEMU + uses: docker/setup-qemu-action@v2 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v2 + + - name: Login to DockerHub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + + - name: cleanup + run: docker system prune -a -f + + - name: GCC image + uses: docker/build-push-action@v5 + with: + context: . + file: './docker/Dockerfile.cpu' + push: true + target: 'gcc' + build-args: 'arch=gcc' + tags: 'zoeleibowitz/bases:cpu-gcc' + deploy-petsc: + name: "petsc" + needs: deploy-cpu-bases + runs-on: ubuntu-latest + env: + DOCKER_BUILDKIT: "1" + steps: + - name: Checkout devito + uses: actions/checkout@v4 + + - name: Check event name + run: echo ${{ github.event_name }} + + - name: Set up QEMU + uses: docker/setup-qemu-action@v2 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v2 + + - name: Login to DockerHub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + + - name: Build and push image + uses: docker/build-push-action@v3 + with: + context: . + file: ./docker/Dockerfile.devito + push: true + tags: petsc-dev + build-args: base=zoeleibowitz/bases:cpu-gcc + + - name: Remove dangling layers + run: docker system prune -f + + - name: Run tests + run: | + docker run ${{ matrix.flag }} --rm -t --name testrun 'devitocodes/devito:petsc-dev' pytest ${{ matrix.test }} diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index 1e6edbfc28..a1ea9b90c1 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -3,7 +3,7 @@ ############################################################## # Base image with compilers -ARG base=devitocodes/bases:cpu-gcc +ARG base=zoeleibowitz/bases:cpu-gcc ARG petscinstall="" FROM $base as copybase From c965d21d08b160bf232a3db6205aefdab17d0cb2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 31 Oct 2024 16:16:29 +0000 Subject: [PATCH 103/107] workflows: edit petsc wf --- .github/workflows/docker-petsc.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/docker-petsc.yml b/.github/workflows/docker-petsc.yml index a96c945ce0..29ed942da9 100644 --- a/.github/workflows/docker-petsc.yml +++ b/.github/workflows/docker-petsc.yml @@ -85,7 +85,7 @@ jobs: context: . file: ./docker/Dockerfile.devito push: true - tags: petsc-dev + tags: 'zoeleibowitz/petsc-dev' build-args: base=zoeleibowitz/bases:cpu-gcc - name: Remove dangling layers From b66d064b6d9f558672210ee7d1af1dbbeb24f640 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 31 Oct 2024 17:11:08 +0000 Subject: [PATCH 104/107] workflows: Add test to petsc docker yml --- .github/workflows/docker-petsc.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/docker-petsc.yml b/.github/workflows/docker-petsc.yml index 29ed942da9..6c1ef228e9 100644 --- a/.github/workflows/docker-petsc.yml +++ b/.github/workflows/docker-petsc.yml @@ -93,4 +93,4 @@ jobs: - name: Run tests run: | - docker run ${{ matrix.flag }} --rm -t --name testrun 'devitocodes/devito:petsc-dev' pytest ${{ matrix.test }} + docker run ${{ matrix.flag }} --rm -t --name testrun 'zoeleibowitz/petsc-dev' pytest tests/test_operator.py From 2c4224d620eb082e26239804a247adc3fba84f77 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 31 Oct 2024 18:40:27 +0000 Subject: [PATCH 105/107] workflows: Only build the cpu bases each week --- .github/workflows/docker-petsc.yml | 40 ------------------------------ 1 file changed, 40 deletions(-) diff --git a/.github/workflows/docker-petsc.yml b/.github/workflows/docker-petsc.yml index 6c1ef228e9..4fea5b88a8 100644 --- a/.github/workflows/docker-petsc.yml +++ b/.github/workflows/docker-petsc.yml @@ -54,43 +54,3 @@ jobs: target: 'gcc' build-args: 'arch=gcc' tags: 'zoeleibowitz/bases:cpu-gcc' - deploy-petsc: - name: "petsc" - needs: deploy-cpu-bases - runs-on: ubuntu-latest - env: - DOCKER_BUILDKIT: "1" - steps: - - name: Checkout devito - uses: actions/checkout@v4 - - - name: Check event name - run: echo ${{ github.event_name }} - - - name: Set up QEMU - uses: docker/setup-qemu-action@v2 - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v2 - - - name: Login to DockerHub - uses: docker/login-action@v2 - with: - username: ${{ secrets.DOCKER_USERNAME }} - password: ${{ secrets.DOCKER_PASSWORD }} - - - name: Build and push image - uses: docker/build-push-action@v3 - with: - context: . - file: ./docker/Dockerfile.devito - push: true - tags: 'zoeleibowitz/petsc-dev' - build-args: base=zoeleibowitz/bases:cpu-gcc - - - name: Remove dangling layers - run: docker system prune -f - - - name: Run tests - run: | - docker run ${{ matrix.flag }} --rm -t --name testrun 'zoeleibowitz/petsc-dev' pytest tests/test_operator.py From 4c133831ce07ade30f450569a3259f2761879032 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 31 Oct 2024 18:44:51 +0000 Subject: [PATCH 106/107] workflows: Edit docker-petsc --- .github/workflows/docker-petsc.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/docker-petsc.yml b/.github/workflows/docker-petsc.yml index 4fea5b88a8..b0058423ce 100644 --- a/.github/workflows/docker-petsc.yml +++ b/.github/workflows/docker-petsc.yml @@ -1,4 +1,4 @@ -name: Build PETSc Docker images +name: Build CPU base for PETSc concurrency: group: ${{ github.workflow }}-${{ github.ref }} From b9b17e29111b9f957f4a254fd32ff1c5075c7401 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 31 Oct 2024 18:50:47 +0000 Subject: [PATCH 107/107] workflows: Remove edit in dockerfile.devito --- docker/Dockerfile.devito | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index a1ea9b90c1..1e6edbfc28 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -3,7 +3,7 @@ ############################################################## # Base image with compilers -ARG base=zoeleibowitz/bases:cpu-gcc +ARG base=devitocodes/bases:cpu-gcc ARG petscinstall="" FROM $base as copybase